from rsf.proj import *
import zomig

nt=800
dt=0.004

nx=1000
dx=13.3333

nz=650
dz=5.005

xmax = 13319.9667
zmax = 3200

# Get data from the server

for case in ('npk','fmg','smb','stk2'):
    elf = 'elf-' + case
    Fetch(elf+'.rsf','masha')
    Flow(case,elf,'dd form=native | put unit1=s')

# Plot

Plot('npk',
     '''
     scale dscale=0.001 | put d2=0.01333333 | 
     grey color=j title="\F2 Time migration velocity"
     label2="\F2 Lateral position" unit2=km parallel2=n
     allpos=y bias=1.9 scalebar=y barreverse=y barunit=km/s
     label1="\F2 Time" labelsz=10 titlesz=12.5
     barlabel= barunit= format2="%3.1f" formatbar="%3.1f"
     ''',local=1)
Plot('fmg',
     '''
     agc rect1=250 | put d2=0.01333333 | 
     grey title="\F2 Prestack time migration" 
     label2="\F2 Lateral position" unit2=km labelsz=10 titlesz=12.5
     parallel2=n label1="\F2 Time" format2="%3.1f"
     ''',local=1)
Result('elf-fmg2','fmg npk','SideBySideAniso',local=1)

# From migration velocity to Dix velocity

for case in ('npk','smb'):
    Flow(case+'-mir',case,'reverse which=2 opt=i')
    Flow(case+'-ext',[case+'-mir',case],
         'cat axis=2 ${SOURCES[1]} ${SOURCES[0]}')

Flow('vdix2 vmig2','npk-ext smb-ext',
     '''
     dix rect1=20 rect2=20
     weight=${SOURCES[1]} vrmsout=${TARGETS[1]}
     niter=40
     ''')

Flow('vdix1','vdix2',
     '''
     
     time2depth velocity=$SOURCE
     nz=640 dz=5.005 intime=y 
     ''')
Flow('vdx','vdix1','window n2=%d f2=%d | put o2=0 | transp' % (nx,nx))

Result('vels','vdix1',
       '''
       window n2=%d f2=%d | put o2=0 d2=0.01333333 d1=0.005005 | 
       scale dscale=0.001 |
       grey color=j scalebar=y allpos=y bias=1.4
       wanttitle=y barlabel="\F2 Velocity" barunit=km/s
       bartype=v label1="\F2 Depth" unit1=km 
       label2="\F2 Lateral position" unit2=km
       title="Dix Velocity Converted to Depth" barreverse=y
       screenht=%g screenwd=%g bartype=v labelsz=6 titlesz=6
       wanttitle=n barreverse=y parallel2=n 
       ''' % (nx,nx,0.001*zmax,0.001*xmax))

# Time to depth conversion

Flow('vxz x0 t0','vdix2',
     '''
     transp |
     cameron2d method=cheb neval=10 nc=1000 nz=%d dz=%g 
     x0=${TARGETS[1]} t0=${TARGETS[2]} |
     put label1=Lateral unit1=m label2=Depth unit2=m
     ''' % (nz,dz),local=1)

Flow('vxz2','vxz','window max2=%g n1=%d f1=%d | put o1=0' % (zmax,nx,nx))

Plot('vxz','vxz2',
     '''
     put o1=0 d1=0.01333333 d2=0.005005 | 
     scale dscale=0.001 |
     grey color=j scalebar=y allpos=y bias=1.4
     label2=Depth unit2=km label1="\F2 Lateral position" unit1=km
     wanttitle=y barlabel="\F2 Velocity" barunit=km/s label2="\F2 Depth"
     screenht=%g screenwd=%g bartype=v labelsz=6 titlesz=6
     title="Estimated Interval Velocity" barreverse=y
     wanttitle=n transp=n parallel2=n
     ''' % (0.001*zmax,0.001*xmax))

Flow('vd x z','vxz2',
     've2d x=${TARGETS[1]} z=${TARGETS[2]} nt=%d dt=%g' % (nt,dt))

Plot('rays','x z',
     '''
     cmplx ${SOURCES[1]} | transp | window j2=10 |
     graph wanttitle=n yreverse=y min1=0 max1=%g min2=0 max2=%g
     plotcol=7 scalebar=y wantaxis=n
     screenht=%g screenwd=%g bartype=v  
     ''' % (xmax,zmax,0.001*zmax,0.001*xmax))

Result('elfvxz','vxz rays','Overlay')

# Map PSTM to depth

Flow('coord','t0 x0',
     '''
     scale dscale=2 | cat axis=3 ${SOURCES[1]} |
     transp plane=13 
     ''')

Flow('fmg2','fmg coord',
     '''
     pad beg2=%d end2=%d | put o2=0 |
     inttest2 interp=spline nw=4 coord=${SOURCES[1]} |
     window n2=%d f2=%d
     ''' % (nx,nx,nx,nx))

Plot('fmg2',
     '''
     window max1=%g | agc | put o2=0 d2=0.01333333 d1=0.005005 | 
     grey title="\F2 Prestack migration: Time -> depth"
     label1="\F2 Depth" unit1=km label2="\F2 Lateral position" unit2=km
     screenht=%g screenwd=%g labelsz=6 titlesz=6 parallel2=n
     ''' % (zmax,0.001*zmax,0.001*xmax))

Result('fmg2','fmg2','Overlay')

# Post-stack depth migration

Result('elfstk','stk2',
       '''
       grey title=Stack
       label1=Time unit1=s label2="Lateral Position" unit2=m
       ''')

par = {
    'nx':nx,'ox':0,'dx':dx,
    'nt':nt,'ot':0,'dt':dt,
    'nz':351,'oz':0,'dz':dz,
    #
    'nw':201, 'ow':0, 'dw':0.625,'jw':1,
    #
    'verb':'y','eps':0.01,'nrmax':10
    }

nw=201

Flow('dat','stk2','fft1 | window n1=%d | transp plane=12 | transp plane=23' % nw)

for vel in ('vxz2','vdx'):
    slo = 'slo-'+vel
    Flow(slo,vel,
         '''
         transp plane=23 |
         math "output=1/input" |
         put label1=x label2=y label3=z
         ''')
    img = 'img-'+vel
    Flow(img,['dat',slo],
         '''
         zomig3 ompnth=1 mode=m --readwrite=y verb=y
         nrmax=10 slo=${SOURCES[1]}
         ''',
         split=[3,nw,[0]],reduce='add')

    Plot(img,
         '''
         window |
         transp | agc |
         put o2=0 d2=0.01333333 d1=0.005005 | 
         grey title="\F2 Post-stack depth migration"
         label1="\F2 Depth" unit1=km label2="\F2 Lateral position" unit2=km
         screenht=%g screenwd=%g labelsz=6 titlesz=6 parallel2=n
         ''' % (0.001*zmax,0.001*xmax),local=1)

Result('img','img-vxz2','Overlay',local=1)
Result('img0','img-vdx','Overlay',local=1)

End()