from rsf.proj import *
from rsf.recipes.tpx import FPX
import math

##################################
# Usefule plotting functions
##################################

def velplot(title,label1='Depth',unit1='km',max1=2):
    return '''
    window max1=%g |
    grey color=j allpos=y title="%s" scalebar=y
    barlabel=Velocity barunit=km/s
    label1="%s" unit1="%s" label2=Lateral unit2=km
    barreverse=y pclip=100 
    ''' % (max1,title,label1,unit1)

def graph(transp,o2,d2,n2,col,fat,extra=''):
    return '''
    window max1=2 |
    graph transp=%d yreverse=y pad=n min2=%g max2=%g
    plotcol=%d plotfat=%d wantaxis=n wanttitle=n %s
    ''' % (transp,o2,o2+(n2-1)*d2,col,fat,extra)

def graphw(min1,max1,transp,o2,d2,n2,col,fat,extra=''):
    return '''
    window min1=%g max1=%g |
    graph transp=%d yreverse=y pad=n min2=%g max2=%g
    plotcol=%d plotfat=%d wantaxis=n wanttitle=n %s
    ''' % (min1,max1,transp,o2,o2+(n2-1)*d2,col,fat,extra)

def wiggle(title):
    return  '''
    window max1=2 j2=2 |
    wiggle transp=y yreverse=y poly=y 
    title="%s" label2=Offset unit2=km zplot=0.5
    wherexlabel=t wheretitle=b labelsz=12 titlesz=15
    ''' % title

Flow('modl',None,
     '''
     spike n1=1501 o1=-10 d1=0.02 n2=4
     nsp=4 k2=1,2,3,4 mag=0.4,0.7,1,1.5
     ''')
Flow('refl',None,
     '''
     spike n1=1501 n2=4 nsp=4 k2=1,2,3,4
     mag=0.0909091,0.1428570,0.1111110,0.2000000
     ''')
Flow('mod1','modl','window min1=0')

################################################################
#1.25 too shallow
#2.0 works perfectly
#1.75 - little little problems on the last offset
#Flow('rmodl','modl',
#     '''
#     pad n2=100 | noise rep=y seed=112012
#     type=n mean=1.25 range=1
#     ''')
Flow('rmodl','modl',
     '''
     pad n2=100 | noise rep=y seed=112012
     type=n mean=2.0 range=1
     ''')
################################################################

Flow('amodl','modl rmodl','cat axis=2 ${SOURCES[1]}')

Flow('rrefl','modl',
     '''
     pad n2=100 | noise rep=y seed=122012 |
     math output="input^3"
     ''')
Flow('prefl','rrefl','clip2 lower=10')
Flow('mrefl','rrefl','clip2 upper=-10')
################################################################
#Flow('drefl','prefl mrefl','add scale=0.01,0.01 ${SOURCES[1]}')
Flow('drefl','prefl mrefl','add scale=0.01,0.01 ${SOURCES[1]}')
################################################################

Flow('arefl','refl drefl','cat axis=2 ${SOURCES[1]}')

Flow('unif','mod1','unif2 n1=101 d1=0.02 v00=5,6,8,10,15')
#Result('modl','unif',velplot('Velocity Model',max1=2))

Flow('mod2','unif','math output=1.5+2*x1')

Plot('modl','mod2',velplot('Velocity Model',max1=2))
Plot('modla','mod1',graph(0,0,0.02,101,0,20,'scalebar=y'))
Plot('modlb','mod1',graph(0,0,0.02,101,7,5,'scalebar=y'))
#Result('modl2','modl modla modlb','Overlay')

#Flow('data','amodl arefl',
#     '''
#     kirmod nt=600 dt=0.004 freq=25 refl=${SOURCES[1]}
#     nh=51  dh=0.05 h0=0  
#     ns=501 ds=0.02 s0=0 cmp=y
#     vel=1.5 gradz=1 type=v |
#     put label2=Offset unit2=km label3=Midpoint unit3=km
#     ''',split=[1,1501],reduce='add')

Flow('diffr','rmodl drefl',
     '''
     kirmod nt=600 dt=0.004 freq=25 refl=${SOURCES[1]}
     nh=24  dh=0.05 h0=0  
     ns=501 ds=0.02 s0=0 cmp=y
     vel=2.0 gradz=1 type=v |
     put label2=Offset unit2=km label3=Midpoint unit3=km |
     window | pow pow1=1
     ''',split=[1,1501],reduce='add')
#vel=1.5 type=c |
#gradz=1 and type=v
#Result('diffr','transp plane=23 | grey gainpanel=all min1=0.5 max1=2.0 min2=1.0 max2=9.0 wanttitle=n')

# Making offsets to be half offsets
# creating CMP_Y additional dimension for
# sfpreconstkirchhoff

# Padding is extremely important
# Looks like if you have edge effects slope decomp
# might be wrong

#| pad end1=100
Flow('tiffr','diffr','costaper nw1=400 nw3=40')

#Result('diffr','transp plane=23 | grey gainpanel=all min1=0.5 max1=2.0 min2=1.0 max2=9.0 wanttitle=n')
#Result('tiffr','transp plane=23 | grey gainpanel=all wanttitle=n')

# dh=0.05 => 0.025
Flow('data','tiffr',
            '''
            spray axis=4 n=1 |
            put label4=CMP_Y |
            transp plane=23 |
            transp plane=34 |
            put d4=0.025 
            ''')

# VC range

v0=2.0#1.5
nv=151
dv=0.01

# first step
# 4th dimension is offset
# 3rd dimension is CMP_Y n3=1

Flow('mig_no_halfint','data','preconstkirch aal=y zero=n vel=%g' %v0, split=[4,24])

# correct phase

Flow('mig','mig_no_halfint','transp plane=24 | transp plane=34 | halfint inv=y adj=y | transp plane=23 | transp plane=34')
#Result('mig','window | grey gainpanel=all min1=0.5 max1=2.0 min2=1.0 max2=9.0 wanttitle=n')
#Result('mig_nw','mig','window | grey gainpanel=all wanttitle=n')

# slope decomposition

pad=1000
padft=501#0
nw=231#501

# get rid of CMP_Y by simple window command
Flow('warp','mig','window | t2warp pad=%i'%(pad))
#Flow('warp','mig','window n4=1 | t2warp pad=%i'%(pad))

# Padding is extremely important
# Looks like if you have edge effects slope decomp
# might be wrong
# so it is a really good idea to check the spectrum
# a really good way to plot spectrum
#<warp.rsf sffft1 | sfmath output="abs(input)" | sfreal | sfgraph | sfpen

# I slope decompose all the offsets at once
np=351
#<mig.rsf sfwindow min1=0.8 max1=1.2 min2=3.5 max2=5.0 f4=-1 | sfgrey | sfpen
p0=-1.7
#dp = -p0/((np-1)/2)
FPX('fpx','warp',np=np,p0=p0,nw=nw,v0=0.0)
#FPX('fpx','warp',np=np,p0=p0,nw=501,v0=0.0)

# check slope decomposition
offset_num = 0#20
Flow('txp','fpx','pad n1=%i | fft1 inv=y | t2warp inv=y | transp plane=23 memsize=1000'%padft)
#Result('txp','window n4=1 f4=%d | grey gainpanel=all min1=0.5 max1=2.0 min2=1.0 max2=9.0 title="Slope Decomposed Data offset#=%d" '%(offset_num,offset_num))
#Result('migcomp','txp','stack axis=3 | grey title="Stacked Slope Decomposed Data"')

# PSOVC

# if you use psovcp instead of psovc you can avert transp highlighted by ^^^

                                                                     #^^^
#Flow('f_hall_pk','fpx','fft3 axis=3 | transp plane=24 memsize=1000 | transp plane=34 memsize=1000')
Flow('f_hall_kp','fpx','fft3 axis=3 | transp plane=24 memsize=30000')
Flow('f_hs_pk','fpx','sfwindow n4=1 f4=%d | fft3 axis=3 | transp plane=24 memsize=1000 | transp plane=34 memsize=1000'%(offset_num))

#sfwindow n4=1 min4=0.5
#Flow('f_hs_kp','fpx','sfwindow n4=1 f4=%d | fft3 axis=3 | transp plane=24 memsize=1000'%(offset_num))

#####################################################################################################
# checking the amplitudes
#<fpx.rsf sfstack | sffft1 inv=y | sfattr
#<fpx.rsf sfwindow n4=1 | sfstack | sffft1 inv=y | sfattr
# amplitudes have one order
# l2 norm is higher for full offset - not sure if it matters
# l2 for zo = 0.004
# l2 for all offsets = 0.018
# l2 is dependent on the # of coefficients
#####################################################################################################

# parallelezation is done for k - lateral wavenumber
#Flow('vc_f_hall_pk_psovc','f_hall_pk','psovc nv=%d dv=%g v0=%g' % (nv,dv,v0),split=[4,1024])
Flow('vc_f_hall_kp_psovc','f_hall_kp','psovcp nv=%d dv=%g v0=%g' % (nv,dv,v0),split=[4,np])
Flow('vc_f_hs_pk_psovc','f_hs_pk','psovc nv=%d dv=%g v0=%g' % (nv,dv,v0),split=[4,1024])
#Flow('vc_f_hs_kp_psovc','f_hs_kp','psovcp nv=%d dv=%g v0=%g' % (nv,dv,v0),split=[4,np])

#Result('vc_psovc','grey gainpanel=all min1=0.5 max1=2.0 min2=1.0 max2=9.0 wanttitle=n')
#Result('vc_psovc15','vc_psovc','window n3=1 min3=1.5 | grey gainpanel=all min1=0.5 max1=2.0 min2=1.0 max2=9.0 wanttitle=n')
#Result('vc_psovc_hs','grey gainpanel=all min1=0.5 max1=2.0 min2=1.0 max2=9.0 wanttitle=n')
#Result('vc_psovc_hs15','vc_psovc_hs','window n3=1 min3=1.5 | grey gainpanel=all min1=0.5 max1=2.0 min2=1.0 max2=9.0 wanttitle=n')

### semblance estimation #####################################################
#^^^
Flow('fvkp_hs','vc_f_hs_pk_psovc',
     '''
     transp memsize=50000 plane=34
     ''')

# export OMP_NUM_THREADS=1 and memsize=100 for fft1 or you will get memory allocation problem and overhead 
# not sure how to make it optimal for parallel
Flow('tvxp_hs','fvkp_hs','fft3 axis=3 inv=y | pad n1=%i | fft1 memsize=100 inv=y | t2warp inv=y'%padft,split=[4,np])
#Flow('tvxp_hs','fvkp_hs','fft3 axis=3 inv=y | pad n1=%i | fft1 memsize=100 inv=y | t2warp inv=y')
Flow('tvxp','vc_f_hall_kp_psovc','fft3 axis=3 inv=y | pad n1=%i | fft1 memsize=100 inv=y | t2warp inv=y'%padft,split=[4,np])

Flow('tvx2_hs','tvxp_hs','mul $SOURCE | stack axis=4 norm=n')
Flow('tvx2','tvxp','mul $SOURCE | stack axis=4 norm=n')

Flow('tvx_hs','tvxp_hs','stack axis=4 norm=n')
Flow('tvx','tvxp','stack axis=4 norm=n')

# rect3 since x is the third dimension
# rect1 since t is the first dimension
Flow('semb_hs','tvx_hs tvx2_hs',
     '''
     mul ${SOURCES[0]} |
     divn den=${SOURCES[1]} niter=50 rect3=1000 rect1=5 |
     clip2 lower=0
     ''')
Flow('semb','tvx tvx2',
     '''
     mul ${SOURCES[0]} |
     divn den=${SOURCES[1]} niter=50 rect3=1000 rect1=5 |
     clip2 lower=0
     ''')

Flow('semb_hs_2','semb_hs','stack axis=3 | math output="input^2"')
Flow('semb_2','semb','stack axis=3 | math output="input^2"')

##############################################################################

Plot('diffr_hs','diffr',
     '''
     window n2=1 f2=%d |
     grey min1=0.75
     '''%(offset_num))

### Picking semblance ########################################################

Flow('muted_hs','semb_hs_2',
     '''
     scale axis=2 |
     mutter x0=2.0 inner=y v0=0.5 half=n t0=0.25  |
     mutter x0=2.2 inner=n v0=0.7 half=n t0=-0.25
     ''')
Flow('muted','semb_2',
     '''
     scale axis=2 |
     mutter x0=2.0 inner=y v0=0.5 half=n t0=0.25  |
     mutter x0=2.2 inner=n v0=0.7 half=n t0=-0.25
     ''')

#Flow('muted','semb_hs_2',
#     '''
#     scale axis=2 |
#     mutter x0=1.5 inner=y v0=0.5 half=n t0=0.25  |
#     mutter x0=1.5 inner=n v0=0.7 half=n t0=-0.25
#     ''')

Flow('pik_hs','muted_hs',
     '''
     pick rect1=25 vel0=1.5 gate=100 an=5 |
     spray axis=2 n=501 d=0.02 o=0
     ''')
Flow('pik','muted',
     '''
     pick rect1=25 vel0=1.5 gate=100 an=5 |
     spray axis=2 n=501 d=0.02 o=0
     ''')

Plot('muted_hs',
     '''
     grey min1=0.75 max1=1.9 color=j allpos=y labelsz=12 titlesz=15 
     label2=Velocity unit2=km/s title="Velocity Scan"
     ''')
Plot('muted',
     '''
     grey min1=0.75 max1=1.9 color=j allpos=y labelsz=12 titlesz=15 
     label2=Velocity unit2=km/s title="Velocity Scan"
     ''')

Plot('semb_hs_2',
     '''
     window min1=0.75 max1=1.9 | 
     grey color=j allpos=y labelsz=12 titlesz=15 
     label2=Velocity unit2=km/s title="Diffraction Velocity Scan"
     ''')
Plot('semb_2',
     '''
     window min1=0.75 max1=1.9 | 
     grey color=j allpos=y labelsz=12 titlesz=15 
     label2=Velocity unit2=km/s title="Diffraction Velocity Scan"
     ''')

# problems with velocity:
# diffractions do not come into focus
# at the true vel
# added 50 m/s  + 0.05
Flow('vtrue','diffr','window n2=1 | math output="2.0*sqrt((exp(x1)-1)/x1)" ')

Flow('slice_hs','tvx_hs pik_hs',
     'slice pick=${SOURCES[1]}')

Flow('slice_hs_true','tvx_hs vtrue',
     'slice pick=${SOURCES[1]}')

Flow('slice','tvx pik',
     'slice pick=${SOURCES[1]}')
Flow('slice_true','tvx vtrue',
     'slice pick=${SOURCES[1]}')

#Result('slice_hs','window min1=0.75 max1=1.9 | grey title="Diffraction Image" ')
#Result('slice_hs_true','window min1=0.75 max1=1.9 | grey title="Diffraction Image vtrue" ')
#Result('slice','window min1=0.75 max1=1.9 | grey title="Diffraction Image" ')
#Result('slice_true','window min1=0.75 max1=1.9 | grey title="DI vtrue" ')

Plot('vela','vtrue',graphw(0.75,1.9,1,v0,dv,nv,0,20))
Plot('velb','vtrue',graphw(0.75,1.9,1,v0,dv,nv,7,5))

Plot('pika_hs','pik_hs',graphw(0.75,1.9,1,v0,dv,nv,0,20))
Plot('pikb_hs','pik_hs',graphw(0.75,1.9,1,v0,dv,nv,7,5))
Plot('pika','pik',graphw(0.75,1.9,1,v0,dv,nv,0,20))
Plot('pikb','pik',graphw(0.75,1.9,1,v0,dv,nv,7,5))

Plot('vtrue','vela velb','Overlay')
Plot('pik_hs','pika_hs pikb_hs','Overlay')
Plot('pik','pika pikb','Overlay')

#Result('smb_hs_true','semb_hs_2  vtrue','Overlay')
#Result('smb_hs_picked','muted_hs  pik','Overlay')
#Result('smb_picked','muted  pik','Overlay')

########################################################################################

# lets plot one dip angle gather above the diffractor - offsets r converted into velocity
# dimension
Flow('dadcig232','tvxp','window n3=1 min3=2.32 | transp plane=23 memsize=10000')
Flow('dadcig232_hs','tvxp_hs','window n3=1 min3=2.32 | transp plane=23 memsize=10000')

#| window n4=1 min4=2.32 | transp plane=23 | fft1 inv=y | t2warp inv=y
#Result('dadcig232_15','grey gainpanel=all min1=0.5 max1=2.0 wanttitle=n')
#Result('dadcig232','grey gainpanel=all min1=0.5 max1=2.0 wanttitle=n')
#Result('dadcig232_hs','grey gainpanel=all min1=0.5 max1=2.0 wanttitle=n')
#Result('dadcig232pi','dadcig232','stack axis=3 | grey gainpanel=all min1=0.5 max1=2.0 wanttitle=n')

########################################################################################

# extracting single non-zero offset 
# and migrating it to three velocitites - true, lower, higher
# to check how gathers look like

# make another single offset - 1.0 [km](0.5 since it is half offset)

Flow('f_co_pk','fpx','window n4=1 f4=-1 | fft3 axis=3 | transp plane=24 memsize=1000 | transp plane=34 memsize=1000')

# transform 1.0 [km] (0.5 since it is half offset) offset to three velocities
# true velocity is 3.0

low = 0.1
true = 1.0
high = 1.9
mint = 0.8#0.5
maxt = 1.8#2.0
minp = -1.0
maxp = 1.0

# v0=2.0 + low/true/high = actual migration velocity

Flow('vc_f_co_pk_low','f_co_pk','psovc nv=%d dv=%g v0=%g' % (1,low,v0),split=[4,1024])
Flow('vc_f_co_pk_true','f_co_pk','psovc nv=%d dv=%g v0=%g' % (1,true,v0),split=[4,1024])
Flow('vc_f_co_pk_high','f_co_pk','psovc nv=%d dv=%g v0=%g' % (1,high,v0),split=[4,1024])

# increasing the highest velocity for zero offset
Flow('vc_f_zo_pk_high','f_hs_pk','psovc nv=%d dv=%g v0=%g' % (1,high,v0),split=[4,1024])

cigplot = '''
        grey min1=%g max1=%g min2=%g max2=%g labelsz=12. titlesz=12. d2num=0.25 n2tic=4 o2num=1.0
        unit2="s\^2\_/km" title="Slope Gather (x=%g km, v=%g km/s)"
        '''

Flow('dadcig232_co_low','vc_f_co_pk_low','window | fft3 axis=3 inv=y | window n3=1 min3=2.32 | pad n1=%i | fft1 inv=y | t2warp inv=y'%padft)
Flow('dadcig232_co_true','vc_f_co_pk_true','window | fft3 axis=3 inv=y | window n3=1 min3=2.32 | pad n1=%i | fft1 inv=y | t2warp inv=y'%padft)
Flow('dadcig232_co_high','vc_f_co_pk_high','window | fft3 axis=3 inv=y | window n3=1 min3=2.32 | pad n1=%i | fft1 inv=y | t2warp inv=y'%padft)

Flow('dadcig232_zo_high','vc_f_zo_pk_high','window | fft3 axis=3 inv=y | window n3=1 min3=2.32 | pad n1=%i | fft1 inv=y | t2warp inv=y'%padft)

Plot('a_dadcig232_co_low','dadcig232_co_low',cigplot %(mint,maxt,minp,maxp,2.32,v0+low))
Plot('a_dadcig232_co_true','dadcig232_co_true',cigplot %(mint,maxt,minp,maxp,2.32,v0+true))
Plot('a_dadcig232_co_high','dadcig232_co_high',cigplot %(mint,maxt,minp,maxp,2.32,v0+high))

Flow('dadcig232_zo_low','dadcig232_hs','window n3=1 min3=%g'%(v0+low))
Flow('dadcig232_zo_true','dadcig232_hs','window n3=1 min3=%g'%(v0+true))
#Flow('dadcig232_zo_high','dadcig232_hs','window n3=1 min3=%g'%(v0+high))

Plot('a_dadcig232_zo_low','dadcig232_zo_low',cigplot %(mint,maxt,minp,maxp,2.32,v0+low))
Plot('a_dadcig232_zo_true','dadcig232_zo_true',cigplot %(mint,maxt,minp,maxp,2.32,v0+true))
Plot('a_dadcig232_zo_high','dadcig232_zo_high',cigplot %(mint,maxt,minp,maxp,2.32,v0+high))

Result('a-cigs-zo','a_dadcig232_zo_low a_dadcig232_zo_true a_dadcig232_zo_high','SideBySideAniso')
Result('a-cigs-co','a_dadcig232_co_low a_dadcig232_co_true a_dadcig232_co_high','SideBySideAniso')

#Result('a_cigs_low','a_dadcig232_zo_low a_dadcig232_co_low','SideBySideAniso')
#Result('a_cigs_true','a_dadcig232_zo_true a_dadcig232_co_true','SideBySideAniso')
#Result('a_cigs_high','a_dadcig232_zo_high a_dadcig232_co_high','SideBySideAniso')

# Lets compare it with fourvc

#Flow('ckx48','mig','cosft sign2=1 | window | transp plane=23')

#Flow('vlf48','ckx48',
#              '''
#              window n2=1 f2=%d | transp plane=23 |              
#              fourvc nv=%g dv=%g v0=%g verb=y
#              '''%(offset_num,nv,dv,v0),split=[3,501])
#Flow('vlf480','vlf48','cosft sign3=-1')

#Result('vlf480','transp plane=23 | grey gainpanel=all min1=0.5 max1=2.0 min2=1.0 max2=9.0 wanttitle=n')
#Result('vlf480_15','vlf480','window n2=1 min2=1.5 | grey gainpanel=all min1=0.5 max1=2.0 min2=1.0 max2=9.0 wanttitle=n')

##############################################################################
### Pictures for article #####################################################
##############################################################################

Result('a-zo','diffr','window min1=0.8 max1=1.6 n2=1 min3=1.0 max3=9.0 | grey wanttitle=n')
#Result('a_vtrue4zo','vtrue','window min1=0.8 max1=1.6 min2=1.0 max2=9.0 | grey color=j allpos=n mean=y wanttitle=n scalebar=y')

Result('a-co','diffr','window min1=0.8 max1=1.6 n2=1 min2=1.0 min3=1.0 max3=9.0 | grey wanttitle=n')

Result('a-semb-zo','semb_hs_2',
        '''
        scale axis=2 | 
        grey min1=0.75 max1=1.9 color=j allpos=n bias=0.5 clip=0.5
        label2=Velocity unit2=km/s wanttitle=n scalebar=y
        ''')

Result('a-semb-fo','semb_2',
        '''
        scale axis=2 | 
        grey min1=0.75 max1=1.9 color=j allpos=n bias=0.5 clip=0.5 
        label2=Velocity unit2=km/s wanttitle=n scalebar=y
        ''')

Result('a-slice-zo','slice_hs',
        '''
        window min1=0.8 max1=1.6 min2=1.0 max2=9.0 |
        grey wanttitle=n clip=0.003
        ''')

Result('a-slice-fo','slice',
        '''
        window min1=0.8 max1=1.6 min2=1.0 max2=9.0 |
        grey wanttitle=n clip=0.08
        ''')

End()