from rsf.proj import *

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

    ns=601,
    dt=0.001,

    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-6,
    npk=30,

    n2w=1,
    o2w=1.0,

    )

# =================================================================================
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'])
    )

# =================================================================================
# extrapolation using Low-rank spectral approach with SVD truncation for wavenumber
# =================================================================================
name1='''
ElasticxPSLR ElasticzPSLR
'''
name2='''
ElasticPxPSLR ElasticPzPSLR ElasticSxPSLR ElasticSzPSLR
'''

Flow(['ElasticxPSLR',  'ElasticzPSLR',  'ElasticPxPSLR',  'ElasticPzPSLR',  'ElasticSxPSLR',  'ElasticSzPSLR'],
      'vp0 vs0 epsi del the',
         '''
         tti2devectorlrsvd
         vp0=${SOURCES[0]}
         vs0=${SOURCES[1]}
         epsi=${SOURCES[2]} 
         del=${SOURCES[3]}
         the=${SOURCES[4]}
         Elasticz=${TARGETS[1]}
         ElasticPx=${TARGETS[2]}
         ElasticPz=${TARGETS[3]}
         ElasticSx=${TARGETS[4]}
         ElasticSz=${TARGETS[5]}
         ns=%d 
         dt=%g 
         eps=%g
         seed=%d
         npk=%d
         ''' % (par['ns'],par['dt'],par['eps'],par['seed'],par['npk'])
    )

Flow('interface','ElasticxPSLR',
     '''
     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=4 
     min2=0 max2=4
     wanttitle=n wantaxis=n screenratio=1
     '''
     )

for qq in Split(name1):
        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):
        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(name1):
    Result(qq+'Interf',[qq,'interface'],'Overlay')

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

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

Result('ElasticzPSLRwave',
        '''
        graph label1=Z title='Lowrank pseudo-spectral' label2=Amplitude screenratio=0.5
        ''')

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

Result('ElasticPzPSLRwave',
        '''
        graph label1=Z title='Lowrank pseudo-spectral' label2=Amplitude screenratio=0.5
        ''')

# =================================================================================
# vector decomp
# =================================================================================
Flow(['ElasticxPSLR1',  'ElasticzPSLR1'],
      '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'])
    )

name3='''
ElasticxPSLR1 ElasticzPSLR1 ElasticPxPSLR1 ElasticPzPSLR1 ElasticSxPSLR1 ElasticSzPSLR1
'''

Flow(['ElasticPxPSLR1', 'ElasticPzPSLR1', 'ElasticSxPSLR1', 'ElasticSzPSLR1'],
      'vp0 vs0 epsi del the ElasticxPSLR1 ElasticzPSLR1',
         '''
         tti2delrdec
         vp0=${SOURCES[0]}
         vs0=${SOURCES[1]}
         epsi=${SOURCES[2]} 
         del=${SOURCES[3]}
         the=${SOURCES[4]}
         Elasticx=${SOURCES[5]}
         Elasticz=${SOURCES[6]}
         ElasticPx=${TARGETS[0]}
         ElasticPz=${TARGETS[1]}
         ElasticSx=${TARGETS[2]}
         ElasticSz=${TARGETS[3]}
         eps=%g
         seed=%d
         npk=%d
         ''' % (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('ElasticzPSLR1wave','ElasticzPSLR1',
        '''
        window n2=%d min2=%g
        ''' % (par['n2w'],par['o2w']))

Result('ElasticzPSLR1wave',
        '''
        graph label1=Z title='Lowrank pseudo-spectral' label2=Amplitude screenratio=0.5
        ''')

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

Flow('Pzdiff','ElasticPzPSLR1wave ElasticPzPSLRwave','math s=${SOURCES[0]} m=${SOURCES[1]} output=m-s')

Flow('Pzdiff1','ElasticPzPSLR1wave Pzdiff','math s=${SOURCES[0]} m=${SOURCES[1]} output=m-s')

Result('ElasticPzPSLR1wave',
        '''
        graph label1=Z title='Lowrank pseudo-spectral' label2=Amplitude screenratio=0.5
        ''')

Result('Pzdiff',
         '''
         graph label1=Z title='Lowrank pseudo-spectral' label2=Amplitude screenratio=0.5
         ''')

Result('Pzdiff1',
         '''
         graph label1=Z title='Lowrank pseudo-spectral' label2=Amplitude screenratio=0.5
         ''')

End()