from rsf.proj import *
import fdmod as f

# ------------------------------------------------------------
# Sigsbee 2A parameters
par = {
    'nx':1601,'ox':25.0,'dx':0.025,'lx':'x',
    'nz':801, 'oz':4.50,'dz':0.025,'lz':'z',
    'nt':5001,'ot':0,   'dt':0.001,'lt':'t'
    }
par['jsnap']=200 # snapshot jump
par['kt']=100    # wavelet delay (samples) 
par['nb']=100    # boundary (grid points)

# ------------------------------------------------------------
# convert coordinates from ft to km
par['ft2km']=0.3048
par['ox']=par['ox']*par['ft2km']
par['dx']=par['dx']*par['ft2km']
par['oz']=par['oz']*par['ft2km']
par['dz']=par['dz']*par['ft2km']

# ------------------------------------------------------------
# set axes labels
par['ux']='km'
par['uz']='km'
par['ut']='s'

# ------------------------------------------------------------
# set plotting parameters
f.param(par)

# ------------------------------------------------------------
# source coordinates (exploding reflectors)
f.boxarray('ss',
           5,   # vertical number
           5,   # vertical origin 
           0.5, # vertical sampling
           14,  # horizontal number
           10,  # horizontal origin
           0.5, # horizontal sampling
           par)

Plot('ss',f.ssplot('plotfat=5 symbol=. plotcol=5',par))

# ------------------------------------------------------------
# horizontal array @ z=1.5km
f.horizontal('tH',1.5,par)
par['jrH']=10   # jump (grid points)
par['orH']=14.0 # origin
par['nrH']=75   # number
par['crH']=1    # color

# vertical array @ x=8.5km
f.vertical('tV',8.5,par)
par['jrV']=20   # jump (grid points)
par['orV']=2.5  # origin
par['nrV']=25   # number
par['crV']=2    # color

for j in ('H','V'):

    # window array
    Flow('r'+j,
         't'+j,
         'window j2=%d min2=%g n2=%d'%
         (par['jr'+j],  # jump
          par['or'+j],  # origin
          par['nr'+j])) # number

    # plot array
    Plot('r'+j,
         f.rrplot('plotcol=%d'%par['cr'+j],par))

# ------------------------------------------------------------
# merge receiver files
Flow('rA',
     'rH rV',
     'cat axis=2 space=n ${SOURCES[1]}')

# overlay receiver arrays
Plot('rA','rH rV','Overlay')

# ------------------------------------------------------------
# get stratigraphic velocity
strvelfile = 'sigsbee2a_stratigraphy.sgy'
Fetch(strvelfile,'sigsbee')
Flow('vstr-raw',strvelfile,'segyread read=data')

Flow('vstr',
     'vstr-raw',
     '''
     scale rscale=%g |
     put o1=%g d1=%g o2=%g d2=%g |
     window n1=%d min1=%g n2=%d min2=%g
     ''' % (par['ft2km']*0.001,
            0              ,0.025*par['ft2km'],
            10*par['ft2km'],0.025*par['ft2km'],
            par['nz'],par['oz'],
            par['nx'],par['ox']
            ))

# ------------------------------------------------------------
# make smooth velocity
Flow('vsmo',
     'vstr',
     'smooth rect1=25 rect2=25 repeat=3')

# ------------------------------------------------------------
# plot velocities
for v in ('vstr','vsmo'):
    Plot(  v,f.cgrey('allpos=y bias=1.43',par))

    # overlay sources and receivers
    Result(v,[v,'ss','rH','rV'],'Overlay')

# ------------------------------------------------------------
# make constant density
Flow('dens','vsmo','math output=1')
Plot('dens',f.cgrey('allpos=y bias=1.43',par))
Result('dens',['dens','ss','rH','rV'],'Overlay')

# ------------------------------------------------------------
# construct wavelet
f.wavelet('wav_',10,par)

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

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

# ------------------------------------------------------------
# run FD modeling
f.awefd1('tmpA',  # data file (all receivers)
         'wfld',  # wavefield snapshots
         'wav',   # source wavelet
         'vsmo',  # velocity
         'dens',  # density
         'ss',    # source coordinates
         'rA',    # receiver coordinates
         'free=n',# optional flags
         par)

# ------------------------------------------------------------
# plot wavefield frames
for i in range( (par['nt']-1)//par['jsnap']):
    tag = '-%02d'%i
    f.wframe('wfld'+tag,
             'wfld',i,'pclip=99',par)
    Result(  'wfld'+tag,
            ['wfld'+tag,'ss','rH','rV'],'Overlay')

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

# ------------------------------------------------------------
# window data from the horizontal array
Flow('datH',
     'datA',
     '''
     window squeeze=n n1=%d |
     put o1=%g d1=%g
     '''%(par['nrH'],
          par['orH'],
          par['jrH']*par['dx']))

Result('datH',       'window j2=4 | transp|'
       + f.dgrey(''%par,par))
Result('wigH','datH','window j2=4 | transp|'
       + f.dwigl('pclip=98'%par,par))

# ------------------------------------------------------------
# window data from the vertical array
Flow('datV',
     'datA',
     '''
     window squeeze=n f1=%d |
     put o1=%g d1=%g
     '''%(par['nrH'],
          par['orV'],
          par['jrV']*par['dz']))

Result('datV',       'window j2=4 |'
       + f.egrey(''%par,par))
Result('wigV','datV','window j2=4 |'
       + f.ewigl('pclip=98'%par,par))

# ------------------------------------------------------------
# run FD migration
for j in ('H','V','A'):
    f.zom('img'+j,  # image
          'dat'+j,  # data
          'vsmo',   # velocity
          'dens',   # density
          'r'+j,    # receiver coordinates
          'free=n', # optional flags
          par)

    # plot image
    Plot(  'img'+j,'bandpass flo=2 |'
           + f.cgrey('pclip=99.99',par))

    # overlay sources and receivers
    Result('img'+j,['img'+j,'ss','r'+j],'Overlay')

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