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

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

    # Time Parameters
    'nt':408,'ot':0, 'dt':0.0005,  'lt':'t', 'ut':'s',
    'kt':100,'frq':35,

    # Modelling parameters
    'snap':'y','jsnap':100,'nb':44, 'verb':'y',
    'nbell':5, 'jdata':8,'ssou':'y',

    # Output
    'height':5,
    }

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

# . . Thomsen parameters
par['vp']=2.0
par['vs']=1.1547
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

# --------------------------------------------------------------------
# . . Source Section
#
# . . Wavelet
fdmod.wavelet('wav_',par['frq'],par)

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

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

# --------------------------------------------------------------------
# . . 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']

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

# . . 3D receivers
fdmod.horizontal3d('rr-3d',0,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)

# --------------------------------------------------------------------
# . .  2D model section

# . . Create zero array size of 2D model
Flow('zero-2d',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 |
     put label1=%(lz)s label2=%(lx)s unit1=%(uz)s unit2=%(ux)s
     ''' % 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)

# . . Generate 3D models of anisotropy
Flow('vp-3d','zero-3d','math output="%(vp)g"' %par)
Flow('vs-3d','zero-3d','math output="%(vs)g"' %par)
Flow('ro-3d','zero-3d','math output="%(ro)g"' %par)
Flow('eps1-3d','zero-3d','math output="%(eps1)g"' %par)
Flow('eps2-3d','zero-3d','math output="%(eps2)g"' %par)
Flow('del1-3d','zero-3d','math output="%(del1)g"' %par)
Flow('del2-3d','zero-3d','math output="%(del2)g"' %par)
Flow('del3-3d','zero-3d','math output="%(del3)g"' %par)
Flow('gam1-3d','zero-3d','math output="%(gam1)g"' %par)
Flow('gam2-3d','zero-3d','math output="%(gam2)g"' %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)

Flow('HTI11','VTIc-3d','window n4=1 f4=0')
Flow('HTI22','VTIc-3d','window n4=1 f4=1')
Flow('HTI33','VTIc-3d','window n4=1 f4=2')
Flow('HTI44','VTIc-3d','window n4=1 f4=3')
Flow('HTI55','VTIc-3d','window n4=1 f4=4')
Flow('HTI66','VTIc-3d','window n4=1 f4=5')
Flow('HTI12','VTIc-3d','window n4=1 f4=6')
Flow('HTI13','VTIc-3d','window n4=1 f4=7')
Flow('HTI23','VTIc-3d','window n4=1 f4=8')
Flow('HTIc-3d','HTI33 HTI22 HTI11 HTI66 HTI55 HTI44 HTI23 HTI13 HTI12','cat axis=4 ${SOURCES[1:9]} ')

sindex=range(0,9,1)
Flow('null',None,'spike n1=16 n2=16 nsp=1 mag=0')

# . . Make different stiffness tensors
for m in ['ISO','VTI','ORT','HTI']:
        for isou in sindex:
                stag="%01d"%isou
                Flow(m+'p1'+stag,m+'c-3d','window n1=1 n2=1 n3=1 n4=1 f4=%d '%isou)
                Flow(m+'p2'+stag,[m+'p1'+stag,  m+'p1'+stag,m+'p1'+stag,m+'p1'+stag,m+'p1'+stag,m+'p1'+stag,m+'p1'+stag,m+'p1'+stag,m+'p1'+stag,m+'p1'+stag,m+'p1'+stag,m+'p1'+stag,m+'p1'+stag,m+'p1'+stag,m+'p1'+stag,m+'p1'+stag], 'cat ${SOURCES[1:16]} axis=1')
                Flow(m+'c'+stag,[m+'p2'+stag,m+'p2'+stag,m+'p2'+stag,m+'p2'+stag,m+'p2'+stag,m+'p2'+stag,m+'p2'+stag,m+'p2'+stag,m+'p2'+stag,m+'p2'+stag,m+'p2'+stag,m+'p2'+stag,m+'p2'+stag,m+'p2'+stag,m+'p2'+stag,m+'p2'+stag],'cat ${SOURCES[1:16]} axis=2')
        Flow(m+'x1',[m+'c0',m+'c6',m+'c7','null','null','null'],'cat axis=2 ${SOURCES[1:6]} ')
        Flow(m+'x2',[m+'c6',m+'c1',m+'c8','null','null','null'],'cat axis=2 ${SOURCES[1:6]} ')
        Flow(m+'x3',[m+'c7',m+'c8',m+'c2','null','null','null'],'cat axis=2 ${SOURCES[1:6]} ')
        Flow(m+'x4',['null','null','null',m+'c3','null','null'],'cat axis=2 ${SOURCES[1:6]} ')
        Flow(m+'x5',['null','null','null','null',m+'c4','null'],'cat axis=2 ${SOURCES[1:6]} ')
        Flow(m+'x6',['null','null','null','null','null',m+'c5'],'cat axis=2 ${SOURCES[1:6]} ')
        Flow(m+'c',[m+'x1',m+'x2',m+'x3',m+'x4',m+'x5',m+'x6'],'cat axis=1 ${SOURCES[1:6]} ')
        Result(m+'c','grey pclip=100 title="" wantaxis=n screenratio=1 allpos=y color=j')

# ------------------------------------------------------------
# . . Elastic Modeling Section
# . . GPU SECTION
for m in ['ISO','VTI','ORT','HTI']:

    # . . Perform 3D modelling
    Flow([m+'d-3d-GPU',m+'w-3d-GPU'], ['wave-3d',m+'c-3d','ro-3d','ss-3d','rr-3d'],
        '''
        ewefd3d_gpu_p2p
        ccc=${SOURCES[1]} 
        den=${SOURCES[2]}
        sou=${SOURCES[3]} 
        rec=${SOURCES[4]}
        wfl=${TARGETS[1]}
        ngpu=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
        ''' %par)

    # . . Make 3D wavefield modelling movie
    Result(m+'w-3d-GPU','window n5=1 f5=%(iframe)d min1=0 min2=0 min3=0 n4=1 n1=%(nz)d n2=%(nx)d n3=%(ny)d | byte gainpanel=a pclip=99.9|' %par + fdmod.cgrey3d('pclip=99.9 screenht=11 label1=Depth label2=Inline label3=Crossline unit1=km unit2=km unit3=km '+center,par))

    for i in range(3):
        tag = "%d"%i
        Result(m+'w-3d-GPU'+tag,m+'w-3d-GPU',
                'window n5=1 f5=4 n4=1 f4=%d min1=0 min2=0 min3=0 n1=%d n2=%d n3=%d | byte pclip=99.5 gainpanel=a |'%(i,par['nz'],par['nx'],par['ny']) + fdmod.cgrey3d('pclip=99.9 screenht=11 '+center+' movie=1 frame1=1',par))

# . . CPU SECTION
for m in ['ISO','VTI','ORT','HTI']:

    # . . Perform 3D modelling
    Flow([m+'d-3d-CPU',m+'w-3d-CPU'], ['wave-3d',m+'c-3d','ro-3d','ss-3d','rr-3d'],
    '''
    ewefd3d_omp
    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
    ''' %par)


    # . . Make 3D wavefield modelling movie
    Result(m+'w-3d-CPU','window n5=1 f5=%(iframe)d min1=0 min2=0 min3=0 n4=1 n1=100 n2=100 n3=100|byte gainpanel=a pclip=99.9|' %par
           + fdmod.cgrey3d('pclip=99.9 screenht=11 '+center,par))

    for i in range(3):
        tag = "%d"%i
        Result(m+'w-3d-CPU'+tag,m+'w-3d-CPU',' window n5=1 f5=4 n4=1 f4=%d min1=0 min2=0 min3=0 n1=%d n2=%d n3=%d | byte pclip=99.5 gainpanel=a |'%(i,par['nz'],par['nx'],par['ny'])+ fdmod.cgrey3d('pclip=99.9 screenht=11 '+center+' movie=1 frame1=1',par))

End()