from rsf.proj import *


## Kirchhoff modeling
Flow('modl',None,
     '''
     spike n1=200 o1=-8 d1=0.08 n2=4
     nsp=4 k2=1,2,3,4 mag=1,2.2,3.5,5
     ''')
Flow('refl',None,
     '''
     spike n1=200 n2=4 nsp=4 k2=1,2,3,4
     mag=0.0308091,0.0638570,0.1011110,0.19000000
     ''')
Flow('mod1','modl','window min1=0')

Flow('unif','mod1','unif2 n1=110 d1=0.05 v00=5,6,8,10,15')


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


Flow('clean','modl refl',
     '''
     kirmod nt=375 dt=0.004 freq=25 refl=${SOURCES[1]}
     ns=12 s0=1.4 ds=-0.1
     nh=64 h0=0.1 dh=0.1
     vel=5 gradz=2 type=v
     ''')

Flow('data','clean',
     '''
     noise seed=1 type=y mean=0 var=0.00000004
     ''')

shots = []
for shot in range(0,12,5):
    data = 'data%d' % shot
    shots.append(data)
    Plot(data,'data',
         '''
         window n3=1 f3=%d |
         wiggle poly=y yreverse=y transp=y title="Shot %d"   
         ''' % (shot,shot+1))
Result('data',shots,'SideBySideAniso',vppen='txscale=2.3')

v0=4
dv=0.075
nv=120


Flow('off','data','window n1=1 | math output="0.5*x1" ')
Flow('cmp','off','math output="input+x2" ')



#####################
# Convert shot to cmp
#####################
Flow('cmpdata cmpmask','data',
     '''
     reverse which=4 |
     put o2=0.05 d2=0.05 |
     shot2cmp mask=${TARGETS[1]} 
     ''')
#####################
# Velocity scan
#####################
Flow('scan','cmpdata cmpmask',
     '''
     vscan semblance=y mask=${SOURCES[1]}
     v0=%g dv=%g nv=%d
     ''' % (v0,dv,nv))
#####################
# Velocity picking
#####################
Flow('pick','scan',
     '''
     mutter x0=4 inner=y v0=3 half=n |
     mutter x0=4 inner=n v0=8 half=n |
     pick rect1=50 rect2=100 vel0=5 | window
     ''')

#####################
# NMO and initial stack
#####################
Flow('nmo','cmpdata cmpmask pick',
     'nmo mask=${SOURCES[1]} velocity=${SOURCES[2]}')
Flow('stack','nmo','stack')

#############################################one cmp 
Flow('onecmp','cmpdata','window f3=23 n3=1 ')

Flow('onenmo', 'nmo', 'window f3=23 n3=1 ')

Plot('onecmp', 'onecmp',
       '''window f2=0 n2=12 | put o2=1 d2=1 label2=Trace unit2="" | sfwiggle poly=y yreverse=y transp=y title='a)' 
       ''')

Plot('onenmo', 'onenmo',
       '''window f2=0 n2=12 | put o2=1 d2=1 label2=Trace unit2="" | sfwiggle poly=y yreverse=y transp=y title='b)' 
       ''')

Flow('onenmostack','onenmo','stack')

Flow('onenmostacks','onenmostack',
     '''
     spray axis=2 n=32
     ''')


Flow('weight1','onenmostacks onenmo',
     '''
     similarity other=${SOURCES[1]} niter=30 rect1=5 rect2=7
     ''')

Flow('tweight1','weight1','threshold pclip=20')

Flow('onesimistack', 'tweight1 onenmostacks onenmo',
     '''
     sfmath y=${SOURCES[1]} output=input*y | stack

     ''')

Flow('norweight1','tweight1','stack')

Flow('noronesimistack','onesimistack norweight1','sfmath y=${SOURCES[1]} output="input/(y+1)"')

Flow('snrstack','onenmo','snrstack w=15 ee=1e-1 esp=1000 ')
Plot('snrstack1','snrstack','put o2=0 label2=Amplitude unit2="" | wiggle poly=y yreverse=y transp=y title="d)"')

Plot('onenmostack1','onenmostack','put o2=0 label2=Amplitude unit2="" | wiggle poly=y yreverse=y transp=y title="c)"')

Plot('onesimistack1','onesimistack','put o2=0 label2=Amplitude unit2="" |wiggle poly=y yreverse=y transp=y title="e)"')

Result('onestack1',' onecmp onenmo onenmostack1  snrstack1 onesimistack1', 'SideBySideAniso',vppen='txscale=1.8')

#######################################all cmps
Flow('wstack','stack', 'window min2=1 max2=3.8')
Plot('wstack','wstack',
       '''wiggle transp=y poly=y yreverse=y  title='a)' 
       ''')
Flow('stacks', 'stack','spray axis=3 n=32 | transp plane=23')
Flow('wstacks','stacks', 'window min3=1 max3=3.8')
Flow('wnmo','nmo', 'window min3=1 max3=3.8 ')
Flow('weights', 'wnmo wstacks',
     '''similarity other=${SOURCES[1]} niter=30 rect1=5 rect2=7  
     ''')



Flow('simistack','weights wnmo',
     '''
     sfmath y=${SOURCES[1]} output=input*y | stack  
     ''')

# Thresholding
Flow('tweights','weights','threshold pclip=20')

Flow('tnormal2','tweights','stack | stack axis=1 | spray n=375 | transp | put o1=0 d1=0.004')

Flow('tsimistack1','tweights wnmo',
     '''
     sfmath y=${SOURCES[1]} output=input*y | stack 
     ''')
Flow('tsimistack','tsimistack1 tnormal2','sfmath y=${SOURCES[1]} output=input/y')

Plot('tsimistack',
       '''wiggle transp=y poly=y yreverse=y  title='c)' 
       ''')
Flow('snrstacks','wnmo','snrstack w=15 ee=1e-1 esp=1000 ')

Plot('snrstacks',
       ''' put label2="Midpoint"  unit2="km" | wiggle transp=y poly=y yreverse=y s title='b)' 
       ''')

Result('stackss','wstack snrstacks tsimistack', 'SideBySideAniso',vppen='txscale=1.8')

End()