from rsf.proj import *

par = dict(
    nx=401,
    nz=401,
    dx=0.005,
    dz=0.005,
    x0=0.0,
    z0=0.0,

    ns=601,
    dt=0.0005,
    da=0.01,
    time=0.3,
    f0=30.0,
    t0=0.04,

    vp0=3000,
    vs0=1500,
    ep=0.1,
    de=0.05,
    the=0.0,

    isep1=0,
    isep2=1,
    itaper1=0,
    itaper2=1,
    ihomo=1,
    tapertype="T",
    nstep=2
    )

# =================================================================================
Flow('vp0', None,
     '''
       math n1=%d n2=%d d1=%g d2=%g o1=%g o2=%g
            label1=x1 unit1=km label2=x2 unit2=km 
            output=%g
      ''' % (par['nz'],par['nx'],par['dz'],par['dx'],par['z0'],par['x0'],par['vp0'])
     )

Flow('vs0','vp0',
        '''
         math output=%g
        ''' % (par['vs0'])
 )

Flow('epsi','vp0',
        '''
         math output=%g
        ''' % (par['ep'])
 )

Flow('del','vp0',
        '''
         math output=%g
        ''' % (par['de'])
 )

# =================================================================================
# Forward modeling using original elastic wave eq.
# =================================================================================

name1='''
Elasticx Elasticz
'''

Flow(['Elasticx',      'Elasticz'],
         'vp0  vs0  epsi del',
         '''
         vti2desep
         vp0=${SOURCES[0]}
         vs0=${SOURCES[1]}
         epsi=${SOURCES[2]}
         del=${SOURCES[3]}
         Elasticz=${TARGETS[1]}
         ns=%d
         dt=%g
         isep=%d
         ihomo=%d
         tapertype=%s
         nstep=%d
         ''' % (par['ns'],par['dt'],par['isep1'],par['ihomo'],par['tapertype'],par['nstep'])
    )

for qq in Split(name1):
        Result(qq,
        '''
        grey polarity=n scalebar=n screenratio=1 wanttitle=n
        axisfat=5 axiscol=7 labelsz=10
        ''')

# =================================================================================
# Having Taper for operators in spatial domain
# =================================================================================
name2='''
PseudoPurePx PseudoPurePz
'''

name3='''
PseudoPureP PseudoPureSepP
'''

Flow(['PseudoPurePx',  'PseudoPurePz',    'PseudoPureP',
         'adx',           'adz',
         'adxx',          'adzz',
         'apx',           'apz',
         'apxx',          'apzz',
         'apvx',          'apvz',
         'apvxx',         'apvzz',
         'PseudoPureSepP'],
         'vp0  vs0  epsi del',
         '''
         vti2dpseudop1
         vp0=${SOURCES[0]}
         vs0=${SOURCES[1]}
         epsi=${SOURCES[2]}
         del=${SOURCES[3]}
         PseudoPurePz=${TARGETS[1]}
         PseudoPureP=${TARGETS[2]}
         adx=${TARGETS[3]}
         adz=${TARGETS[4]}
         adxx=${TARGETS[5]}
         adzz=${TARGETS[6]}
         apx=${TARGETS[7]}
         apz=${TARGETS[8]}
         apxx=${TARGETS[9]}
         apzz=${TARGETS[10]}
         apvx=${TARGETS[11]}
         apvz=${TARGETS[12]}
         apvxx=${TARGETS[13]}
         apvzz=${TARGETS[14]}
         PseudoPureSepP=${TARGETS[15]}
         ns=%d
         dt=%g
         isep=%d
         ihomo=%d
         itaper=%d
         tapertype=%s
         ''' % (par['ns'],par['dt'],par['isep2'],par['ihomo'],par['itaper2'],par['tapertype'])
    )

for qq in Split(name2):
        Result(qq,
        '''
        grey polarity=n scalebar=n screenratio=1 wanttitle=n
        axisfat=5 axiscol=7 labelsz=10
        ''')

for qq in Split(name3):
        Result(qq,
        '''
        grey polarity=n scalebar=n screenratio=1 wanttitle=n
        axisfat=5 axiscol=7 labelsz=10
        ''')

# =================================================================================
# calculate geometric wavefront of qP and qSV waves 
# =================================================================================
Flow(['WF',  'WFp',   'WFs'], None,
        '''
         kine2dvti 
         WFp=${TARGETS[1]}
         WFs=${TARGETS[2]}
         nx=%d
         nz=%d
         dx=%g
         dz=%g
         da=%g
         time=%g
         vp0=%g
         vs0=%g
         eps=%g
         del=%g
         the=%g
         f0=%g
         t0=%g
         ''' % (par['nx'],par['nz'],par['dx'],par['dz'],par['da'],par['time'],
                par['vp0'],par['vs0'],par['ep'],par['de'],par['the'],
                par['f0'],par['t0'])
    )

Result('WF',
        '''
        grey polarity=y backcol=0 scalebar=n screenratio=1 wanttitle=n
        ''')

Result('WFp',
        '''
        grey polarity=y backcol=7 scalebar=n screenratio=1 wanttitle=n
        ''')

Result('WFs',
        '''
        grey polarity=y backcol=0 scalebar=n screenratio=1 wanttitle=n
        ''')

End()