import sys
import rsf.recipes.tau as tau

from rsf.proj import *

ft = .3048
nx = 751;  x0 = 0.; dx = 20. * ft
nz = 601;  z0 = 0.; dz = 20. * ft
nT = 501;  T0 = 0.; dT = .00324
nt = 5001; t0 = 0.; dt = .0005

b = [40   ,20   ,20   ,20   ]
c = [4.e-5,1.e-4,1.e-4,1.e-4]

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

freq = 25

spk0= 1 + int(1./(freq*dt))
dspk= int(.5/ dt)
spk = range(spk0,spk0+3*dspk,dspk)

par = {'nx' : nx, 'x0' : x0, 'dx' : dx,
           'nz' : nz, 'z0' : z0, 'dz' : dz,
           'nT' : nT, 'T0' : T0, 'dT' : dT,
           'nt' : nt, 't0' : t0, 'dt' : dt}

window  = '''window n1=%d f1=0 n2=%d f2=0 |
put o1=%g d1=%g label1=Depth unit1=m
        o2=%g d2=%g label2=Distance unit2=m
''' % (nz,nx,z0,dz,x0,dx)

for m in Split('vp delta epsilon'):
    sgy = 'timodel_%s.segy' % m
    sgyz = sgy + '.gz'
    Fetch(sgyz,dir='Hess_VTI',server='ftp://software.seg.org',
          top='pub/datasets/2D')
    Flow(sgy,sgyz,'zcat $SOURCE',stdin=0)

Flow('vverC','timodel_vp.segy','''
segyread tfile=/dev/null verb=y | scale rscale=.3048 |
''' + window)
Flow('epsiC','timodel_epsilon.segy','''
segyread tfile=/dev/null verb=y |
''' + window)
Flow('deltC','timodel_delta.segy','''
segyread tfile=/dev/null verb=y |
''' + window)

Flow('vmapC','vverC','smooth rect1=8 rect2=8 repeat=2')

Flow('data',None,'''
spike nsp=3 mag=1 n1=%d o1=0 d1=%g k1=%d,%d,%d |
ricker1 frequency=%g | scale axis=1 | reverse which=1 opt=i |
transp | put label1=Distance unit1=m
''' % (nt,dt,spk[0],spk[1],spk[2],freq))

tau.init(par)

tau.compute_vnmo('vnmoC','vverC','deltC')
tau.compute_eta('hetaC','epsiC','deltC')

tau.compute_tau('tauC','vmapC')
tau.compute_sigma('sigmC','tauC')
tau.coord_xmid('cr')

tau.mapping('vverT','vverC','tauC',inv=False)
tau.mapping('vnmoT','vnmoC','tauC',inv=False)
tau.mapping('hetaT','hetaC','tauC',inv=False)
tau.mapping('vmapT','vmapC','tauC',inv=False)
tau.mapping('sigmT','sigmC','tauC',inv=False)
tau.mapping('vverB','vverT','tauC',inv=True)
tau.mapping('vnmoB','vnmoT','tauC',inv=True)
tau.mapping('hetaB','hetaT','tauC',inv=True)
tau.mapping('vmapB','vmapT','tauC',inv=True)

tau.rtm_vti('imagC','waveC','data','vverC','vnmoC','hetaC',None   ,None   ,'cr',n3,b,c,tau=False)
tau.rtm_vti('imagT','waveT','data','vverT','vnmoT','hetaT','vmapT','sigmT','cr',n3,b,c,tau=True)

m2kmC = 'put d2=%g unit2=km d1=%g unit1=km' % (.001*dx,.001*dz)
m2kmT = 'put d2=%g unit2=km' % (.001*dx)
screen = 'screenwd=12 screenht=10 labelsz=10'
contt = 'contour allpos=y plotcol=7 plotfat=5 wantaxis=n wanttitle=n c0=0 dc=.12 ' + screen
contx = 'contour allpos=y plotcol=7 plotfat=5 wantaxis=n wanttitle=n c0=0 dc=%g ' % (45*dx) + screen
vgrey = 'grey color=i title= allpos=y ' + screen
agrey = 'grey color=i title= allpos=y ' + screen
wgrey = 'grey grid=y title= clip=.4 ' + screen
vgrey2= 'scale rscale=.001 | ' + vgrey + ' scalebar=y barreverse=y barlabel=Velocity barunit="km/s"'
agrey2= agrey + ' scalebar=y barreverse=y barlabel=Anellipticity barunit=1'

Result('hesm1','vverC',m2kmC + ' | ' + vgrey2)
Result('hesm2','vnmoC',m2kmC + ' | ' + vgrey2)
Result('hesm3','hetaC',m2kmC + ' | ' + agrey2)

Plot('gridC1','tauC',contt)
Plot('gridC2','tauC','math output=x2 | ' + contx)
Plot('gridC3','vverC',m2kmC + ' | ' + vgrey)

Plot('gridT1','vverT','math output=x1 | ' + contt)
Plot('gridT2','vverT','math output=x2 | ' + contx)
Plot('gridT3','vverT',m2kmT + ' | ' + vgrey)

Result('hesgC','gridC3 gridC2 gridC1','Overlay')
Result('hesgT','gridT3 gridT2 gridT1','Overlay')

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

Flow('imagC1','imagC','scale axis=2')
Flow('imagT1','imagT','scale axis=2')
Flow('imagB1','imagB','scale axis=2')

Result('hesiC','imagC1',m2kmC + ' | ' + wgrey)
Result('hesiT','imagT1',m2kmT + ' | ' + wgrey)
Result('hesiB','imagB1',m2kmC + ' | ' + wgrey)

Result('hesiE','imagC1 imagB1','add scale=-1,1 ${SOURCES[1]} | ' + m2kmC + ' | ' + wgrey)

End()