from rsf.proj import *
import math
sys.path.append('..')
import tau

nx = 501; x0 = 0.; dx = 5.
nz = 401; z0 = 0.; dz = 5.
nZ = 401; Z0 = 0.; dZ = .0025 # equal amplitude
nT = 301; T0 = 0.; dT = .003  # attenuated amplitude
nt = 1601;t0 = 0.; dt = .0005

b = [0,0,0,0]
c = [0,0,0,0]

j3 = 25; n3 = 1 + int((nt-1)/j3)

freq = 25

parZ = {'nx' : nx, 'x0' : x0, 'dx' : dx,
            'nz' : nz, 'z0' : z0, 'dz' : dz,
            'nT' : nZ, 'T0' : Z0, 'dT' : dZ,
            'nt' : nt, 't0' : t0, 'dt' : dt}
parT = {'nx' : nx, 'x0' : x0, 'dx' : dx,
            'nz' : nz, 'z0' : z0, 'dz' : dz,
            'nT' : nT, 'T0' : T0, 'dT' : dT,
            'nt' : nt, 't0' : t0, 'dt' : dt}

epsilon = .32
delta   = .1
velz    = 2000
veln    = velz * math.sqrt(1. + 2.*delta)
eta     = (epsilon - delta) / (1. + 2.*delta)

xr = .5 * (nx-1) * dx
zr = .5 * (nz-1) * dz
Tr = zr / velz

Flow('vverC',None,'''
math n1=%d o1=0 d1=%g n2=%d o2=0 d2=%g output="%g" |
put label1=Depth unit1=m label2=Distance unit2=m
''' % (nz,dz,nx,dx,velz))
Flow('vnmoC','vverC','math output="%g"' % veln)
Flow('hetaC','vverC','math output=%g' % eta)

Flow('data',None,'''
spike nsp=1 mag=1 n1=%d o1=0 d1=%g k1=%d |
ricker1 frequency=%g | scale axis=1 | reverse which=1 opt=i| 
transp | put label1=Distance unit1=m
''' % (nt,dt,1+int(1./(freq*dt)),freq))

tau.init(parZ)

Flow('tauC','vverC','math output="x1/input"')

Flow('crC crC.asc',None,'''
echo %g %g > ${TARGETS[1]} &&
echo n1=2 in=${TARGETS[1]} data_format=ascii_float |
dd form=native
''' % (zr,xr))

Flow('crT crT.asc',None,'''
echo %g %g > ${TARGETS[1]} &&
echo n1=2 in=${TARGETS[1]} data_format=ascii_float |
dd form=native
''' % (Tr,xr))

tau.mapping('vverZ','vverC','tauC',inv=False)
tau.mapping('vnmoZ','vnmoC','tauC',inv=False)
tau.mapping('hetaZ','hetaC','tauC',inv=False)
tau.mapping('vverB','vverZ','tauC',inv=True)
tau.mapping('vnmoB','vnmoZ','tauC',inv=True)
tau.mapping('hetaB','hetaZ','tauC',inv=True)

Flow('sigmZ','vverZ','math output=0')

tau.rtm_vti('imagC','waveC','data','vverC','vnmoC','hetaC',None   ,None   ,'crC',n3,b,c,tau=False)
tau.rtm_vti('imagZ1','waveZ','data','vverZ','vnmoZ','hetaZ','vverZ','sigmZ','crT',n3,b,c,tau=True)

tau.mapping('imagZ','imagZ1','tauC',inv=True)

tau.init(parT)

tau.mapping('vverT','vverC','tauC',inv=False)
tau.mapping('vnmoT','vnmoC','tauC',inv=False)
tau.mapping('hetaT','hetaC','tauC',inv=False)

Flow('sigmT','vverT','math output=0')

tau.rtm_vti('imagT1','waveT','data','vverT','vnmoT','hetaT','vverT','sigmT','crT',n3,b,c,tau=True)

tau.mapping('imagT','imagT1','tauC',inv=True)

Plot('waveC','grey gainpanel=a')
Plot('waveZ','grey gainpanel=a')
Plot('waveT','grey gainpanel=a')

m2kmC = 'put d2=%g unit2=km d1=%g unit1=km' % (.001*dx,.001*dz)
m2kmT = 'put d1=%g unit1=km' % (.001*dx)
screen = 'screenwd=10 screenht=10 labelsz=8'
window = 'window n1=1 f1=200'
wgrey = 'grey grid=y title= labelsz=10 pclip=98 ' + screen
wline = 'graph wantaxis=n wanttitle=n plotfat=8 ' + screen

for c in ['C','Z','T']:
        imag = 'imag' + c
        line = 'line' + c
        arte = 'arte' + c
        Plot(imag,m2kmC + ' | ' + wgrey)
        Plot(line + '_',imag,window + ' | ' + wline)
        Plot(line,line + '_','Overlay',vppen='yscale=.5 yshift=5.5')

Result('arteC','imagC lineC','OverUnderIso')
Result('arteZ','imagZ lineZ','OverUnderIso')
Result('arteT','imagT lineT','OverUnderIso')


End()