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,

    vp0_1=2500,
    vs0_1=1200,
    eps_1=0.2,
    del_1=-0.2,
    the_1=0.0,
    vp0_2=3600,
    vs0_2=1800,
    eps_2=0.2,
    del_2=0.1,
    the_2=30.0,

    seed=2012,
    eps=1.e-10,
    npk=30,

        n2w=1,
    o2w=0.75,

    )

# =================================================================================
Flow(['vp0','vs0','epsi','del','the'], None,
     '''
         twolayer2dti
         vs0=${TARGETS[1]}
         epsi=${TARGETS[2]}
         del=${TARGETS[3]}
         the=${TARGETS[4]}
         nx=%d
         nz=%d
         dx=%g
         dz=%g
         vp0_1=%g
         vs0_1=%g
         eps_1=%g
         del_1=%g
         the_1=%g
         vp0_2=%g
         vs0_2=%g
         eps_2=%g
         del_2=%g
         the_2=%g
         ''' % (par['nx'],par['nz'],par['dx'],par['dz'],
                par['vp0_1'],par['vs0_1'],par['eps_1'],par['del_1'],par['the_1'],
                par['vp0_2'],par['vs0_2'],par['eps_2'],par['del_2'],par['the_2'])
    )

# =================================================================================
# calculate separated wave-mode:         yes
# =================================================================================
name0='''
ElasticxFD ElasticzFD
'''

Flow(['ElasticxFD',  'ElasticzFD'],
         'vp0  vs0  epsi del the',
         '''
         tti2de
         vp0=${SOURCES[0]}
         vs0=${SOURCES[1]}
         epsi=${SOURCES[2]} 
         del=${SOURCES[3]}
         the=${SOURCES[4]}
         Elasticz=${TARGETS[1]}
         ns=%d 
         dt=%g
         ''' % (par['ns'],par['dt'])
    )

Flow('interface','ElasticxFD',
     '''
     math output="(%d-1)*%g*7/12"
     ''' % (par['nz'],par['dz'])
    )

Plot('interface',
     '''
     graph yreverse=y color= transp=n pad=n plotfat=1 scalebar=n
     min1=0 max1=2 
     min2=0 max2=2
     wanttitle=n wantaxis=n screenratio=1
     '''
     )

for qq in Split(name0):
      Plot(qq,
        '''
        grey color=j polarity=n scalebar=n screenratio=1 wanttitle=n pclip=99.9
        axisfat=5 axiscol=7 labelsz=10
        ''')

for qq in Split(name0):
    Result(qq+'Interf',[qq,'interface'],'Overlay')

Flow('ElasticzFDwave','ElasticzFD',
        '''
        window n2=%d min2=%g
        ''' % (par['n2w'],par['o2w']))

Result('ElasticzFDwave',
        '''
        graph label1=X title='10th-order finite difference' label2=Amplitude screenratio=0.5
        ''')

# =================================================================================
# extrapolation using Low-rank spectral approach
# =================================================================================
name='''
ElasticxPSLR ElasticzPSLR
'''

Flow(['ElasticxPSLR',  'ElasticzPSLR'],
      'vp0 vs0 epsi del the',
         '''
         tti2delr
         vp0=${SOURCES[0]}
         vs0=${SOURCES[1]}
         epsi=${SOURCES[2]} 
         del=${SOURCES[3]}
         the=${SOURCES[4]}
         Elasticz=${TARGETS[1]}
         ns=%d 
         dt=%g 
         eps=%g
         seed=%d
         npk=%d
         ''' % (par['ns'],par['dt'],par['eps'],par['seed'],par['npk'])
    )

for qq in Split(name):
        Plot(qq,
        '''
        grey color=j polarity=n scalebar=n screenratio=1 wanttitle=n pclip=99.9
        axisfat=5 axiscol=7 labelsz=10
        ''')

for qq in Split(name):
    Result(qq+'Interf',[qq,'interface'],'Overlay')

Flow('ElasticzPSLRwave','ElasticzPSLR',
        '''
        window n2=%d min2=%g
        ''' % (par['n2w'],par['o2w']))

Result('ElasticzPSLRwave',
        '''
        graph label1=X title='Low-rank pseudo-spectral' label2=Amplitude screenratio=0.5
        ''')

# =================================================================================
# extrapolation using k-space Low-rank spectral approach
# =================================================================================
name2='''
ElasticxKS ElasticzKS
'''

Flow(['ElasticxKS',  'ElasticzKS'],
      'vp0 vs0 epsi del the',
         '''
         tti2dekspacelr
         vp0=${SOURCES[0]}
         vs0=${SOURCES[1]}
         epsi=${SOURCES[2]} 
         del=${SOURCES[3]}
         the=${SOURCES[4]}
         Elasticz=${TARGETS[1]}
         ns=%d 
         dt=%g 
         eps=%g
         seed=%d
         npk=%d
         ''' % (par['ns'],par['dt'],par['eps'],par['seed'],par['npk'])
    )

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

for qq in Split(name2):
    Result(qq+'Interf',[qq,'interface'],'Overlay')

Flow('ElasticzKSwave','ElasticzKS',
        '''
        window n2=%d min2=%g
        ''' % (par['n2w'],par['o2w']))

Result('ElasticzKSwave',
        '''
        graph label1=X title='Low-rank pseudo-spectral (k-space)' label2=Amplitude screenratio=0.5
        ''')

# =================================================================================
# extrapolation using k-space Low-rank spectral approach with adjusted source term
# =================================================================================
name3='''
ElasticxKSS ElasticzKSS
'''

#         tti2dekspacelrsource_double_zp1
Flow(['ElasticxKSS',  'ElasticzKSS'],
      'vp0 vs0 epsi del the',
         '''
         tti2dekspacelrsource
         vp0=${SOURCES[0]}
         vs0=${SOURCES[1]}
         epsi=${SOURCES[2]} 
         del=${SOURCES[3]}
         the=${SOURCES[4]}
         Elasticz=${TARGETS[1]}
         ns=%d 
         dt=%g 
         eps=%g
         seed=%d
         npk=%d
         ''' % (par['ns'],par['dt'],par['eps'],par['seed'],par['npk'])
    )

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

for qq in Split(name3):
    Result(qq+'Interf',[qq,'interface'],'Overlay')

Flow('ElasticzKSSwave','ElasticzKSS',
        '''
        window n2=%d min2=%g
        ''' % (par['n2w'],par['o2w']))

Result('ElasticzKSSwave',
        '''
        graph label1=X title='Low-rank pseudo-spectral (k-space/source)' label2=Amplitude screenratio=0.5
        ''')

End()