from rsf.proj import *

Fetch('elf-stk2.rsf','masha')
Flow('data','elf-stk2.rsf','dd form=native | put o1=0')

def grey(title):
    return '''
    window n1=751 |
    grey  pclip=99 title="%s" label2=" " crowd=0.85
    ''' % title

def zoom(title):
    return '''
    window min1=0.81 max1=1.60 min2=4800 max2=7600 |
    grey  pclip=99 title="%s" label2=" " crowd=0.8 grid=y
    ''' % title

Plot('data',grey('(a) DMO-stack data'))

Flow('fk','data','cosft sign2=1')
Flow('stolt','fk','stolt vel=2000 pad=2049 minstr=0.5 | cosft sign2=-1')

Plot('stolt',grey('(b) Stolt migration (v0 = 2000 m/s)'))

Fetch('vsmooth_levset.dat','masha')

nx=1000
dx=13.3333

nz=1000
dz=5.005

Flow('vel','vsmooth_levset.dat',
     '''
     echo in=$SOURCE
     n1=%d o1=0 d1=%g
     n2=%d o2=0 d2=%g
     label1=Lateral unit1=m
     label2=Depth unit2=m
     label=Velocity unit="km/s"
     data_format=ascii_float |
     dd form=native |
     transp 
     ''' % (nx,dx,nz,dz),stdin=0)

Flow('vt','vel','depth2time velocity=$SOURCE nt=800 dt=0.004 t0=0')
Flow('bot','vt','window n1=1 f1=650 | spray axis=1 n=150')
Flow('vt2','vt bot','window n1=650 | cat axis=1 ${SOURCES[1]}')
Flow('vt1','vt2','stack | spray axis=2 n=1000 d=13.3333 o=0')

v0 = 1800

Flow('str','data vt1',
     'stoltstretch velocity=${SOURCES[1]} vel=%g pad=1250' % v0)
Flow('stfk','str','cosft sign2=1')

par = {
    'good': '',
    'bad': 'stretch=0.5'
    }

for case in par.keys():
    stolt = 'stolt-'+case
    Flow(stolt,'stfk vt1',
         '''
         stolt vel=%g pad=2049 minstr=0.5 %s |
         cosft sign2=-1 |
         stoltstretch velocity=${SOURCES[1]} vel=%g inv=y
         ''' % (v0,par[case],v0))

Plot('stolt-bad',grey('(c) Stolt-stretch migration (W=0.5)'))
Plot('stolt-bad-zoom','stolt-bad',
     zoom('(c) Stolt-stretch migration (W=0.5)'))
Plot('stolt-good',grey('(a) Stolt-stretch migration (optimal W)'))
Plot('stolt-good-zoom','stolt-good',
     zoom('(a) Stolt-stretch migration (optimal W)'))

Flow('pshift','fk vt1',
     'gazdag velocity=${SOURCES[1]} pad=2048 | cosft sign2=-1')
Plot('pshift',grey('(b) Phase-shift migration'))
Plot('pshift-zoom','pshift',
     zoom('(b) Phase-shift migration'))

ntcut = [0,125,275,350,400,800]
vref = [1949,1678,1751,1377,627.9]
ncut = len(ntcut)-1

Flow('casc','vt1','cascade ncut=%d ntcut=%s' %
     (ncut-1,string.join(map(str,ntcut[1:ncut]),',')))

Flow('casc2','vt1','cascade ncut=2 ntcut=200,400')

Plot('vt1',
     '''
     window n2=1 n1=751 | graph crowd=0.8
     transp=y yreverse=y xinch=7 titlesz=14 labelsz=10 
     title="(a)" label1=Time unit1=s label2=Velocity unit2=m/s
     wheretitle=b wherexlabel=t
     ''')

for case in ('casc','casc2'):
    Plot(case,['vt1',case],
         '''
         transp plane=23 |
         cat axis=2 ${SOURCES[1]} |
         window n3=1 n1=751 | graph crowd=0.8
         transp=y yreverse=y titlesz=14 labelsz=10 
         title="(%c)" label1=Time unit1=s label2=Velocity unit2=m/s
         wheretitle=b wherexlabel=t dash=0,1,2,3,4,5
         ''' % 'bc'[case=='casc'])

Result('velocities','vt1 casc2 casc','SideBySideAniso')

data = 'data'
wins = []
migs = []

for ic in range(ncut):
    # migrate
    migr = 'migr%d' % ic
    vel = 'vel%d' % ic
    v0 = vref[ic]

    Flow(vel,'casc','window n2=1 f2=%d' % ic)
    Flow(migr,[data,vel],
         '''
         stoltstretch velocity=${SOURCES[1]} vel=%g pad=1250 |
         cosft sign2=1 |
         stolt vel=%g pad=2049 minstr=0.5 |
         cosft sign2=-1 |
         stoltstretch velocity=${SOURCES[1]} vel=%g inv=y
         ''' % (v0,v0,v0))
    migs.append(migr)
    data = migr

    # window
    wind = 'wind%d' % ic
    wins.append(wind)
    Flow(wind,'data',
         'math output=%d | window f1=%d n1=%d' %
         (ic,ntcut[ic],ntcut[ic+1]-ntcut[ic]))

Flow('slice',wins,
     '''
     cat axis=1 ${SOURCES[1:%d]} | 
     math output=input-0.01 | 
     smooth rect1=50
     ''' % ncut)
Flow('cmig',migs+['slice'],
     '''
     cat ${SOURCES[1:%d]} |
     put o3=0 d3=1 |
     transp plane=23 |
     slice pick=${SOURCES[%d]} |
     window
     ''' % (ncut,ncut))

Plot('cmig',grey('(c) Cascaded Stolt-Stretch (%d velocities)' % ncut))
Plot('cmig-zoom','cmig',
     zoom('(d) Cascaded Stolt-Stretch (%d velocities)' % ncut))

Result('data-stolt-ststr','data stolt stolt-bad','OverUnderAniso',
       vppen='xscale=1.7')

Result('data-ststr-pshift-casc','stolt-good pshift cmig','OverUnderAniso',
       vppen='xscale=1.7')

Result('dip-zoom','stolt-good-zoom pshift-zoom stolt-bad-zoom cmig-zoom',
       'TwoRows')

End()