from rsf.proj import *
from rsf.recipes.velcon import velcon

# fetch data
Fetch('beinew.HH','midpts')
Flow('bei','beinew.HH',
     '''
     dd form=native | transp plane=23 | transp plane=34 |
     put o1=0 label1=Time unit1=s label2=Midpoint unit2=km
     label3=Crossline unit3=km label4=Offset unit4=km
     ''')

# velocity continuation
velcon('bei',
       nv=125,      # continuation steps
       v0=1.5,      # initial velocity
       dv=0.01,     # velocity step
       nx=250,      # lateral dimension
       nh=48,       # number of offsets
       padt=1024,   # time padding
       padt2=2048,  # extra time padding
       padx=521,    # lateral padding
       dx=0.0335,   # lateral sampling
       n1=1000,     # time dimension
       x0=7.705,    # lateral origin
       srect1=15,
       srect2=5)

Plot('tmig','bei-agc',
     '''
     grey title="Time Migration"
     label1=Time unit1=s label2=Distance unit2=km
     labelsz=15 titlesz=16 titlefat=8 labelfat=6
     ''')

Plot('rms','bei-npk',
     '''
     grey title="Picked Migration Velocity"
     color=j scalebar=y barreverse=y mean=y
     label1=Time unit1=s label2=Distance unit2=km barunit=km/s
     labelsz=15 titlesz=16 titlefat=8 labelfat=6
     ''')

# from migration velocity to Dix velocity
Flow('bei-npk-mir','bei-npk','reverse which=2 opt=i')
Flow('bei-npk-ext','bei-npk-mir bei-npk','cat axis=2 ${SOURCES[1]} ${SOURCES[0]}')

Flow('vdix0 vmig0','bei-npk-ext',
     'dix rect1=15 rect2=5 vrmsout=${TARGETS[1]} niter=80')

# Consider only a portion of dix velocity
Flow('dix','vdix0',
     'window n2=250 f2=250 | put o2=7.705 d2=0.0335 d1=0.002 | put label1=Time unit1=s')

Plot('dix',
     '''
     put d1=0.004 | grey title="Dix Velocity"
     color=j scalebar=y barreverse=y
     label1=Time unit1=s label2=Distance unit2=km barunit=km/s
     minval=1.5 maxval=2.9 bias=2.1
     labelsz=15 titlesz=16 titlefat=8 labelfat=6
     ''')

Result('vdix','rms tmig','OverUnderAniso')

# convert vdix to depth
Flow('dixdepth','dix',
     '''
     time2depth velocity=$SOURCE nz=801 dz=0.005 intime=y twoway=n|
     put label1=Depth unit1=km | window max1=3.9 
     ''')

Plot('dixdepth',
     '''
     grey title="Dix-inverted Model"
     color=j scalebar=y barreverse=y
     label1=Depth unit1=km label2=Distance unit2=km barunit=km/s
     minval=1.5 maxval=2.9 bias=2.1
     labelsz=15 titlesz=16 titlefat=8 labelfat=6
     ''')
# time-to-depth conversion (Sripanich and Fomel, 2017)

# Derivatives of Dix velocity squared
Flow('dv2dt0','dix','math output="input^2" | smoothder')
Flow('dv2dx0','dix','math output="input^2" | transp | smoothder | transp')
Flow('beta','dv2dt0 dix',
     '''
     time2depth velocity=${SOURCES[1]} nz=801 dz=0.005 intime=y twoway=n|
     put label1=Depth unit1=km | window max1=3.9 
     ''')
Flow('alpha','dv2dx0 dix',
     '''
     time2depth velocity=${SOURCES[1]} nz=801 dz=0.005 intime=y twoway=n|
     put label1=Depth unit1=km | window max1=3.9 
     ''')
Plot('alpha',
     '''
     grey color=j scalebar=y title="dw\_d\^/dx\_0" pclip=100
     labelsz=15 titlesz=16 titlefat=8 labelfat=6 allpos=y minval=-1.7 maxval=1.7 bias=-1.7 clip=3.4
     screenratio=0.75 screenht=9 allpos=y barlabel="dw\_d\^/dx\_0\^" barunit="km/s\^2"
     ''')
Plot('beta',
     '''
     grey color=j scalebar=y title="dw\_d\^/dt\_0" pclip=100
     labelsz=15 titlesz=16 titlefat=8 labelfat=6 allpos=y minval=-5 maxval=12 bias=-5 clip=17
     screenratio=0.75 screenht=9 allpos=y barlabel="dw\_d\^/dt\_0\^" barunit="km\^2\_/s\^3"
     ''')

# Reference velocity squared
Flow('refdix','dixdepth','math output="input^2" ')
Flow('refvz','dixdepth','window n2=1 f2=125| math output="input^2" | spray axis=2 n=250 o=7.705 d=0.0335 ')

Plot('refdix',
     '''
     grey color=j scalebar=y title="Ref w\_dr\^(x,z)" 
     labelsz=15 titlesz=16 titlefat=8 labelfat=6 allpos=y minval=2.25 maxval=7 bias=2.25 clip=4.75
     screenratio=0.75 screenht=9 allpos=y barlabel="v\^2\_" barunit="km^2\_/s\^2"
     ''')
Plot('refvz',
     '''
     grey color=j scalebar=y title="Ref w(z)" 
     labelsz=15 titlesz=16 titlefat=8 labelfat=6 allpos=y minval=2.25 maxval=7 bias=2.25 clip=4.75
     screenratio=0.75 screenht=9 allpos=y barlabel="v\^2\_" barunit="km\^2\_/s\^2"
     ''')

Result('input-field','refdix alpha beta','OverUnderAniso')

# Find dv2 dt0 and dx0 from sftime2depthweak #####################################################
Flow('depth dx0 dt0 dv','refdix refvz alpha beta',
        '''
        time2depthweak zsubsample=50 nsmooth=16
        velocity=$SOURCE refvelocity=${SOURCES[1]} dvdx0=${SOURCES[2]} dvdt0=${SOURCES[3]}
        outdx0=${TARGETS[1]} outdt0=${TARGETS[2]} outdv=${TARGETS[3]}
        ''')
Flow('finalv','refvz dv','math est=${SOURCES[1]} output="sqrt(input+est)" | put d3=1 o3=0 ')
Plot('finalv',
     '''
     grey color=j scalebar=y title=" Estimated interval v(x,z)" 
     labelsz=15 titlesz=16 titlefat=8 labelfat=6  
     label1=Depth unit1=km label2=Distance unit2=km barunit=km/s
     allpos=y minval=1.5 maxval=2.6 bias=1.5 clip=1.1 
     screenratio=0.75 screenht=9  barlabel="v" barunit="km/s"
     ''')

Flow('diff1','finalv dixdepth','math est=${SOURCES[1]} output="input-est" | put d3=1 o3=0 ')
Flow('diff','finalv dixdepth','math est=${SOURCES[1]} output="input^2-est^2" | put d3=1 o3=0 ')
Plot('diff1',
     '''
     grey color=j scalebar=y title=" Estimated interval v(x,z) - v\_dr\^(x,z) " 
     labelsz=15 titlesz=16 titlefat=8 labelfat=6 
     label1=Depth unit1=km label2=Distance unit2=km barunit=km/s
     allpos=y minval=-0.17 maxval=0.17 clip=0.34 bias=-0.17
     screenratio=0.75 screenht=9  barlabel="v" barunit="km/s"
     ''')

Flow('reft0','refvz','math output="2*1/sqrt(input)*0.005" | causint') # two-way
Flow('finalt0','reft0 dt0','math dt=${SOURCES[1]} output="input+dt"')

Flow('refx0','refvz','math output="x2"')
Flow('finalx0','refx0 dx0','math dx=${SOURCES[1]} output="input+dx"')

Flow('dixcoord','reft0 refx0',
     '''
     cat axis=3 ${SOURCES[1]} |
     transp plane=23 | transp plane=12
     ''')

Flow('finalcoord','finalt0 finalx0',
     '''
     cat axis=3 ${SOURCES[1]} |
     transp plane=23 | transp plane=12
     ''')

Flow('dixmapd','bei-fmg dixcoord','inttest2 interp=spline nw=8 coord=${SOURCES[1]}')
Flow('finalmapd','bei-fmg finalcoord','inttest2 interp=spline nw=8 coord=${SOURCES[1]}')

Plot('dixmapd',
     '''
     agc rect1=200 | window  max1=3.9 | grey title="Time -> Depth (Dix)"
     label1=Depth unit1=km label2=Distance unit2=km
     labelsz=10 titlesz=12 titlefat=4 labelfat=4
     ''')      #labelsz=8 titlesz=10 titlefat=2 labelfat=2 screenratio=0.5 screenht=7.3 labelsz=10 titlesz=12 titlefat=6 labelfat=6 screenratio=0.75 screenht=9.5
Plot('finalmapd',
     '''
     agc rect1=200 | window  max1=3.9 | grey title="Time -> Depth (Proposed)"
     label1=Depth unit1=km label2=Distance unit2=km
     labelsz=10 titlesz=12 titlefat=4 labelfat=4
     ''')      #labelsz=8 titlesz=10 titlefat=2 labelfat=2 screenratio=0.5 screenht=7.3 labelsz=10 titlesz=12 titlefat=6 labelfat=6 screenratio=0.75 screenht=9.5

Result('finalvcompare','diff1 finalv inv','OverUnderAniso')
Result('dixcompare','dixmapd finalmapd','OverUnderAniso')
Result('finalcompare','finalmapd dmigfinalv','OverUnderAniso')
##########################################################################################
# time-to-depth conversion (Li and Fomel, 2015)
niter=5
cgiter=2500
rect1=25
rect2=10
eps=2

Flow('inv it ix if ig ic','dixdepth dix',
     '''
     tdconvert niter=%d cgiter=%d eps=%g shape=y rect1=%d rect2=%d dix=${SOURCES[1]} 
     t0=${TARGETS[1]} x0=${TARGETS[2]} f0=${TARGETS[3]} grad=${TARGETS[4]} cost=${TARGETS[5]}
     ''' % (niter,cgiter,eps,rect1,rect2))
# x0
Plot('rinit','ix',
     '''
     window n3=1 | contour nc=200 plotcol=7 plotfat=7
     wantaxis=n wanttitle=n scalebar=y      
     ''')
Plot('pinit','dixdepth rinit','Overlay')

# v
Plot('inv',
     '''
     window max1=3.9 |
     grey title="Inverted v(x,z) (Li and Fomel, 2015)"
     color=j scalebar=y 
     label1=Depth unit1=km label2=Distance unit2=km 
     barlabel="v" barunit="km/s"
     allpos=y minval=1.5 maxval=2.6 bias=1.5 clip=1.1 
     labelsz=15 titlesz=16 titlefat=8 labelfat=6 
     screenratio=0.75 screenht=9
     ''')
Plot('rinv','ix',
     '''
     window max1=3.9 |
     window n3=1 f3=%d | contour nc=200 plotcol=7 plotfat=7
     wantaxis=n wanttitle=n scalebar=y 
     ''' % niter)
Plot('pinv','inv rinv','Overlay')

# cost
Plot('cinit','ic',
     '''
     window n3=1 max1=3.9 |
     grey title="Initial f" color=j scalebar=y barreverse=y
     label1=Depth unit1=km label2=Distance unit2=km barlabel=Cost barunit=
     minval=-0.75 maxval=6 clip=2
     labelsz=15 titlesz=16 titlefat=8 labelfat=6
     ''')

Plot('cinv','ic',
     '''
     window n3=1 f3=%d max1=3.9 | grey title="Final f" color=j scalebar=y barreverse=y
     label1=Depth unit1=km label2=Distance unit2=km barlabel=Cost barunit=
     minval=-0.75 maxval=6 clip=2
     labelsz=15 titlesz=16 titlefat=8 labelfat=6
     ''' % niter)

Plot('dinv','inv dixdepth',
     '''
     add scale=1,-1 ${SOURCES[1]} | window max1=3.9 |
     grey title=Update
     color=j scalebar=y barreverse=y mean=y pclip=99.5
     label1=Depth unit1=km label2=Distance unit2=km barunit=km/s
     labelsz=15 titlesz=16 titlefat=8 labelfat=6
     ''')

Result('init','dix pinit','OverUnderAniso')
Result('inv','cinit cinv','OverUnderAniso')
Result('dinv','pinv dinv','OverUnderAniso')

# map time to depth
Flow('t0','it','window n3=1 f3=%d | scale dscale=2' % niter)
Flow('x0','ix','window n3=1 f3=%d' % niter)

Flow('coord','t0 x0',
     '''
     cat axis=3 ${SOURCES[1]} |
     transp plane=23 | transp plane=12
     ''')

Flow('mapd','bei-fmg coord','inttest2 interp=spline nw=8 coord=${SOURCES[1]}')

Plot('mapd',
     '''
     agc rect1=200 | window  max1=3.9 | grey title="Time -> Depth (Li and Fomel, 2015)"
     label1=Depth unit1=km label2=Distance unit2=km
     labelsz=10 titlesz=12 titlefat=6 labelfat=6
     ''')

# Depth migration ###################################################################################
Flow('flatvz','refvz','math output="sqrt(input)"| put d3=1 o3=0')

Flow('shot','bei','window | transp plane=23 | cmp2shot | put o1=0 d2=0.067 o2=0.264')

Flow('ys',None,'math n1=297 o1=5.9985 d1=0.0335 output=x1')
Flow('zs','ys','math output=0')
Flow('sht','zs ys','cat axis=2 ${SOURCES[1]} ${SOURCES[0]} | transp')

Flow('yr',None,'math n1=196 o1=6.2625 d1=0.067 output=x1')
Flow('zr','yr','math output=0')
Flow('rcv','zr yr','cat axis=2 ${SOURCES[1]} ${SOURCES[0]} | transp')

for modl in ('flatvz','dixdepth','inv','finalv'): # v(z), dix v(x,z), siwei, proposed
        Flow('l'+modl,modl,'window n2=1 f2=0 | spray axis=2 n=51 d=0.0335 o=5.9965')
        Flow('r'+modl,modl,'window n2=1 f2=249 | spray axis=2 n=98 d=0.0335 o=16.08')
        Flow('p'+modl,['l'+modl,modl,'r'+modl],
            '''
            cat axis=2 ${SOURCES[1]} ${SOURCES[2]} |
            transp plane=12 | spline o1=5.9965 d1=0.008375 n1=1593 | transp plane=12 |
            transp plane=34
            ''')
        # eikonal
        Flow([modl+'times',modl+'tdls',modl+'tdss'],['p'+modl,'sht'],
            '''
            eikods shotfile=${SOURCES[1]} tdl1=${TARGETS[1]} tds1=${TARGETS[2]} b1=2 b2=2 |
            put o4=5.9985 d4=0.0335 | window
            ''')
        Flow([modl+'timer',modl+'tdlr',modl+'tdsr'],['p'+modl,'rcv'],
            '''
            eikods shotfile=${SOURCES[1]} tdl1=${TARGETS[1]} tds1=${TARGETS[2]} b1=2 b2=2 |
            put o4=6.2625 d4=0.067 | window
            ''')

        # Kirchhoff with surface offset CIG
        Flow('dmig'+modl,['shot',modl+'times',modl+'tdss',modl+'timer',modl+'tdsr'],
            '''
            kirmigsr aperture=5 antialias=1 cig=y
            stable=${SOURCES[1]} sderiv=${SOURCES[2]}
            rtable=${SOURCES[3]} rderiv=${SOURCES[4]}
            ''')

        Plot('dmig'+modl,
            '''
            window j2=4 | window n2=250 f2=51 | stack axis=3 norm=n | window max1=3.9 |
            agc rect1=200 | grey title="Prestack Kirchhoff Depth Migration"
            label1=Depth unit1=km label2=Distance unit2=km
            labelsz=10 titlesz=12 titlefat=4 labelfat=4
            ''')
        # zoom
        if modl=='dixdepth':
            title='with w\_dr\^(x,z)'
        else:
            title='with w\_r\^(z) + \F10 D\F3 w'

        Plot('dmig'+modl+'0','dmig'+modl,
            '''
            stack axis=3 norm=n | window min1=2 max1=3.9 min2=10 max2=13 |
            agc rect1=200 | grey title="PSDM %s"
            label1=Depth unit1=km label2=Distance unit2=km
            labelsz=7 titlesz=9 titlefat=4 labelfat=4 screenht=9 screenratio=0.55
            ''' % title)

        # CIGs
        if modl=='dixdepth':
            title='Before'
        else:
            title='After'

        Plot('cig1'+modl+'a','dmig'+modl, # all
            '''
            window n2=1 min2=8 | window max1=3.9 |
            grey title="%s" pclip=90
            label1=Depth unit1=km label2=Offset unit2=km
            labelsz=20 titlesz=22 titlefat=15 labelfat=15
            ''' % (title+' at 8km'))
        Plot('cig1'+modl,'dmig'+modl, # cut
            '''
            window n2=1 min2=8 | window min1=3 max1=3.9 |
            grey title="%s" pclip=90 screenht=35 screenratio=2.5
            label1=Depth unit1=km label2=Offset unit2=km
            labelsz=20 titlesz=22 titlefat=15 labelfat=15
            ''' % (title+' at 8km'))
        Plot('cig2'+modl+'a','dmig'+modl,
            '''
            window n2=1 min2=11.25| window max1=3.9 |
            grey title="%s" pclip=90
            label1=Depth unit1=km label2=Offset unit2=km
            labelsz=20 titlesz=22 titlefat=15 labelfat=15
            ''' % (title+' at 11.25km'))
        Plot('cig2'+modl,'dmig'+modl,
            '''
            window n2=1 min2=11.25| window min1=3 max1=3.9 |
            grey title="%s" pclip=90 screenht=35 screenratio=2.5
            label1=Depth unit1=km label2=Offset unit2=km
            labelsz=20 titlesz=22 titlefat=15 labelfat=15
            ''' % (title+' at 11.25km')) # screenht=35

 # overlay CIGs with reference lines
Plot('cigref1dixdepth','dmigfinalv',
    '''
    window n2=1 min2=11 | math output=x1 | window min1=3 max1=3.9 |
    contour nc=3 c=3.375,3.785 wanttitle=n wantaxis=n
    plotcol=3 plotfat=20 dash=9 screenht=35 screenratio=2.5
    ''')
Plot('cig1init','cig1dixdepth cigref1dixdepth','Overlay')
Plot('cigref1finalv','dmigfinalv',
    '''
    window n2=1 min2=11 | math output=x1 | window min1=3 max1=3.9 |
    contour nc=3 c=3.39,3.81 wanttitle=n wantaxis=n
    plotcol=3 plotfat=20 dash=9 screenht=35 screenratio=2.5
    ''')
Plot('cig1invert','cig1finalv cigref1finalv','Overlay')

Plot('cigref2dixdepth','dmigfinalv',
    '''
    window n2=1 min2=11 | math output=x1 | window min1=3 max1=3.9 |
    contour nc=3 c=,3.73,3.795 wanttitle=n wantaxis=n
    plotcol=3 plotfat=20 dash=9 screenht=35 screenratio=2.5
    ''')
Plot('cig2init','cig2dixdepth cigref2dixdepth','Overlay')
Plot('cigref2finalv','dmigfinalv',
    '''
    window n2=1 min2=11 | math output=x1 | window min1=3 max1=3.9 |
    contour nc=3 c=,3.705,3.775 wanttitle=n wantaxis=n
    plotcol=3 plotfat=20 dash=9 screenht=35 screenratio=2.5
    ''')
Plot('cig2invert','cig2finalv cigref2finalv','Overlay')

# Result('dmig','mapd dmiginv','OverUnderAniso')
# Result('ddmig0','dmiginit0 dmiginv0','OverUnderIso')
# Result('dmigcompare','dmigdixdepth dmigfinalv','OverUnderAniso')
Result('cig','cig1init cig1invert cig2init cig2invert','SideBySideIso')
Result('cig1','cig1init cig1invert','SideBySideAniso')
Result('cig2','cig2init cig2invert','SideBySideAniso')





End()