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=10 titlesz=12 titlefat=6 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=10 titlesz=12 titlefat=6 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')

Flow('vdix','vdix0',
     'window n2=250 f2=250 | put o2=7.705 d2=0.0335 d1=0.002 | put label1=Time unit1=s')

Plot('vdix',
     '''
     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=10 titlesz=12 titlefat=6 labelfat=6
     ''')

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

# convert vdix to depth
Flow('init','vdix0',
     '''
     time2depth velocity=$SOURCE nz=801 dz=0.005 intime=y |
     window n2=250 f2=250 | put o2=7.705 d2=0.0335  |
     put label1=Depth unit1=km
     ''')

Plot('init',
     '''
     window max1=3.9 |
     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=10 titlesz=12 titlefat=6 labelfat=6
     ''')

# time-to-depth conversion
niter=5
cgiter=2500
rect1=25
rect2=10
eps=2

Flow('inv it ix if ig ic','init vdix',
     '''
     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))

Plot('rinit','ix',
     '''
     window n3=1 | contour nc=200 plotcol=7 plotfat=7
     wantaxis=n wanttitle=n scalebar=y      
     ''')
Plot('pinit','init rinit','Overlay')

Plot('inv',
     '''
     window max1=3.9 |
     grey title="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=10 titlesz=12 titlefat=6 labelfat=6
     ''')
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')

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=10 titlesz=12 titlefat=6 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=10 titlesz=12 titlefat=6 labelfat=6
     ''' % niter)

Plot('dinv','inv init',
     '''
     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=10 titlesz=12 titlefat=6 labelfat=6
     ''')

Flow('hist',None,
     '''
     echo 0.98 0.35 0.11 0.03 0.01
     n1=5 o1=1 d1=1 data_format=ascii_float in=$TARGET
     ''')
Result('hist',
       '''
       graph wanttitle=n plotfat=4 labelsz=7 titlesz=9 labelfat=3 titlefat=3
       label1="Number of Linearization Updates" unit1= label2="Normalized l2 Misfit" unit2=
       screenratio=0.8 screenht=7.5
       ''')

Result('init','vdix 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"
     label1=Depth unit1=km label2=Distance unit2=km
     labelsz=10 titlesz=12 titlefat=6 labelfat=6
     ''')

# depth migration
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 ('init','inv'):
    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=6 labelfat=6
         ''')

    # zoom
    if modl=='init':
        title='Initial'
    else:
        title='Inverted'

    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=='init':
        title='Initial'
    else:
        title='Inverted'

    Plot('cig1'+modl+'a','dmig'+modl,
         '''
         window n2=1 min2=11 | window max1=3.9 |
         grey title="%s" pclip=90
         label1=Depth unit1=km label2=Offset unit2=km
         labelsz=10 titlesz=12 titlefat=6 labelfat=6
         ''' % (title+' at 11km'))
    Plot('cig1'+modl,'dmig'+modl,
         '''
         window n2=1 min2=11 | window min1=2 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 11km'))
    Plot('cig2'+modl+'a','dmig'+modl,
         '''
         window n2=1 min2=12 | window max1=3.9 |
         grey title="%s" pclip=90
         label1=Depth unit1=km label2=Offset unit2=km
         labelsz=10 titlesz=12 titlefat=6 labelfat=6
         ''' % (title+' at 12km'))
    Plot('cig2'+modl,'dmig'+modl,
         '''
         window n2=1 min2=12 | window min1=2 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 12km'))

# overlay CIGs with reference lines
Plot('cigref1init','dmiginv',
     '''
     window n2=1 min2=11 | math output=x1 | window min1=2 max1=3.9 |
     contour nc=3 c=2.69,3.74,3.86 wanttitle=n wantaxis=n
     plotcol=3 plotfat=15 dash=9 screenht=35 screenratio=2.5
     ''')
Plot('cig1init0','cig1init cigref1init','Overlay')
Plot('cigref1inv','dmiginv',
     '''
     window n2=1 min2=11 | math output=x1 | window min1=2 max1=3.9 |
     contour nc=3 c=2.71,3.73,3.84 wanttitle=n wantaxis=n
     plotcol=3 plotfat=15 dash=9 screenht=35 screenratio=2.5
     ''')
Plot('cig1inv0','cig1inv cigref1inv','Overlay')

Plot('cigref2init','dmiginv',
     '''
     window n2=1 min2=11 | math output=x1 | window min1=2 max1=3.9 |
     contour nc=3 c=2.61,3.05,3.64 wanttitle=n wantaxis=n
     plotcol=3 plotfat=15 dash=9 screenht=35 screenratio=2.5
     ''')
Plot('cig2init0','cig2init cigref2init','Overlay')
Plot('cigref2inv','dmiginv',
     '''
     window n2=1 min2=11 | math output=x1 | window min1=2 max1=3.9 |
     contour nc=3 c=2.58,3.00,3.59 wanttitle=n wantaxis=n
     plotcol=3 plotfat=15 dash=9 screenht=35 screenratio=2.5
     ''')
Plot('cig2inv0','cig2inv cigref2inv','Overlay')

Result('dmig','mapd dmiginv','OverUnderAniso')
Result('ddmig0','dmiginit0 dmiginv0','OverUnderIso')
Result('cig0','cig1init0 cig1inv0 cig2init0 cig2inv0','SideBySideIso')

End()