from rsf.proj import *
from rsf.recipes.m2d import *


model = {
    'X':    5000,
    'Z':    5000,
    'T':    1.0,
    'dt':   0.0015,
    'SelT':  1.0,
    'dx':   10.0,
    'dz':   10.0,
    'snpintvl': 1.0,
    'size'    : 4,   # FD order
    'frqcut'  : 1.0,
    'pml'     : 0,
    'vel'     : '4000',
    'den'     :  '2100'
}

# source & receiver
srp = {
     'bgn'   : 0.1,     # s, time of maximum ricker
     'frq'   : 20.0,     # source domain frequence
     'srcmms': 'y',      # point source
     'inject': 'y',      # if y, inject; if n, Initiate conditon
     'slx'   : 2500,   # source location (x), meter
     'slz'   : 2000,   # source location (x), meter
     'gdep'  : 200      # receiver location (z), meter
     }

par = setpar(model, srp)
# --------------------------------------------------------------------
Flow('vel1',None,
     '''
     spike n1=250 k1=98 n2=%(nx)d k2=98 d1=%(dz)g d2=%(dx)g|
     math output="1300"
     ''' %par)

Flow('vel2',None,
     '''
     spike n1=251 k1=98 n2=%(nx)d k2=98 d1=%(dz)g d2=%(dx)g|
     math output="3200"
     '''%par)
#3200

Flow('vel','vel1 vel2',
     '''
     cat axis=1 ${SOURCES[1]} | 
     smooth rect1=3 rect2=3   
     ''')


Flow('den1',None,
     '''
     spike n1=250 k1=98 n2=%(nx)d k2=98 d1=%(dz)g d2=%(dx)g|
     math output="1700"
     '''%par)

Flow('den2',None,
     '''
     spike n1=251 k1=98 n2=%(nx)d k2=98 d1=%(dz)g d2=%(dx)g|
     math output="2500"
     '''%par)

Flow('den','den1 den2', 'cat axis=1 ${SOURCES[1]} |smooth rect1=3 rect2=3 ')

for m in ['den','vel']:
    pml  = m+'_pml'
    iwname = 'iw'+m
    Flow(pml,m,
         '''
         expand left=%(bd)d right=%(bd)d 
                top=%(bd)d  bottom=%(bd)d
         '''%par)
    Flow(iwname, m,'math output="input/1000" ')

Result('den',
       '''
       math output="input/1000" |
       put d1=%g d2=%g unit1="km" unit2="km" label2="Distance" label1="Depth"  |
       grey color=j scalebar=y allpos=y transp=y screenratio=1
       wanttitle=n barlabel='Density' barunit='g/cm\^3\_'
       bartype=h   labelsz=6  
       '''%(par['dz']/1000,par['dx']/1000))
Result('tlvel','vel',
       '''
       math output="input/1000" |
       put  d1=%g d2=%g unit1="km" unit2="km" label2="Distance" label1="Depth"|
       grey color=i scalebar=n allpos=y transp=y screenratio=1
       wanttitle=n labelsz=6  polarity=y
       bartype=h    barlabel='Velocity' barunit='km/s'
       '''%(par['dz']/1000,par['dx']/1000))



# --------------------------------------------------------------------
Flow('srcp', None,
         '''
         spike n1=%(nt)d d1=%(dt)g k1=%(bgnp)g |
         ricker1 frequency=%(frq)g | 
         scale axis=1 
         '''%par)
# --------------------------------------------------------------------


# lr
lrw = 'lrwav'
lrr = 'lrrec'
lrs = 'lrsnap'
lrtr= 'lrtr'
lrtrfft = 'lrtrfft'
sglr2(lrw, lrr, 'srcp', 'vel', 'den', par, '', 0)
Flow(lrs, lrw, 'window n3=1 f3=%(snt)d ' %par)
Flow(lrtr,lrs, 'window n2=1 f2=250 | scale axis=1 ')

Result(lrs,
       '''
       put unit1="km" unit2="km" d1=%g d2=%g label2="Distance" label1="Depth"|
       grey  screenratio=1 pclip=95 
       title="SGL method"
       '''%(par['dz']/1000,par['dx']/1000))

# lfd
lfdw = 'lfdwav'
lfdr = 'lfdrec'
lfds = 'lfdsnap'
lfdtr= 'lfdtr'
lfdtrfft = 'lfdtrfft'
sglfd2(lfdw, lfdr, 'srcp', 'vel_pml', 'den_pml', par, '', 0)
Flow(lfds, lfdw, 'window n3=1 f3=%(snt)d ' %par)
Flow(lfdtr,lfds, 'window n2=1 f2=250 | scale axis=1 ')

Result(lfds,
       '''
       put unit1="km" unit2="km" d1=%g d2=%g label2="Distance" label1="Depth" |
       grey  screenratio=1 pclip=95
        title="Fourth-order SGLFD method" 
       '''%(par['dz']/1000,par['dx']/1000))
# iwave
iww = 'iwwav'
iwr = 'iwrec'
iws = 'iwsnap'
iwtr= 'iwtr'
iwtrfft = 'iwtrfft'
#iw(iww, iwr, 'srcp', 'iwvel', 'iwden', par, '', 0)
#Flow(iws, 'movie_p', 'window n3=1 f3=%(snt)d ' %par)
#Flow(iwtr,iws, 'window n2=1 f2=250 | scale axis=1 ')
#Result(iws,
#       '''
#       put unit1="km" unit2="km" d1=%g d2=%g label2="Distance" label1="Depth" |
#       grey screenratio=1 pclip=95
#       title="Fourth-order SGFD method"
#       '''%(par['dz']/1000,par['dx']/1000))


#Plot('iwtr',
#     '''
#     put label1='Depth' unit1='km' d1=%g |
#     graph  wantaxis2=n wherexlabel=top title="Fourth-order SGFD method"  labelsz=12 titlesz=12
#            wheretitle=bottom  yreverse=n transp=n plotfat=7 plotcol=7  screenratio=0.3 screenht=8 
#     '''%(par['dz']/1000))
#Plot('lrtr',
#     '''
#     put label1='Depth' unit1='km' d1=%g | 
#     graph  wantaxis2=n wherexlabel=top title="SGL method"  labelsz=12 titlesz=12
#            wheretitle=bottom yreverse=n transp=n plotfat=7 plotcol=7  screenratio=0.3 screenht=8 
#     '''%(par['dz']/1000))

#Result('trace','iwtr lrtr','OverUnderIso')

# ------------------------------------------------------------------------------

model['dt'] = 0.002
model['size'] = 4
par = setpar(model, srp)

Flow('srcp1', None,
         '''
         spike n1=%(nt)d d1=%(dt)g k1=%(bgnp)g |
         ricker1 frequency=%(frq)g | 
         scale axis=1 
         '''%par)

# lfd
lfdw = 'lfd4wav2'
lfdr = 'lfd4rec2'
lfds = 'lfd4snap2'
lfdtr= 'lfd4tr2'
lfdtrfft = 'lfdtrfft2'
sglfd2(lfdw, lfdr, 'srcp1', 'vel_pml', 'den_pml', par, '', 1)
Flow(lfds, lfdw, 'window n3=1 f3=%(snt)d ' %par)
Flow(lfdtr,lfds, 'window n2=1 f2=250 | scale axis=1 ')
Result(lfds,
       '''
       put unit1="km" unit2="km" d1=%g d2=%g label2="Distance" label1="Depth"|
       grey  screenratio=1 pclip=95
       wanttitle=y  title="Fourth-order SGLFD method" 
       '''%(par['dz']/1000,par['dx']/1000))
# iwave
#iww = 'iwwav2' 
#iwr = 'iwrec2'
#iws = 'iwsnap2' 
#iwtr= 'iwtr2'
#iwtrfft = 'iwtrfft2'
#iw(iww, iwr, 'srcp1', 'iwvel', 'iwden', par, '', 1)
#Flow(iws, 'movie_p', 'window n3=1 f3=%(snt)d ' %par)
#Flow(iwtr,iws, 'window n2=1 f2=250 | scale axis=1 ')
#Result(iws,
#       '''
#       put unit1="km" unit2="km" d1=%g d2=%g label2="Distance" label1="Depth"|
#       grey  screenratio=1 pcip=95
#       wanttitle=y title="Fourth-order SGFD method" 
#       '''%(par['dz']/1000,par['dx']/1000))


# --------------------------------------------------------------------
model['dt']=0.002
model['size']=8
par = setpar(model, srp)

Flow('srcp2', None,
         '''
         spike n1=%(nt)d d1=%(dt)g k1=%(bgnp)g  |
         ricker1 frequency=%(frq)g | 
         scale axis=1 
         '''%par)

for m in ['den','vel']:
    pml  = m+'8_pml'
    Flow(pml,m,
         '''
         expand left=%(bd)d right=%(bd)d 
                top=%(bd)d  bottom=%(bd)d
         '''%par)

# lfd
lfdw = 'lfd8wav2'
lfdr = 'lfd8rec2'
lfds = 'lfd8snap2'
lfdtr= 'lfd8tr2'
lfdtrfft = 'lfd8trfft2'
sglfd2(lfdw, lfdr, 'srcp2', 'vel8_pml', 'den8_pml', par, '', 2)
Flow(lfds, lfdw, 'window n3=1 f3=%(snt)d ' %par)
Flow(lfdtr,lfds, 'window n2=1 f2=250 | scale axis=1 ')
Result(lfds,
       '''
       put unit1="km" unit2="km" d1=%g d2=%g label2="Distance" label1="Depth"|
       grey  screenratio=1 pclip=95
       wanttitle=y title="Eighth-order SGLFD method"
       '''%(par['dz']/1000,par['dx']/1000))


# --------------------------------------------------------------------
Plot('lfd4tr2',
       '''graph  wantaxis2=n wherexlabel=top title="Fourth-order SGLFD method"  wheretitle=bottom  yreverse=n transp=n plotfat=7 plotcol=7  screenratio=0.3 screenht=8 
       ''')
Plot('lfd8tr2',
       '''graph  wantaxis2=n wherexlabel=top title="Eighth order SGLFD method"  wheretitle=bottom  yreverse=n transp=n plotfat=7 plotcol=7 screenratio=0.3 screenht=8 
       ''')

Result('tracelfd','lfd4tr2 lfd8tr2','OverUnderIso')




End()