from rsf.proj import *
import fdmod,wefd,pcsutil,marm2,adcig,stiff,wemig

# ------------------------------------------------------------
par = {
    'nt':10000,'ot':0.00, 'dt':0.00025,  'lt':'Time',     'ut':'s',
    'nx':1200, 'ox':5.00, 'dx':0.002498, 'lx':'Position', 'ux':'km',
    'nz':500,  'oz':0.50, 'dz':0.002498, 'lz':'Depth',    'uz':'km',
    'kt':200,
    'frq':30,
    'jsnap':500,
    'jdata':8,
    'nb':100,
    'oa':-85, 'na':340, 'da':0.5,
    'oap':0, 'nap':80, 'dap':1,
    'ng':1800,'dg':0.2,'og':-90,
    'wclip':0.5,
    'wweight':100,
    'clip':[99,99,99,99]
    }
par['jsnap']=par['nt']
fdmod.param(par)
par['labelattr']=' wantaxis=y '

par['xsou']=7.0
par['zsou']=par['oz']
par['zrec']=0.50

par['xmin'] = par['ox']
par['xmax'] = par['ox'] + (par['nx']-1)*par['dx']

# taper params
par['ltap']=101
par['rtap']=par['nx']-100

# ------------------------------------------------------------
# image coordinates
par['nqz']=par['nz']#/3*2
par['oqz']=par['oz']#+par['nz']/3*par['dz']
par['dqz']=par['dz']

par['nqx']=par['nx']/4
par['oqx']=par['ox']+3*par['nx']/8*par['dx']
par['dqx']=par['dx']

fdmod.boxarray('qq',
               par['nqz'],par['oqz'],par['dqz'],
               par['nqx'],par['oqx'],par['dqx'],
               par)
Plot('qq','window j2=61 |' + fdmod.qqplot('',par))

# ------------------------------------------------------------
# receiver coordinates
fdmod.horizontal('rr',par['zrec'],par)
Plot('rr','window j2=10 |' + fdmod.rrplot('symbol=o plotfat=10',par))

# ------------------------------------------------------------
marm2.data(par)

# ------------------------------------------------------------
marm2.dip('dipall','rx',par)

# ------------------------------------------------------------
for file in (['vp','rx']):
    Plot(file,fdmod.cgrey('allpos=y pclip=99 labelsz=8 bias=1',par))


# elastic stiffness tensor
Flow('zero','vp','math output=0')
stiff.tti2d('co','vp','vs','ro','zero','zero','zero',par)
stiff.tti2d('cx','vp','vs','rx','zero','zero','zero',par)
# ------------------------------------------------------------
# wavelet
fdmod.wavelet('wav0',par['frq'],par)

Flow('ver','wav0','math output=input*1' % par)
Flow('hor','wav0','math output=input*0' % par)
Flow('wave','ver hor',
     '''
     cat axis=2 space=n ${SOURCES[1:2]} |
     transp plane=12 |
     spray axis=1 n=1 o=%(ox)g d=%(dx)g
     ''' % par)
wefd.ewavelet('wave','',par)

# ------------------------------------------------------------
shots = range(par['nx']//2-400,par['nx']//2+420,40)
#shots = range(par['nx']/2,par['nx']/2+20,20)
#print shots
# ------------------------------------------------------------
for ix in shots:
    tag = "-%04d-" % ix
    xsou = par['ox']+ix*par['dx']

    fdmod.point('ss'+tag,xsou,par['zsou'],par)
    Plot(       'ss'+tag,'window       |' + fdmod.ssplot('',par))

allsou =  map(lambda x: 'ss-%04d-' % x,shots)
Plot('ss',allsou,'Overlay')
Result('vp','vp ss ','Overlay')
Result('rx','rx rr ss ','Overlay')

vptop = 1.67
# ------------------------------------------------------------
for ix in shots:
    tag = "-%04d-" % ix
    xsou = par['ox']+ix*par['dx']

    tmin = (xsou-par['xmin'])/vptop
    tmax = (par['xmax']-xsou)/vptop

    marm2.mask('me'+tag,xsou,tmin,tmax,par)

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

nhx=120
nhz=0
nht=0

xcig=6.5
par['nhx']=120
par['nhz']=0
par['nht']=0
par['dht']=par['dt']*4
par['xcig']=6.5
# ------------------------------------------------------------
# cig coordinates
fdmod.vertical('cc',xcig,par)
Plot('cc',fdmod.rrplot('symbol=o plotfat=10',par) )
par['ncz']=par['nz'];par['ocz']=par['oz'];par['dcz']=par['dz'];
par['ncx']=1;par['ocx']=1;par['dcx']=0;


#dip angle at cig location x
Flow('dipone','dipall','window n2=1 min2=%g'%xcig)

#vpvs ratio at cig location x
Flow('vratioPP','vratio','window n3=1 f3=0 n2=1 min2=%g'%xcig)
Flow('vratioPS','vratio','window n3=1 f3=1 n2=1 min2=%g'%xcig)

# ------------------------------------------------------------
# migration
for ix in shots:
    tag = "-%04d-" % ix

    # modeling
    fdmod.ewefd2('de'+tag,
                'we'+tag,
                'wave',
                'cx','rx',
                'ss'+tag,
                'rr',' ssou=y opot=n free=n ',par)
    Flow('ee'+tag,['de'+tag,'me'+tag],
         '''
         add mode=p ${SOURCES[1]} |
         transp plane=13 | bandpass flo=10 | transp plane=13
         ''')

    wefd.edata('de'+tag,'de'+tag,'pclip=98',par)
    wefd.edata('ee'+tag,'ee'+tag,'pclip=98',par)

    # elastic RTM
    wefd.ewfld('je'+tag,'wave','ee'+tag,'co','ro','ss'+tag,'rr','qq',' opot=y jdata=%(jdata)d '%par,par)
    wefd.ecic('je'+tag,'ss','rr','cc','',par)

    # angle gathers
    wefd.eeic('je'+tag,'ss','rr','cc',xcig,'',par)







# ------------------------------------------------------------
# all shots CIG
allppang =  map(lambda x: 'je-%04d-Eang11' % x,shots)
allpsang =  map(lambda x: 'je-%04d-Eang12' % x,shots)
nshots=len(allppang)
Flow('PPang',allppang,'add ${SOURCES[0:%d]} | smooth rect2=1'%nshots,stdin=0)
Flow('PSang',allpsang,'add ${SOURCES[0:%d]} | smooth rect2=1'%nshots,stdin=0)
Result('PPang', adcig.agrey(' title=  grid=y ',par))
Result('PSang', adcig.agrey(' title=  grid=y ',par))


# stack CIG, then do decomp
allppcig =  map(lambda x: 'je-%04d-Ecig11' % x,shots)
allpscig =  map(lambda x: 'je-%04d-Ecig12' % x,shots)

Flow('PPcig',allppcig,
     '''
     add ${SOURCES[0:%d]} |
     window f2=80 n2=80 
     '''%nshots,stdin=0)
Flow('PScig',allpscig,
     '''
     add ${SOURCES[0:%d]} |
     window f2=80 n2=80 
     '''%nshots,stdin=0)
Result('PPcig',adcig.xgrey('pclip=98',par))
Result('PScig',adcig.xgrey('pclip=98',par))

Flow( 'PPcig-ang',
             ['PPcig','dipone','vratioPP'],
             adcig.cig2ssk(300,-1.5,0.01) + '|' +
             adcig.xsk2ang(240, -60,0.50))
Flow( 'PScig-ang',
             ['PScig','dipone','vratioPS'],
             adcig.cig2ssk(300,-1.5,0.01) + '|' +
             adcig.xsk2ang(240, -60,0.50))




for icig in ('PPcig-ang','PScig-ang'):
    Result(icig,
           '''
           bandpass flo=20 |
           smooth rect2=20 |
           ''' + adcig.agrey('pclip=99',par))






End()