from rsf.proj import *
import fdmod

# Fetch Files from repository
modelFiles=['vp_marmousi-ii.segy'] #,'density_marmousi-ii.segy']
outputFiles=['vp'] #,'density']
for file in modelFiles:
    Fetch(file,"marm2")
# Convert Files to RSF
counter=0
for file in modelFiles:
    Flow(outputFiles[counter],file, ''' segyread tape=$SOURCE | put
        d1=.001249  d2=.001249 o1=0 o2=0 label1=Depth label2=Distance
        unit1=km unit2=km
                                    ''')
    Result(outputFiles[counter],'grey scalebar=y allpos=y barreverse=y wanttitle=n')
    counter = counter+1

# ------------------------------------------------------------
par = {
    'nt':10000, 'dt':0.0001,'ot':0,     'lt':'t','ut':'s',
    'kt':800,    # wavelet delay
    'nx':13601, 'ox':0,  'dx':0.001249, 'lx':'x','ux':'km',
    'nz':2801,  'oz':0,  'dz':0.001249, 'lz':'z','uz':'km',
    }
# add F-D modeling parameters
fdmod.param(par)

# ------------------------------------------------------------
# wavelet
Flow('wav',None,
         '''
         spike nsp=1 mag=1 n1=%(nt)d d1=%(dt)g o1=%(ot)g k1=%(kt)d |
         ricker1 frequency=15 |
         scale axis=123 |
         put label1=t label2=x label3=y
         ''' % par, split=[2,123], reduce='add')
Result('wav','window n1=1000 | graph title="" label1="t" label2=')
# ------------------------------------------------------------

# experiment setup
Flow('r_',None,'math n1=13601  d1=.001249  o1=0   output=0' % par)
Flow('s_',None,'math n1=1      d1=0        o1=0   output=0' % par)

# receiver positions
Flow('zr','r_','math output=".450" ') #'0'
Flow('xr','r_','math output="x1"')
Flow('rr',['xr','zr'],
     '''
     cat axis=2 space=n
     ${SOURCES[0]} ${SOURCES[1]} | transp
     ''', stdin=0, split=[2,1000],reduce='add')
Plot('rr',fdmod.rrplot('',par))

# source positions
Flow('zs','s_','math output=0')
Flow('xs','s_','math output=2')
Flow('rs','s_','math output=1')
Flow('ss',['xs','zs','rs'],
     '''
     cat axis=2 space=n
     ${SOURCES[0]} ${SOURCES[1]} ${SOURCES[2]} | transp
     ''', stdin=0, split=[2, 1000], reduce='add')
Plot('ss',fdmod.ssplot('',par))

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

# density
Flow('vel','vp',
     ''' 
     put o1=%(oz)g d1=%(dz)g  o2=%(oz)g d2=%(dz)g 
     ''' % par)

Plot('vp',fdmod.cgrey('allpos=y bias=1.5 pclip=98 scalebar=y',par))
Result('vel',['vp','rr','ss'],'Overlay')

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

# density
Flow('den','vel','math output=1')

#Flow('den','density','''
#    put o1=%(oz)g  d1=%(dz)g  o2=%(oz)g d2=%(dz)g
#    ''' % par)
#Plot('den',fdmod.cgrey('allpos=y bias=1.5 pclip=98',par))
#Result('den',['den','rr','ss'],'Overlay')

# ------------------------------------------------------------
# finite-differences modeling
fdmod.awefd('dat','wfl','wav','vel','den','ss','rr','free=y dens=n',par)

Result('wfl',fdmod.wgrey('pclip=99 min2=0 max2=4 min1=0 max1=1',par))
for item in ['10','25','45','65']:
    Result('time'+item,'wfl','''window f3=%s n3=1 min1=0 min2=0 min2=0 max2=4 min1=0 max1=1
              | grey gainpanel=a pclip=99 wantframenum=y 
              title=Wavefield\ at\ %s\ \(ms\) labelsz=4
              titlesz=6 screenratio=.18 screenht=2 wheretitle=t''' % (item,item))

Result('dat','transp |' + fdmod.dgrey('pclip=99 min2=0 max2=4 min1=0 max1=1',par))
End()