from rsfproj import *
import zomig,rwezo

par = {
    'nz':250,  'dz':0.04,    'oz':0,  'lz':'z', 'uz':'km',             # CC axis 1
    'nx':400,  'dx':0.04,    'ox':0,  'lx':'x', 'ux':'km',             # CC axis 2
    'nt':1800, 'dt':0.004,   'ot':0,  'lt':'"\F10 t\F3 "', 'ut':'s',   # RC axis 1
    'ng':225,  'dg':0.04,    'og':0.5, 'lg':'x', 'ug':'km',            # RC axis 2
    'nT':2400, 'dT':0.008,   'oT':0,      # time
    'tmx':10,'eps':0.0000001, 'verb':'y','nrmax':3,
    'ow':1, 'nw': 894, 'incore':'y',
    'prefix':'D-'
    }

# compute parameters
rwezo.param(par)

# velocity
Flow('vv',None,
     '''
     spike
     n1=%(nz)d o1=%(oz)g d1=%(dz)g
     n2=%(nx)d o2=%(ox)g d2=%(dx)g |
     math output="2.5+0.15*x1-0.20*x2" |
     scale rscale=-1 |
     add add=10 |
     clip clip=9.5 |
     add add=-10 |
     scale rscale=-1
     ''' % par)

Flow('va',None,
     '''
     spike nsp=2 mag=0,0
     n1=%(nz)d o1=%(oz)g d1=%(dz)g k1=100,200 l1=150,250
     n2=%(nx)d o2=%(ox)g d2=%(dx)g k2=100,100 l2=150,150 |
     smooth rect1=126 rect2=126 |
     smooth rect1=126 rect2=126 |
     scale axis=123
     ''' % par)
Flow('vb',None,
     '''
     spike nsp=1 mag=-0.25
     n1=%(nz)d o1=%(oz)g d1=%(dz)g k1=125 l1=150
     n2=%(nx)d o2=%(ox)g d2=%(dx)g k2=160 l2=170 |
     smooth rect1=41 rect2=41 |
     scale axis=123
     ''' % par)
Flow('vel','vv vb',
     '''
     math
     a=${SOURCES[0]}
     b=${SOURCES[1]}
     output="a+b"
     ''')
Plot('vel',rwezo.cgrey('allpos=y bias=1',par))

# smooth velocity for HWT
Flow('vsm','vv','smooth')

# ------------------------------------------------------------
# wavefronts by HWT

Flow('mzs',None,'spike nsp=1 mag=0 n1=%(ng)d d1=%(dg)g o1=%(og)g | math output="0"'  % par)
Flow('mxs',None,'spike nsp=1 mag=0 n1=%(ng)d d1=%(dg)g o1=%(og)g | math output="x1"' % par)
Flow('mrs',None,'spike nsp=1 mag=0 n1=%(ng)d d1=%(dg)g o1=%(og)g | math output="1"'  % par)
Flow('sou','mxs mzs','cat axis=2 space=n ${SOURCES[1]} | transp | dd type=complex | window')

Flow('hwt','vsm sou',
     '''
     hwtex verb=n sou=${SOURCES[1]}
     nt=%(nt)d ot=%(ot)g dt=%(dt)g
     ''' % par)

# coordinate system: rays(g,t)
Flow('cos','hwt','window min2=%(ot)g n2=%(nt)d' % par)

# coordinate system plot
rwezo.cos('cos',40,100,'plotcol=0',par)
Result(par['prefix']+'cos' % par,'vel cos','Overlay')

# slowness
rwezo.slow('sloCC','sloRC','vel','cos',par)
Result(par['prefix']+'sloRC','sloRC','transp |' % par
       + rwezo.rgrey('pclip=100 allpos=y bias=0.12',par))

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

# ------------------------------------------------------------
# time data CC: (t,x)
Flow('datCC',None,
     '''
     spike nsp=3 mag=1,1,1
     n1=%(nT)d o1=%(oT)g d1=%(dT)g k1=500,1000,1500
     n2=%(nx)d o2=%(ox)g d2=%(dx)g k2=75 l2=225 |
     ricker1 frequency=5 |
     smooth rect2=51 |
     put label1=t label2=x
     ''' % par )
Result(par['prefix']+'datCC' % par,'datCC','grey pclip=100 wanttitle=n')

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

# ------------------------------------------------------------
# WEM image CC: (z,x)
zomig.image('imgCC','sloCC','frqCC',par)

# ------------------------------------------------------------
# RWE image
rwezo.mig('migCC','migRC','frqRC','abmRC','abrRC','cos',par)
# rwezo.mod('modCC','modRC','migRC','abmRC','abrRC','cos',par)

# ------------------------------------------------------------
# PLOTS
rwezo.plots(par)

Plot('imgCC-bp','imgCC','window | transp | bandpass fhi=3|'
     + rwezo.cgrey('pclip=100',par))
Result(par['prefix']+'imgCC-bp',['imgCC-bp','cos'],'Overlay')

End()