from rsf.proj import *
import sigsbee
import rsf.recipes.fdmod as fdmod

# ------------------------------------------------------------
par=sigsbee.paramwin() # Sigsbee2A parameters
par['nt']=2001         # time steps (samples)
par['kt']=100          # wavelet delay (samples) 
par['dt']=0.001        # time sampling (ms)
par['nb']=100          # boundary (grid points)
fdmod.param(par)       # plotting parameters

# ------------------------------------------------------------
# source coordinates (exploding reflectors)
fdmod.boxarray('ss',5,2,1,12,8,1,par)

# plot sources
Plot('ss',fdmod.ssplot('plotfat=10 symbol=.',par))

# ------------------------------------------------------------
par['jr']=4   # receiver jump (grid points)
par['nr']=100 # number of receivers
par['fr']=500 # receivers origin (grid points)

# receiver coordinates
fdmod.horizontal('tt',par['oz']+par['dz'],par)
Flow('rr',
     'tt',
     'window n2=%(nr)d j2=%(jr)d f2=%(fr)d'%par)

# plot receivers
Plot('rr',fdmod.rrplot('plotfat=10',par))

# ------------------------------------------------------------
# get velocity
sigsbee.getstrvelwin('vstr',par)
Flow(  'velo',
       'vstr',
       'smooth rect1=100 rect2=100 repeat=1')

# plot velocity
Plot(  'velo',fdmod.cgrey('allpos=y bias=1.43',par))
Result('velo',['velo','ss','rr'],'Overlay')

# density
Flow('dens','velo','math output=1')

# ------------------------------------------------------------
# construct wavelet
fdmod.wavelet('wav_',15,par)

# transpose wavelet
Flow(  'wav','wav_','transp')

# plot wavelet
Result('wav','window n2=1000 |'
       + fdmod.waveplot('',par))

# ------------------------------------------------------------
# run FD modeling
fdmod.awefd1('temp','wfld',
             'wav','velo','dens',
             'ss','rr',
             'free=n',par)

# generate wavefield movie
Plot('wfld',fdmod.wgrey('pclip=99',par),view=1)

# plot wavefield frames
for i in range(7):
    tag = '-%02d' %(i)
    fdmod.wframe('wfld'+tag,
                 'wfld',i,'pclip=99',par)
    Result( 'wfld'+tag,
           ['wfld'+tag,'ss','rr'],'Overlay')

# undo wavelet delay
Flow('data','temp',
     '''
     window squeeze=n f2=%(kt)d |
     pad end2=%(kt)d |
     put o2=%(ot)g
     ''' %par)

# plot data
Result('data','window j2=4 | transp |'
       + fdmod.dgrey('',par))

# ------------------------------------------------------------
# run FD migration
fdmod.zom('imag','jdat',
          'data','velo','dens',
          'rr','rr',
          'free=n',par)

# plot image
Plot(  'imag','bandpass flo=2 |'
       + fdmod.cgrey('pclip=99.99',par))
Result('imag',['imag','rr'],'Overlay')

# ------------------------------------------------------------
End()