from rsf.proj import *
from rsf.recipes import rwezo,zomig,fdmod

par = {
    'nz':1201, 'dz':0.025,  'oz':0,     'lz':'z', 'uz':'km',
    'nx':2133, 'dx':0.0375, 'ox':10.925,'lx':'x', 'ux':'km',
    'nt':30000,'dt':0.0005, 'ot':0,     'lt':'\F10 t\F3', 'ut':'s',
    'ng':800,  'dg':1,      'og':0,     'lg':'\F10 g\F3', 'ug': '.',
    'kt':100,    # wavelet delay 
    'nT':1600, 'dT':0.004,  'oT':0,
    'ow':1, 'nw': 795,
    'nbz':250,'nbx':250, 'j2': 5
    }

par['ox']*=0.3048
par['dx']*=0.3048
par['oz']*=0.3048
par['dz']*=0.3048

# compute parameters
rwezo.param(par)
zomig.migpar(par)
fdmod.param(par)

# ------------------------------------------------------------
# velocity
velo = 'sigsbee2a_migvel.sgy'
Fetch(velo,'sigsbee')

Flow('zvelo tzvelo zvelo.asc zvelo.bin',velo,
     '''
     segyread
     tfile=${TARGETS[1]}
     hfile=${TARGETS[2]}
     bfile=${TARGETS[3]}
     ''')

# ------------------------------------------------------------
par['nxcut']=1600
par['nxpad']=2000
par['nzpad']=799

Flow('vraw','zvelo',
     '''
     scale rscale=0.001 |
     scale rscale=0.3048 |
     put d1=%(dz)g o2=%(ox)g d2=%(dx)g label1=z label2=x |
     window n2=%(nxcut)d
     ''' % par)

par['ox']=par['ox']-par['nxpad']*par['dx']
par['nx']=par['nxcut']+par['nxpad']
par['nz']=par['nz']+par['nzpad']

par['xmin']=par['ox']
par['xmax']=par['ox'] + (par['nx']-1) * par['dx']
par['zmin']=par['oz']
par['zmax']=par['oz'] + (par['nz']-1) * par['dz']

Flow('padx','vraw',
     '''
     window n2=1 |
     spray axis=2 n=%(nxpad)d d=%(dx)g o=%(ox)g |
     cat ${SOURCES[0]} axis=2 space=n
     ''' % par)

Flow('padz','padx',
     '''
     window n1=1 f1=1200 |
     spray axis=1 n=%(nzpad)d d=%(dz)g o=0 |
     math "output=input+x1*0.25"
     ''' % par)

Flow('velo','padx padz',
     '''
     cat axis=1 space=n ${SOURCES[1]}
     ''')
Flow('dens','velo','math output=2.5')

Plot  ('velo',rwezo.cgrey('allpos=y pclip=95',par))
Result('velo',rwezo.cgrey('allpos=y pclip=95',par))

Flow(  'salt','velo','add add=-4.26 | clip clip=0.3048')
#Result('salt',rwezo.cgrey('allpos=y pclip=100',par))

# ------------------------------------------------------------

Flow(  'edgez','salt','         igrad square=y')
Flow(  'edgex','salt','transp | igrad square=y | transp')
Flow(  'edge','edgez edgex',
       'math z=${SOURCES[0]} x=${SOURCES[1]} output="z+x"')
Result('edge','sfdd type=float |' + rwezo.cgrey('pclip=100',par))

Flow(  'mask','edge','mask min=0.0006')
Result('mask','sfdd type=float |' + rwezo.cgrey('pclip=100',par))

for x in ('x1','x2'):
    if(x=='x1'): o='zs'
    if(x=='x2'): o='xs'

    Flow(o,'edge mask',
         '''
         math output=%s |
         put n1=1 n2=%d |
         headerwindow mask=${SOURCES[1]} |
         window
         ''' % (x,par['nz']*par['nx']))

Flow('rs','edge mask',
     '''
     scale axis=123 |
     put n1=1 n2=%d |
     headerwindow mask=${SOURCES[1]} |
     window
     ''' % (par['nz']*par['nx']))

Flow('ss',['xs','zs'],
     '''
     cat axis=2 space=n
     ${SOURCES[0]} ${SOURCES[1]} | transp |
     window j2=2
     ''', stdin=0)

Plot(  'ss',fdmod.ssplot('',par))
Result('ss',['velo','ss'],'Overlay')

# ------------------------------------------------------------

# ------------------------------------------------------------
# receiver (flat)

fdmod.horizontal('oo',par['oz'],par)
Plot(  'oo',fdmod.rrplot('',par))
Result('oo',['velo','oo'],'Overlay')

Result('model','velo ss oo','Overlay')

Flow('surf',None,
     '''
     spike nsp=1 mag=1 n1=%(nz)d d1=%(dz)g o1=%(oz)g k1=2 |
     spray axis=2       n=%(nx)d  d=%(dx)g  o=%(ox)g
     ''' % par)

# ------------------------------------------------------------
# wavelet
fdmod.wavelet('wav_',20,par)
Flow(  'wav', 'wav_','transp')
Result('wav','window n2=1000 |' + fdmod.waveplot('',par))

# ------------------------------------------------------------
# modeling
fdmod.awefd1('do','wo','wav','velo','dens','ss','oo','free=n',par)

# ------------------------------------------------------------

Plot('wo',
     'window j3=20 min1=%(oz)g n1=%(nz)d min2=%(ox)g n2=%(nx)d |' % par
     + rwezo.cgrey('pclip=99 gainpanel=a',par),view=1)
Plot('do','window j2=4 | transp |' + rwezo.dgrey('',par),view=1)

for j in range(20,300,20):
    par['j'] = j

    Plot('wo'+str(j//20),'wo',
         '''
         window n3=1 f3=%(j)d
         min1=%(oz)g n1=%(nz)d min2=%(ox)g n2=%(nx)d |
         ''' % par
         + rwezo.cgrey('pclip=99',par))
    Result('wo'+str(j//20),['wo'+str(j//20),'cos1','ss'],'Overlay')

Flow('dd','do','window f2=%(kt)d | put o2=0 d2=0.001 | window j2=4' % par)
Plot('dd','transp |' + rwezo.dgrey('',par),view=1)

# ------------------------------------------------------------

# ------------------------------------------------------------
Flow('vel1','velo','smooth rect1=351 rect2=351 repeat=3')

Flow('s_',None,'spike nsp=1 n1=%(ng)d d1=%(dg)g o1=%(og)g' % par)

# coordinate system: rays(g,t)
Flow('mzs1','s_','math output="01.0+0.009*x1"' % par)
Flow('mxs1','s_','math output="17.0+0.008*x1"' % par)
Flow('sou1','mxs1 mzs1',
     'cat axis=2 space=n ${SOURCES[1]} | transp | dd type=complex | window')

# wavefronts (by HWT)
Flow('hwt1','vel1 sou1',
     '''
     hwtex verb=n sou=${SOURCES[1]}
     nt=%(nt)d ot=%(ot)g dt=%(dt)g
     ''' % par)
Flow('cos1','hwt1',
     '''
     reverse which=2 |
     put o2=0 d2=%(dt)g |
     window j1=2 j2=8
     ''' % par)

# ------------------------------------------------------------
Flow('vel2','velo','math output=2.75')

# coordinate system: rays(g,t)
Flow('mzs2','s_','math output="-06.5+0.025*x1"' % par)
Flow('mxs2','s_','math output="-17.0-0.005*x1"' % par)
Flow('sou2','mxs2 mzs2',
     'cat axis=2 space=n ${SOURCES[1]} | transp | dd type=complex | window')

# wavefronts (by HWT)
Flow('hwt2','vel2 sou2',
     '''
     hwtex verb=n sou=${SOURCES[1]}
     nt=%(nt)d ot=%(ot)g dt=%(dt)g
     ''' % par)
Flow('cos2','hwt2','window j2=4')

# ------------------------------------------------------------
# time data CC: (t,x)
Flow  ('datCC','dd','transp memsize=1000')
Result('datCC','window j1=5 j2=5 | grey pclip=99 wanttitle=n label1=t unit1=s label2=x unit2=km')

# ------------------------------------------------------------
# WEM image CC: (z,x)
zomig.image3('imgCC','sloCC','frqCC',par)
Plot  ('imgCC','window | transp |' + rwezo.cgrey('pclip=99',par))
Result('imgCC',['imgCC','cos1'],'Overlay')

# loop over coordinate systems
# 1 - Riemannian coordinate system
# 2 - Cartesian  coordinate system
for i in ('1','2'):
    # coordinate system plot
    rwezo.cos(   'cos'+i,20,100,'',par)
    Result('cos'+i,['velo','cos'+i],'Overlay')

    # slowness
    rwezo.slow(  'sloCC','sloRC'+i,'velo','cos'+i,par)
    Result('sloRC'+i,'window j2=5 | transp |'
           + rwezo.rgrey('pclip=99 allpos=y color=j bias=0.061',par))

    # surface and salt
    Result('edgeRC'+i,['edge','surf','cos'+i],
           '''
           add ${SOURCES[1]} | transp |
           c2r rays=${SOURCES[2]} adj=n | transp | bandpass fhi=10 |
           scale axis=123 | window j1=5 |
           '''+ rwezo.rgrey('pclip=99 scalebar=y',par))

    # ------------------------------------------------------------
    # coordinate system coefficients (A,B,M=mask) and references (A,B)
    rwezo.abm(    'abmRC'+i,'abrRC'+i,'sloRC'+i,'cos'+i,par)
    rwezo.abmplot('abmRC'+i,par)

    # ------------------------------------------------------------
    # frequency data RC: (g,w)
    rwezo.frq('frqRC'+i,'frqCC','datCC','cos'+i,par)

    # ------------------------------------------------------------
    # RWE image
    rwezo.mig('migCC'+i,'migRC'+i,'frqRC'+i,'abmRC'+i,'abrRC'+i,'cos'+i,par)

# ------------------------------------------------------------
# PLOTS
#plots(par)

# ------------------------------------------------------------
zoompar = {
    'nz':7,    'dz':1,      'oz':2,     'lz':'z', 'uz':'km',
    'nx':10,   'dx':1,      'ox':12,    'lx':'x', 'ux': 'km',
    'nt':30000,'dt':0.0005, 'ot':0,     'lt':'\F10 t\F3 ', 'ut': 's',
    'ng':400,  'dg':1,      'og':0,     'lg':'\F10 g\F3 ', 'ug': '.',
    'nT':1600, 'dT':0.004,  'oT':0
    }
rwezo.param(zoompar)

Plot('zoomvel','velo',rwezo.cgrey('',zoompar))
Flow('vscaleCC','sloCC',
     'window | transp | math output=1/input | scale axis=123 | add add=-0.6')

for i in ('1','2'):
    Flow('vscaleRC'+i,'sloRC'+i,
         'transp | math output=1/input | scale axis=123 | add add=-0.6')

    Flow(     'zoomcos'+i,'cos'+i,'window')
    rwezo.cos('zoomcos'+i,5,25,'',zoompar)
#    Result(   'zoomcos'+i,  ['zoomvel','zoomcos'+i],'Overlay')
    Plot('zoomvel-ovl'+i,['zoomvel','zoomcos'+i],'Overlay')

    for j in (['SSF','FFD','PSC',
               'F15','F45','F60']):
        sfx = i + '-' + j

        Flow('iscaleCC'+sfx,'migCC'+sfx,'transp | scale axis=123')
        Flow('compCC'+sfx,['vscaleCC','iscaleCC'+sfx],
             'math a=${SOURCES[0]} b=${SOURCES[1]} output="0.5*a+b"')
        Plot(  'compCC'+sfx,rwezo.cgrey('pclip=100',zoompar))
        Result('compCC'+sfx,['compCC'+sfx,'zoomcos'+i],'Overlay')

        Flow('iscaleRC'+sfx,'migRC'+sfx,'transp | scale axis=123')
        Flow('compRC'+sfx,['vscaleRC'+i,'iscaleRC'+sfx],
             'math a=${SOURCES[0]} b=${SOURCES[1]} output="0.5*a+b"')
        Result('compRC'+sfx,rwezo.rgrey('pclip=100 min1=10',zoompar))

End()