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

# . . Set modelling parameters 
par = {
# . . Space parameters
'nx':100, 'ox':0, 'dx':0.01,  'lx':'x', 'ux':'km',
'ny':100, 'oy':0, 'dy':0.01,  'ly':'y', 'uy':'km',
'nz':100, 'oz':0, 'dz':0.01,  'lz':'z', 'uz':'km',

# . . Time parameters
'nt':1000,'ot':0, 'dt':0.0005,  'lt':'t', 'ut':'s',
'kt':60,'frq':40,

# . . Modelling parameters
'snap':'y','jsnap':100, 'nb':26,'dabc':'y',
'nbell':5,'jdata':16,'verb':'y', 'ssou':'y','ngpu':1,

# . . Output parameters
'height':10, 'labelattr':' ','nframe':10,'iframe':7,
}

# . . Initialize parameters in fdm module   
fdmod.param(par)

# -----------------------------
#
# . . Model Section
#
# . . Background density model
fdmod.gauss3d('DBG',0.5,1,0.5,1,0.5,1,par)
Flow('DenBG','DBG','math output="0*input+2" ')

# . . Point perturbations to generate diffractors
dindex=range(100,127,1)
ii=100
for yy in [0.4, 0.6, 0.8]:
    for xx in [0.4, 0.6, 0.8]:
        for zz in [0.4, 0.6, 0.8]:
            ztag="-%04g"%zz
            xtag="-%04g"%xx
            ytag="-%04g"%yy
            itag="-%03d"%ii
            fdmod.gauss3d('D'+itag+'.rsf',xx,yy,zz,0.01,0.01,0.01,par)
            ii=ii+1
Flow('DenPOINT',map(lambda x: 'D-'+'%03d' % x,dindex),'add ${SOURCES[1:27]}' %par )

# . . Total Density Field
Flow('Density','DenBG DenPOINT','math points=${SOURCES[1]} output="input+points" | put label1=Depth label2=X-axis label3=Y-axis unit1=km unit2=km unit3=km')

#. . Create Model
Flow('Model','Density','pad axis=4 n4=9')

# . . Isotropic Field
Flow('ISO','Model Density',
        '''
        modelcreate 
        den=${SOURCES[1]} 
        aniso=n ein=0 din=0 gin=0 verb=y a=2 b=2 d=60 
        ''' %par)
Flow('ISOMODEL','ISO','smooth rect1=4 rect2=4 rect3=4 repeat=2')

# . . Aniostropic model
Flow('ANISO','Model Density',
        '''
        modelcreate 
        den=${SOURCES[1]} 
        allaniso=y aniso=y ein=0.2 din=0.2 gin=0.2 verb=y a=2 b=2 d=60 
        '''%par)
Flow('ANISOMODEL','ANISO','smooth rect1=4 rect2=4 rect3=4 repeat=2')

# --------------------------------------------------------------------
# . . Source Section
Flow('wav_',None,'spike nsp=1 mag=1 n1=%(nt)d d1=%(dt)g o1=%(ot)g k1=%(kt)d | ricker1 frequency=%(frq)g' %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 | scale axis=123')

# --------------------------------------------------------------------
# . . 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 Receivers
par['nrec']=par['nx']*par['ny']/16
fdmod.horizontal3d('rr-3d',0,par) # . . 3D 
Flow('rr','rr-3d','window j3=4 j2=4 | put n2=%(nrec)d n3=1' %par)

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

# ------------------------------------------------------------
Flow('t1','ISOMODEL','window n4=1')

for m in ['ISO','ANISO']:

    stiffness.cplot3d(m+'MODEL',1,1,1,par)

    # . . Perform 3D modelling
    Flow([m+'data',m+'wfld'], ['wave-3d',m+'MODEL','Density','ss-3d','rr-3d'],
        '''
        ewefd3d_gpu_p2p
        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 ngpu=%(ngpu)d
        dabc=%(dabc)s nb=%(nb)d nbell=%(nbell)d
        ''' %par)

    Flow(m+'datacut',m+'wfld',
                'window n1=1 n2=%(nx)d n3=%(ny)d min1=%(oz)g min2=%(ox)g min3=%(oy)g n4=1 | transp plane=23 | transp plane=12  ' % par)

    # . . Make 3D wavefield modelling movie
    Result(m+'wfld',[m+'wfld','t1'],'window n5=1 f5=%(iframe)d n4=1 min1=0 min2=0 min3=0 n4=1 n1=%(nz)d n2=%(nx)d n3=%(ny)d | scale axis=123 | math vel=${SOURCES[1]} output="input-(vel-5.13439)/100" | byte gainpanel=a pclip=99|' %par
           + fdmod.cgrey3d('pclip=99.9 screenht=11 label1=Depth label2=Inline label3=Crossline unit1=km unit2=km unit3=km '+center,par))

End()