# . . Set up generic project files
from rsf.proj import *
from rsf.recipes import fdmod,encode,wemig,stiffness
import os
from rsf.gallery import french
french.get_refl('refl')

# . . Set parameter for the modelling.  Applies for all situations.
par = {
    # Space parameters
    'nx':161,  'ox':0, 'dx':0.010,  'lx':'x', 'ux':'km',
    'ny':161,  'oy':0, 'dy':0.010,  'ly':'y', 'uy':'km',
    'nz':161,  'oz':0, 'dz':0.010,  'lz':'z', 'uz':'km',
    'kz':50, 'lz':100, 'kx':1, 'lx':100, 'ky':1, 'ly':100,

    # Time Parameters
    'nt':300,'ot':0, 'dt':0.002,  'lt':'t', 'ut':'s',
    'kt':50,'frq':35,

    # Modelling parameters
    'snap':'y','jsnap':4,'nb':0,'cb':2,'verb':'y',
    'nbell':2, 'jdata':1,'ssou':'2',

    # Output
    'height':5,
    }

# . . Initialize parameters in fdm module   
fdmod.param(par)
par['nframe']=10
par['iframe']=4
par['dabc']='n'


# --------------------------------------------------------------------
# . . Source Section
#
# . . Wavelet
#fdmod.wavelet('wav_',par['frq'],par)
Flow('wav_',None, 'spike n1=%(nt)d d1=%(dt)f k1=%(kt)d | ricker1 frequency=%(frq)f' %par)
Flow('cwav_',None, 'spike n1=%(nt)d d1=%(dt)f k1=%(kt)d | imagsrc frequency=%(frq)f | rtoc' %par)
#srcgen('cwav_',par)

# . . 3D Elastic source
Flow('souz','cwav_','math output=input*1')
Flow('soux','cwav_','math output=input*1')
Flow('souy','cwav_','math output=input*1')

Flow('wave-3d',['souz','soux','souy'],
     '''
     cat axis=2 space=n ${SOURCES[1:3]} |
     transp plane=23 | 
     transp plane=14
     ''')

# . . 3D Elastic source
Flow('souz0','wav_','math output=input*1')
Flow('soux0','wav_','math output=input*1')
Flow('souy0','wav_','math output=input*1')

Flow('wave-3d0',['souz0','soux0','souy0'],
     '''
     cat axis=2 space=n ${SOURCES[1:3]} |
     transp plane=23 | 
     transp plane=14
     ''')

# --------------------------------------------------------------------
# . . Coordinate Section
# . . Locate source position
xsou=par['ox']+(par['nx']-1)*par['dx']/2.
ysou=par['oy']+(par['ny']-1)*par['dy']/2.
zsou=par['oz']+(par['nz']-1)*par['dz']/2.
#zsou=par['oz']+40*par['dz']

# . . 3D Sources
fdmod.point3d('ss-3d',xsou,ysou,zsou,par) # . . 3D  Sources

# . . 3D receivers
fdmod.horizontal3d('rr-3d',0.02,par) # . . 3D 

# . . Create a 3D point location for plotting
par['zlook'] = 0.2
par['nzcut'] = par['nz']/2
center=fdmod.center3d(xsou,ysou,par['zlook'],par)

# --------------------------------------------------------------------
# Set up Cijs
m = 'ORT'
Flow(m+'c-11','refl',
       '''
       unif3 v00=4.5,9 n1=401 d1=10.265  | window f1=80 n1=161 | put d1=%f d2=%f d3=%f o1=0 | smooth rect1=4 rect2=4 rect3=4 repeat=1
       '''% (par['dz'],par['dx'],par['dy']))
Flow(m+'c-12','refl',
       '''
       unif3 v00=1.8,3.6 n1=401 d1=10.265  | window f1=80 n1=161  | put d1=%f d2=%f d3=%f o1=0 | smooth rect1=4 rect2=4 rect3=4 repeat=1
       '''% (par['dz'],par['dx'],par['dy']))
Flow(m+'c-13','refl',
       '''
       unif3 v00=1.125,2.25 n1=401 d1=10.265  | window f1=80 n1=161  | put d1=%f d2=%f d3=%f o1=0 | smooth rect1=4 rect2=4 rect3=4 repeat=1
       '''% (par['dz'],par['dx'],par['dy']))
Flow(m+'c-22','refl',
       '''
       unif3 v00=4.92,9.84 n1=401 d1=10.265   | window f1=80 n1=161  | put d1=%f d2=%f d3=%f o1=0 | smooth rect1=4 rect2=4 rect3=4 repeat=1
       '''% (par['dz'],par['dx'],par['dy']))
Flow(m+'c-23','refl',
       '''
       unif3 v00=1.2,2.4 n1=401 d1=10.265   | window f1=80 n1=161  | put d1=%f d2=%f d3=%f o1=0 | smooth rect1=4 rect2=4 rect3=4 repeat=1
       '''% (par['dz'],par['dx'],par['dy']))
Flow(m+'c-33','refl',
       '''
       unif3 v00=2.9687,5.9375 n1=401 d1=10.265  | window f1=80 n1=161  | put d1=%f d2=%f d3=%f o1=0 | smooth rect1=4 rect2=4 rect3=4 repeat=1
       '''% (par['dz'],par['dx'],par['dy']))
Flow(m+'c-44','refl',
       '''
       unif3 v00=1,2 n1=401 d1=10.265   | window f1=80 n1=161  | put d1=%f d2=%f d3=%f o1=0 | smooth rect1=4 rect2=4 rect3=4 repeat=1
       '''% (par['dz'],par['dx'],par['dy']))
Flow(m+'c-55','refl',
       '''
       unif3 v00=0.8,1.6 n1=401 d1=10.265  | window f1=80 n1=161  | put d1=%f d2=%f d3=%f o1=0 | smooth rect1=4 rect2=4 rect3=4 repeat=1
       '''% (par['dz'],par['dx'],par['dy']))
Flow(m+'c-66','refl',
       '''
       unif3 v00=1.091,2.182 n1=401s d1=10.265  | window f1=80 n1=161  | put d1=%f d2=%f d3=%f o1=0 | smooth rect1=4 rect2=4 rect3=4 repeat=1
       '''% (par['dz'],par['dx'],par['dy']))

Flow(m+'-the',m+'c-11','math output=45')
Flow(m+'-phi',m+'c-11','math output=30')

Result(m+'c-11'+'-f',m+'c-11','byte gainpanel=a bar=bar.rsf barreverse=n bias=5 allpos=y | grey3 color=j frame1=40 frame2=50 frame3=50 point1=0.5 point2=0.5 screenratio=1 scalebar=y title="Orthorhombic Model" label1=Z unit1=km label2=X unit2=km label3=Y unit3=km barlabel=c11 barunit="km\^2\_/s\^2\_" ')


Flow('rank-p left-p right-p',[m+'c-11',m+'c-12',m+'c-13',m+'c-22',m+'c-23',m+'c-33',m+'c-44',m+'c-55',m+'c-66'],
     '''
     ewelr3 mode=1 jump=2
     dt=%(dt)f eps=1e-4 npk=10 verb=%(verb)s nb=%(nb)d
     c12=${SOURCES[1]} c13=${SOURCES[2]}
     c22=${SOURCES[3]} c23=${SOURCES[4]} c33=${SOURCES[5]} 
     c44=${SOURCES[6]} c55=${SOURCES[7]} c66=${SOURCES[8]} 
     tilt=n
     tric=n
     left=${TARGETS[1]} right=${TARGETS[2]}
     '''%par)

Flow('rank-s left-s right-s',[m+'c-11',m+'c-12',m+'c-13',m+'c-22',m+'c-23',m+'c-33',m+'c-44',m+'c-55',m+'c-66'],
     '''
     ewelr3 mode=2 jump=2
     dt=%(dt)f eps=1e-4 npk=10 verb=%(verb)s nb=%(nb)d
     c12=${SOURCES[1]} c13=${SOURCES[2]}
     c22=${SOURCES[3]} c23=${SOURCES[4]} c33=${SOURCES[5]} 
     c44=${SOURCES[6]} c55=${SOURCES[7]} c66=${SOURCES[8]} 
     tilt=n
     tric=n
     left=${TARGETS[1]} right=${TARGETS[2]}
     '''%par)

Flow([m+'d-lr',m+'w-lr-p',m+'w-lr-s'], ['wave-3d',m+'c-11','rank-p','left-p','right-p','rank-s','left-s','right-s','ss-3d','rr-3d'],
     '''
     ewedc3d 
     back=n
     ccc=${SOURCES[1]}
     rkp=${SOURCES[2]}
     ltp=${SOURCES[3]}
     rtp=${SOURCES[4]}
     rks=${SOURCES[5]}
     lts=${SOURCES[6]}
     rts=${SOURCES[7]}
     sou=${SOURCES[8]}
     rec=${SOURCES[9]}
     wfp=${TARGETS[1]}
     wfs=${TARGETS[2]}
     jdata=%(jdata)d verb=%(verb)s free=n
     snap=%(snap)s jsnap=%(jsnap)d
     nb=%(nb)d cb=%(cb)g nbell=%(nbell)d
     dabc=%(dabc)s
     esou=n
     '''%par)


Result(m+'w-lr-p-z',[m+'w-lr-p',m+'c-11'],'window n4=1 f4=0 f5=49 n5=1 | real | math c=${SOURCES[1]} output="input*3e5+c*0" | byte bias=0 gainpanel=a | grey3 frame1=80 frame2=80 frame3=80 point1=0.5 point2=0.5 flat=y screenratio=1 title="P-wave (z)" label1=Z unit1=km label2=X unit2=km label3=Y unit3=km')
Result(m+'w-lr-p-x',[m+'w-lr-p',m+'c-11'],'window n4=1 f4=1 f5=49 n5=1 | real | math c=${SOURCES[1]} output="input*3e5+c*0" | byte bias=0 gainpanel=a | grey3 frame1=80 frame2=80 frame3=80 point1=0.5 point2=0.5 flat=y screenratio=1 title="P-wave (x)" label1=Z unit1=km label2=X unit2=km label3=Y unit3=km')
Result(m+'w-lr-p-y',[m+'w-lr-p',m+'c-11'],'window n4=1 f4=2 f5=49 n5=1 | real | math c=${SOURCES[1]} output="input*3e5+c*0" | byte bias=0 gainpanel=a | grey3 frame1=80 frame2=80 frame3=80 point1=0.5 point2=0.5 flat=y screenratio=1 title="P-wave (y)" label1=Z unit1=km label2=X unit2=km label3=Y unit3=km')
Result(m+'w-lr-s-z',[m+'w-lr-s',m+'c-11'],'window n4=1 f4=0 f5=49 n5=1 | real | math c=${SOURCES[1]} output="input*3e5+c*0" | byte bias=0 gainpanel=a | grey3 frame1=80 frame2=80 frame3=80 point1=0.5 point2=0.5 flat=y screenratio=1 title="Coupled S-waves (z)" label1=Z unit1=km label2=X unit2=km label3=Y unit3=km')
Result(m+'w-lr-s-x',[m+'w-lr-s',m+'c-11'],'window n4=1 f4=1 f5=49 n5=1 | real | math c=${SOURCES[1]} output="input*3e5+c*0" | byte bias=0 gainpanel=a | grey3 frame1=80 frame2=80 frame3=80 point1=0.5 point2=0.5 flat=y screenratio=1 title="Coupled S-waves (x)" label1=Z unit1=km label2=X unit2=km label3=Y unit3=km')
Result(m+'w-lr-s-y',[m+'w-lr-s',m+'c-11'],'window n4=1 f4=2 f5=49 n5=1 | real | math c=${SOURCES[1]} output="input*3e5+c*0" | byte bias=0 gainpanel=a | grey3 frame1=80 frame2=80 frame3=80 point1=0.5 point2=0.5 flat=y screenratio=1 title="Coupled S-waves (y)" label1=Z unit1=km label2=X unit2=km label3=Y unit3=km')

# generate source movies
nmov = 7
for i in range (nmov):
    movp = m+'w-lr-p-x-%d' %i
    movs = m+'w-lr-s-x-%d' %i
    Result(movp,[m+'w-lr-p',m+'c-11'],'window n4=1 f4=1 f5=%d n5=1 | real | math c=${SOURCES[1]} output="input*3e5+c" | byte bias=9 gainpanel=a clip=16.4246 | grey3 frame1=80 frame2=80 frame3=80 point1=0.5 point2=0.5 flat=y screenratio=1 title="P-wave (x)"' %((i+1)*7))
    Result(movs,[m+'w-lr-s',m+'c-11'],'window n4=1 f4=1 f5=%d n5=1 | real | math c=${SOURCES[1]} output="input*3e5+c" | byte bias=9 gainpanel=a clip=16.4246 | grey3 frame1=80 frame2=80 frame3=80 point1=0.5 point2=0.5 flat=y screenratio=1 title="Coupled S-waves (x)"' %((i+1)*7))

# display the angle
Flow('snap-p-x-cut',m+'w-lr-p','sfwindow n4=1 f4=1 n5=1 f5=49 | sfcut f1=50 | sfcut n2=40 | sfcut f2=80 | sfcut n3=40 | sfcut f3=80')
Result('snap-p-x-cut','sfreal | sfbyte gainpanel=a | sfgrey3 frame1=40 frame2=55 frame3=58 point1=0.5 point2=0.5 flat=y screenratio=1 title="A Portion of P-wavefield" label1=Z unit1=km label2=X unit2=km label3=Y unit3=km')
Flow('freq-p-x','snap-p-x-cut','fft3 axis=1 pad=1 | fft3 axis=2 pad=1 | fft3 axis=3 pad=1')
Result('freq-p-x','sfmath output="abs(input)" | sfreal | sfbyte mean=y gainpanel=a pclip=98 bar=scbar.rsf | sfgrey3 color=j frame1=81 frame2=81 frame3=81 point1=0.5 point2=0.5 flat=y screenratio=1 title="Wavenumber Spectrum" label1=kz label2=kx label3=ky unit1=1/km unit2=1/km unit3=1/km scalebar=y barlabel= barunit=')

Flow('freq-p-x2','snap-p-x-cut','real | rtoc | fft3 axis=1 pad=1 | fft3 axis=2 pad=1 | fft3 axis=3 pad=1')
Result('freq-p-x2','sfmath output="abs(input)" | sfreal | sfbyte mean=y gainpanel=a pclip=98 bar=scbar.rsf | sfgrey3 color=j frame1=81 frame2=81 frame3=81 point1=0.5 point2=0.5 flat=y screenratio=1 title="Wavenumber Spectrum" label1=kz label2=kx label3=ky unit1=1/km unit2=1/km unit3=1/km scalebar=y barlabel= barunit=')

# up-down separation
Flow('snap-p-x',m+'w-lr-p','window n4=1 f4=1 n5=1 f5=49')
Result('snap-p-x','real | byte gainpanel=a clip=5e-5 | grey3 frame1=80 frame2=80 frame3=80 point1=0.5 point2=0.5 flat=y screenratio=1 title="P Wavefield(x)" label1=Z unit1=km label2=X unit2=km label3=Y unit3=km')
#Result('snap-p-x',['snap-p-x',m+'c-11'],'real | math c=${SOURCES[1]} output="input*3e5+c" | byte bias=9 gainpanel=a clip=16.4246 | grey3 frame1=80 frame2=80 frame3=80 point1=0.5 point2=0.5 flat=y screenratio=1 title="P Wavefield(x)" label1=Z unit1=km label2=X unit2=km label3=Y unit3=km')

Flow('snap-p-x-up','snap-p-x','fft3 axis=1 pad=1 | fft3 axis=2 pad=1 | fft3 axis=3 pad=1 | cut n1=81 | fft3 axis=3 pad=1 inv=y | fft3 axis=2 pad=1 inv=y | fft3 axis=1 pad=1 inv=y')
Result('snap-p-x-up','real | byte gainpanel=a clip=5e-5 | grey3 frame1=80 frame2=80 frame3=80 point1=0.5 point2=0.5 flat=y screenratio=1 title="Up-going P-wave(x)" label1=Z unit1=km label2=X unit2=km label3=Y unit3=km')
#Result('snap-p-x-up',['snap-p-x-up',m+'c-11'],'real | math c=${SOURCES[1]} output="input*3e5+c" | byte bias=9 gainpanel=a clip=16.4246 | grey3 frame1=80 frame2=80 frame3=80 point1=0.5 point2=0.5 flat=y screenratio=1 title="Up-going P-wave(x)" label1=Z unit1=km label2=X unit2=km label3=Y unit3=km')

Flow('snap-p-x-dn','snap-p-x','fft3 axis=1 pad=1 | fft3 axis=2 pad=1 | fft3 axis=3 pad=1 | cut f1=81 | fft3 axis=3 pad=1 inv=y | fft3 axis=2 pad=1 inv=y | fft3 axis=1 pad=1 inv=y')
Result('snap-p-x-dn','real | byte gainpanel=a clip=5e-5 | grey3 frame1=80 frame2=80 frame3=80 point1=0.5 point2=0.5 flat=y screenratio=1 title="Down-going P-wave(x)" label1=Z unit1=km label2=X unit2=km label3=Y unit3=km')
#Result('snap-p-x-dn',['snap-p-x-dn',m+'c-11'],'real | math c=${SOURCES[1]} output="input*3e5+c" | byte bias=9 gainpanel=a clip=16.4246 | grey3 frame1=80 frame2=80 frame3=80 point1=0.5 point2=0.5 flat=y screenratio=1 title="Down-going P-wave(x)" label1=Z unit1=km label2=X unit2=km label3=Y unit3=km')

Flow('snap-p-x-up2','snap-p-x','real | rtoc | fft3 axis=1 pad=1 | fft3 axis=2 pad=1 | fft3 axis=3 pad=1 | cut n1=81 | fft3 axis=3 pad=1 inv=y | fft3 axis=2 pad=1 inv=y | fft3 axis=1 pad=1 inv=y')
Result('snap-p-x-up2','real | byte gainpanel=a clip=5e-5 | grey3 frame1=80 frame2=80 frame3=80 point1=0.5 point2=0.5 flat=y screenratio=1 title="Up-going P-wave(x)" label1=Z unit1=km label2=X unit2=km label3=Y unit3=km')
#Result('snap-p-x-up2',['snap-p-x-up2',m+'c-11'],'real | math c=${SOURCES[1]} output="input*3e5+c" | byte bias=9 gainpanel=a clip=16.4246 | grey3 frame1=80 frame2=80 frame3=80 point1=0.5 point2=0.5 flat=y screenratio=1 title="Up-going P-wave(x)" label1=Z unit1=km label2=X unit2=km label3=Y unit3=km')

Flow('snap-s-x',m+'w-lr-s','window n4=1 f4=1 n5=1 f5=49')
Result('snap-s-x','real | byte gainpanel=a clip=5e-5 | grey3 frame1=80 frame2=80 frame3=80 point1=0.5 point2=0.5 flat=y screenratio=1 title="S Wavefield(x)" label1=Z unit1=km label2=X unit2=km label3=Y unit3=km')
#Result('snap-s-x',['snap-s-x',m+'c-11'],'real | math c=${SOURCES[1]} output="input*3e5+c" | byte bias=9 gainpanel=a clip=16.4246 | grey3 frame1=80 frame2=80 frame3=80 point1=0.5 point2=0.5 flat=y screenratio=1 title="S Wavefield(x)" label1=Z unit1=km label2=X unit2=km label3=Y unit3=km')

Flow('snap-s-x-up','snap-s-x','fft3 axis=1 pad=1 | fft3 axis=2 pad=1 | fft3 axis=3 pad=1 | cut n1=81 | fft3 axis=3 pad=1 inv=y | fft3 axis=2 pad=1 inv=y | fft3 axis=1 pad=1 inv=y')
Result('snap-s-x-up','real | byte gainpanel=a clip=5e-5 | grey3 frame1=80 frame2=80 frame3=80 point1=0.5 point2=0.5 flat=y screenratio=1 title="Up-going S-wave(x)" label1=Z unit1=km label2=X unit2=km label3=Y unit3=km')
#Result('snap-s-x-up',['snap-s-x-up',m+'c-11'],'real | math c=${SOURCES[1]} output="input*3e5+c" | byte bias=9 gainpanel=a clip=16.4246 | grey3 frame1=80 frame2=80 frame3=80 point1=0.5 point2=0.5 flat=y screenratio=1 title="Up-going S-wave(x)" label1=Z unit1=km label2=X unit2=km label3=Y unit3=km')

Flow('snap-s-x-dn','snap-s-x','fft3 axis=1 pad=1 | fft3 axis=2 pad=1 | fft3 axis=3 pad=1 | cut f1=81 | fft3 axis=3 pad=1 inv=y | fft3 axis=2 pad=1 inv=y | fft3 axis=1 pad=1 inv=y')
Result('snap-s-x-dn','real | byte gainpanel=a clip=5e-5 | grey3 frame1=80 frame2=80 frame3=80 point1=0.5 point2=0.5 flat=y screenratio=1 title="Down-going S-wave(x)" label1=Z unit1=km label2=X unit2=km label3=Y unit3=km')
#Result('snap-s-x-dn',['snap-s-x-dn',m+'c-11'],'real | math c=${SOURCES[1]} output="input*3e5+c" | byte bias=9 gainpanel=a clip=16.4246 | grey3 frame1=80 frame2=80 frame3=80 point1=0.5 point2=0.5 flat=y screenratio=1 title="Down-going S-wave(x)" label1=Z unit1=km label2=X unit2=km label3=Y unit3=km')

Flow('snap-x','snap-p-x snap-s-x','add ${SOURCES[1]}')
Result('snap-x',['snap-x',m+'c-11'],'real | math c=${SOURCES[1]} output="input*3e5+c" | byte bias=9 gainpanel=a clip=16.4246 | grey3 frame1=80 frame2=80 frame3=80 point1=0.5 point2=0.5 flat=y screenratio=1 title="Elastic Wavefield(x)" label1=Z unit1=km label2=X unit2=km label3=Y unit3=km')

End()