# . . Set up generic project files
from rsf.proj import *
from rsf.recipes import fdmod,encode,wemig,stiffness
import os

def srcgen(out,par):
    dtt=0.0001
    fac=par['dt']/dtt
    ntt=par['nt']*fac
    ktt=par['kt']*fac
    dt=str(int(par['dt']*10000))

    #i/(2*phi)=i/(2|omega|)=i/2 * (hilb) [(int)source] 
    Flow('src-'+dt,None,
         '''
         spike n1=%d d1=%f k1=%d |
         ricker1 frequency=%f
         '''%(ntt,dtt,ktt,par['frq']))
    Flow('realsrc-'+dt,'src-'+dt,'math "output=0"')
    Flow('imagsrc-'+dt,'src-'+dt,'envelope hilb=y order=500 | halfint | halfint | math output="input/2" ')

    Flow('csrc-'+dt,['realsrc-'+dt,'imagsrc-'+dt],'cmplx ${SOURCES[1]}')
    Flow(out,'csrc-'+dt,'window j1=%d'% fac)

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

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

    # Modelling parameters
    'snap':'y','jsnap':8,'nb':0, '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'

# . . Thomsen parameters
par['vp']=2
par['vs']=1.5
par['ro']=2000
par['eps1']=+0.2
par['eps2']=+0.3
par['del1']=-0.1
par['del2']=-0.05
par['del3']=-0.075
par['gam1']=+0.2
par['gam2']=+0.5

# --------------------------------------------------------------------
# . . 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']+40*par['dz']
#zsou=par['oz']+(par['nz']-1)*par['dz']/2.

# . . 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)

# --------------------------------------------------------------------
# . . Create zero array size of 3D model
Flow('zero-3d',None,
     '''
     spike nsp=1 mag=0.0
     n1=%(nz)d o1=%(oz)g d1=%(dz)g
     n2=%(nx)d o2=%(ox)g d2=%(dx)g
     n3=%(ny)d o3=%(oy)g d3=%(dy)g |
     put label1=%(lz)s label2=%(lx)s label3=%(ly)s unit1=%(uz)s unit2=%(ux)s unit3=%(uy)s
     ''' % par)

Flow('box-3d',None,
     '''
     spike nsp=1 mag=0.8
     n1=%(nz)d o1=%(oz)g d1=%(dz)g
     n2=%(nx)d o2=%(ox)g d2=%(dx)g
     n3=%(ny)d o3=%(oy)g d3=%(dy)g 
     k1=%(kz)d l1=%(lz)d k2=%(kx)d l2=%(lx)d k3=%(ky)d l3=%(ly)d |
     smooth rect1=4 rect2=4 rect3=4 repeat=1 |
     put label1=%(lz)s label2=%(lx)s label3=%(ly)s unit1=%(uz)s unit2=%(ux)s unit3=%(uy)s
     ''' % (par))

# . . Generate 3D models of anisotropy
#Flow('vp-3d','box-3d','math output="%(vp)g+x1*1.5"' %par)
#Flow('vs-3d','box-3d','math output="%(vs)g+x1*1.5"' %par)
#Flow('ro-3d','box-3d','math output="%(ro)g"' %par)
#Flow('eps1-3d','box-3d','math output="%(eps1)g"' %par)
#Flow('eps2-3d','box-3d','math output="%(eps2)g"' %par)
#Flow('del1-3d','box-3d','math output="%(del1)g"' %par)
#Flow('del2-3d','box-3d','math output="%(del2)g"' %par)
#Flow('del3-3d','box-3d','math output="%(del3)g"' %par)
#Flow('gam1-3d','box-3d','math output="%(gam1)g"' %par)
#Flow('gam2-3d','box-3d','math output="%(gam2)g"' %par)

#Flow('vp-3d','box-3d','math output="%(vp)g*(1.+input)"' %par)
#Flow('vs-3d','box-3d','math output="%(vs)g*(1.+input)"' %par)
#Flow('ro-3d','box-3d','math output="%(ro)g*(1.+input)"' %par)
#Flow('eps1-3d','box-3d','math output="%(eps1)g*(1.+input)"' %par)
#Flow('eps2-3d','box-3d','math output="%(eps2)g*(1.+input)"' %par)
#Flow('del1-3d','box-3d','math output="%(del1)g*(1.+input)"' %par)
#Flow('del2-3d','box-3d','math output="%(del2)g*(1.+input)"' %par)
#Flow('del3-3d','box-3d','math output="%(del3)g*(1.+input)"' %par)
#Flow('gam1-3d','box-3d','math output="%(gam1)g*(1.+input)"' %par)
#Flow('gam2-3d','box-3d','math output="%(gam2)g*(1.+input)"' %par)
#
## . . Make 3D stiffness matricies
#Flow('zerostiff','vp-3d','math output="input*0" ')
#stiffness.iso3d('ISOc-3d','vp-3d','vs-3d','ro-3d',par)
#stiffness.vti3d('VTIc-3d','vp-3d','vs-3d','ro-3d','eps1-3d','del1-3d','gam1-3d',par)
#stiffness.ort3d('ORTc-3d','vp-3d','vs-3d','ro-3d','eps1-3d','eps2-3d','del1-3d','del2-3d','del3-3d','gam1-3d','gam2-3d',par)

m = 'ORT'
#for ij in ['11','12','13','22','23','33','44','55','66']:
#    Cij = m+'c-3d-'+ij
#    cij = m+'c-' + ij
#    Flow(cij,[Cij,'ro-3d'],'math ro=${SOURCES[1]} output="input/ro"')

# Shoenberg and Hilbig
c11=9.0
c12=3.6
c13=2.25
c22=9.84
c23=2.4
c33=5.9375
c44=2.0
c55=1.6
c66=2.182

Flow(m+'c-11','box-3d','math output="%f*(1.+input)"' % c11)
Flow(m+'c-12','box-3d','math output="%f*(1.+input)"' % c12)
Flow(m+'c-13','box-3d','math output="%f*(1.+input)"' % c13)
Flow(m+'c-22','box-3d','math output="%f*(1.+input)"' % c22)
Flow(m+'c-23','box-3d','math output="%f*(1.+input)"' % c23)
Flow(m+'c-33','box-3d','math output="%f*(1.+input)"' % c33)
Flow(m+'c-44','box-3d','math output="%f*(1.+input)"' % c44)
Flow(m+'c-55','box-3d','math output="%f*(1.+input)"' % c55)
Flow(m+'c-66','box-3d','math output="%f*(1.+input)"' % c66)
Flow(m+'c-3d',[m+'c-11',m+'c-22',m+'c-33',m+'c-44',m+'c-55',m+'c-66',m+'c-12',m+'c-13',m+'c-23'],'cat axis=4 ${SOURCES[1:9]} | math output=input*1000')
Flow('ro-3d','box-3d','math output="1000"')

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

Result(m+'c-11','byte gainpanel=a bar=bar.rsf barreverse=y bias=6 allpos=y | grey3 frame1=40 frame2=50 frame3=50 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('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)

mint = 0.18
nt0 = par['nt']
dt0 = par['dt']
kt0 = par['kt']
jt0 = par['jsnap']
for i in [1,2,4,8]:
    dt = str(int(i))
    par['dt']=dt0*i;
    par['nt']=nt0/i;
    par['kt']=kt0/i;
    par['jsnap']=jt0/i;
    #if (i==8):
    #   par['nbell']=2;
    # Source wavelet
    wav = 'wav-'+dt
    wav0= 'wav0-'+dt
    cwav = 'cwav-'+dt
    cwav0= 'cwav0-'+dt
    Flow(wav,'wav','window j1=%d' %i)
    Flow(cwav,'cwav','window j1=%d' %i)
    #srcgen(cwav,par)
    Flow(wav0,wav,'math output=0')
    Flow(cwav0,cwav,'math output=0')

    sou= 'sou-'+dt
    csou= 'csou-'+dt
    Flow(csou,[cwav,cwav,cwav],
    '''
    cat axis=2 space=n ${SOURCES[1:3]} |
    transp plane=23 | 
    transp plane=14
    ''')

    Flow(sou,[wav,wav,wav],
    '''
    cat axis=2 space=n ${SOURCES[1:3]} |
    transp plane=23 | 
    transp plane=14
    ''')

    # ------------------------------------------------------------
    # Lowrank Decomposition

    Flow(['rank-'+dt,'left-'+dt,'right-'+dt],[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',m+'-the',m+'-phi'],
    '''
    ewelr3 dt=%(dt)f eps=1e-4 npk=15 verb=%(verb)s nb=0
    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
    theta=${SOURCES[9]} phi=${SOURCES[10]}
    left=${TARGETS[1]} right=${TARGETS[2]}
    mode=0 jump=2
    '''%par)

    # Lowrank Elastic Modeling
    Flow([m+'d-lr-'+dt,m+'w-lr-'+dt], [csou,m+'c-11','rank-'+dt,'left-'+dt,'right-'+dt,'ss-3d','rr-3d'],
    '''
    ewelr3d 
    back=n
    ccc=${SOURCES[1]}
    rnk=${SOURCES[2]}
    lft=${SOURCES[3]}
    rht=${SOURCES[4]}
    sou=${SOURCES[5]}
    rec=${SOURCES[6]}
    wfl=${TARGETS[1]}
    jdata=%(jdata)d verb=%(verb)s free=n
    snap=%(snap)s jsnap=%(jsnap)d
    nb=0 nbell=%(nbell)d
    esou=n
    oqz=%(oz)f oqx=%(ox)f oqy=%(oy)f
    nqz=%(nz)f nqx=%(nx)f nqy=%(ny)f
    '''%par)


    Result(m+'w-lr-'+dt,'window n4=1 min5=%f n5=1 | real | byte gainpanel=a | grey3 frame1=40 frame2=50 frame3=50 point1=0.5 point2=0.5 screenratio=1 title="Lowrank RITE" label1=Z unit1=km label2=X unit2=km label3=Y unit3=km' %(mint) )

    # ------------------------------------------------------------
    # Pseudospectral Elastic Modeling

    Flow([m+'d-ps-'+dt,m+'w-ps-'+dt], [sou,m+'c-3d','ro-3d','ss-3d','rr-3d'],
    '''
    eweks3d
    back=n
    kspace=y
    ccc=${SOURCES[1]}
    den=${SOURCES[2]}
    sou=${SOURCES[3]}
    rec=${SOURCES[4]}
    wfl=${TARGETS[1]}
    jdata=%(jdata)d verb=%(verb)s free=n ssou=%(ssou)s
    opot=n snap=%(snap)s jsnap=%(jsnap)d
    dabc=%(dabc)s nb=%(nb)d nbell=%(nbell)d
    oqz=%(oz)f oqx=%(ox)f oqy=%(oy)f
    nqz=%(nz)f nqx=%(nx)f nqy=%(ny)f
    ''' %par)


    Result(m+'w-ps-'+dt,'window n4=1 min5=%f n5=1 | byte gainpanel=a | grey3 frame1=40 frame2=50 frame3=50 point1=0.5 point2=0.5 screenratio=1 title=Pseudospectral wanttitle=y label1=Z unit1=km label2=X unit2=km label3=Y unit3=km' %(mint))

    # ------------------------------------------------------------
    # Finite-difference Elastic Modeling

    Flow([m+'d-fd-'+dt,m+'w-fd-'+dt], [sou,m+'c-3d','ro-3d','ss-3d','rr-3d'],
    '''
    eweks3d
    back=n
    kspace=n
    ccc=${SOURCES[1]}
    den=${SOURCES[2]}
    sou=${SOURCES[3]}
    rec=${SOURCES[4]}
    wfl=${TARGETS[1]}
    jdata=%(jdata)d verb=%(verb)s free=n ssou=%(ssou)s
    opot=n snap=%(snap)s jsnap=%(jsnap)d
    dabc=%(dabc)s nb=%(nb)d nbell=%(nbell)d
    oqz=%(oz)f oqx=%(ox)f oqy=%(oy)f
    nqz=%(nz)f nqx=%(nx)f nqy=%(ny)f
    ''' %par)


    Result(m+'w-fd-'+dt,'window n4=1 min5=%f n5=1 | byte gainpanel=a | grey3 frame1=40 frame2=50 frame3=50 point1=0.5 point2=0.5 screenratio=1 title=Finite-difference wanttitle=y label1=Z unit1=km label2=X unit2=km label3=Y unit3=km' %(mint))

End()