## 
 # Simple example of time-shift imaging condition (flat reflector)
 ##

from rsf.proj import *
import math
import zomig,spmig,pplot

# ------------------------------------------------------------
# PLOTTING
# ------------------------------------------------------------
def igrey(custom):
    return '''
    grey  labelrot=n wantaxis=y wanttitle=y title=""
    labelsz=12 titlesz=16
    %s 
    ''' % (custom)
def igraph(custom):
    return '''
    graph labelrot=n wantaxis=y wanttitle=y title=""
    yreverse=y wantaxis=n plotfat=6 plotcol=1 %s
    ''' % (custom)

def reflector(formula,par):
    return '''
    math n1=%d o1=%g d1=%g n2=1 output="%s"
    ''' % (par['nx']/2,par['ox'],par['dx'],formula)
# ------------------------------------------------------------
par = {
    'nz':400,  'dz':2.0,   'oz':0,
    'nx':400,  'dx':10.0,  'ox':0,
    'nt':5000, 'dt':0.0005,'ot':0, 'kt':201, 'f':15,
    'ns':1,    'ds':8,     'os':2000, 'js':1, 'fs':0,
#
    'nw':125, 'ow':0.8, 'jw':1,
#    
    'verb':'y', 'eps':0.01, 'nrmax':1, 'dtmax':0.00005,
    'tmx':32,'tmy':0,'pmx':32,'pmy':0,
#
    'kov-':0.9, 'kov0':1.0, 'kov+':1.1
    }
par['zmin']=par['oz']
par['zmax']=par['oz'] + (par['nz']-1) * par['dz']
par['xmin']=par['ox']
par['xmax']=par['ox'] + (par['nx']-1) * par['dx']

# ------------------------------------------------------------
# REFLECTIVITY
# ------------------------------------------------------------
Flow('left',None,reflector('402',par))
Flow('rght',None,reflector('402',par))
#Flow('left',None,reflector('x1',  par))
#Flow('rght',None,reflector('1000',par))

Flow('msk',None,
     '''
     spike nsp=1 mag=-1
           n1=%(nz)d d1=%(dz)g o1=%(oz)g
           n2=%(nx)d d2=%(dx)g o2=%(ox)g k2=1 l2=%(nx)d |
     math output=input+2 |
     smooth rect2=21
     ''' % par)

Flow('spk','left rght',
     '''
     cat ${SOURCES[1:2]} axis=1 |
     unif3 n1=%(nz)d d1=%(dz)g v00=1,2 |
     ai2refl | scale axis=12
     ''' % par)
Flow('ref','spk msk','math s=${SOURCES[0]} m=${SOURCES[1]} output=s*m ')
Plot  ('ref','ref',igrey('pclip=100 label1="z (m)" label2="x (m)" '))
#Result('ref','ref',igrey('pclip=100'))
Flow('r0','ref','transp plane=12 | transp plane=23')
# ------------------------------------------------------------
# VELOCITY/SLOWNESS
# ------------------------------------------------------------
Flow('vp','left rght',
     '''
     cat ${SOURCES[1:2]} axis=1 |
     unif3 n1=%(nz)d d1=%(dz)g v00=2000,2000
     ''' % par)
#Result('vp','vp',igrey('pclip=100 bias=900 allpos=y'))
Flow('sp','vp',
     '''
     transp |
     math "output=1/input" |
     spray axis=2 n=1 |
     put label2=y
     ''' % par )
Flow('ss','sp','math "output=2*input"')
# ------------------------------------------------------------
# WAVELET
# ------------------------------------------------------------
# wavelet for shot-profile modeling
Flow('wvlt',None,
     '''
     spike nsp=1 mag=1 k1=%(kt)d
     n1=%(nt)d d1=%(dt)g o1=0
     n2=1      d2=%(dx)g o2=%(os)g |
     ricker1 frequency=%(f)s |
     put label1=t label2=x label3=y
     unit2=m unit3=m
     ''' % par )
Result('wvlt','wvlt','window n1=500 | graph grid=y title=%s' % 'wvlt')
# wavelet for shot-profile migration
Flow('wave',None,
     '''
     spike nsp=1 mag=1 k1=1
     n1=%(nt)d d1=%(dt)g o1=0|
     put label1=t label2=x label3=y 
     ''' % par )
#Result('wave','wave','graph title=%s' % 'wave')

# ------------------------------------------------------------
# shot-profile modeling
# ------------------------------------------------------------
Flow('w0','wvlt',
     '''
     fft1 |
     window squeeze=n n1=%d min1=%g j1=%g |
     pad beg2=%d n2out=%d |
     put label1=w label2=x label3=y |
     transp plane=12 | transp plane=23
     ''' % (par['nw'],par['ow'],par['jw'],par['nx']/2,par['nx']))
Result('w0','w0','window | real | transp | grey pclip=100 wanttitle=n')

spmig.model   ('dpw','sp',     'w0','r0',par)
spmig.model_cw('dcw','sp','ss','w0','r0',par)

for k in ('p'):
    wdd = 'd' + k + 'w'
    tdd = 't' + k + 'w'

    Flow(tdd,wdd,
         '''
         window | transp | pad beg1=2 n1out=2501 |
         fft1 inv=y |
         window min1=0.1 |
         pad n1out=%(nt)d |
         put o1=0 o2=-2000 o3=%(os)g n3=%(ns)d d3=%(ds)g label3=s
         ''' % par)
    Plot  (tdd,'window | '
           + igrey('pclip=100 min1=0 max1=1.24975 label1="t" label2="h" '))
#    Result(tdd,'window | ' + igrey('pclip=100 label1=t grid=y title=%s' % tdd))

    dww = 'd' + k + 's'
    uww = 'd' + k + 'r'
    spmig.wflds(dww,uww,'wave',tdd,par)

    img = 'j' + k
    Plot(img,'window n4=1 n5=1 n6=1| transp | '
       + igrey('grid=y pclip=100 g2num=500 title=%s') % img)
# END k (wave type)
#Result('data',['tpw','tcw'],'Movie')
#Result('imag',['jp' ,'jc' ],'Movie')

Result('modeling','ref tpw','SideBySideIso')

#-((-1 + p^2)*Sqrt[h0^2 + z0^2]) + 
# p*Sqrt[h^2 + h0^2*(-1 + p^2) + p^2*z0^2]

v0=2000
z0=800
h0=1500

for n in ('-','0','+'):            # velocity: +,0,-
    v = par['kov' + n]
    hyper = 'hyper'+n
    Flow(hyper,'tpw',
         '''
         window n1=1 |
         math output="(%d+%g*sqrt(%d+x1*x1))/2000"
         ''' % ((1-v*v)*math.hypot(h0,z0),v,h0*h0*(v*v-1)+v*v*z0*z0))
    Plot(hyper+'_',hyper,igraph('min2=0 max2=1.24975 pad=n dash=1'))
    Plot(hyper,['tpw',hyper+'_'],'Overlay')
Result('hyper','hyper- hyper+','SideBySideIso')


# ------------------------------------------------------------
# shot-profile migration
# ------------------------------------------------------------
spmig.image   ('jp','sp',     'dps','dpr',par)
spmig.image_cw('jc','sp','ss','dcs','dcr',par)
# ------------------------------------------------------------

# ------------------------------------------------------------
# different velocities for imaging
# ------------------------------------------------------------
Flow('sp-','sp','math "output=%(kov-)g*input"' % par)
Flow('ss-','ss','math "output=%(kov-)g*input"' % par)

Flow('sp0','sp','math "output=%(kov0)g*input"' % par)
Flow('ss0','ss','math "output=%(kov0)g*input"' % par)

Flow('sp+','sp','math "output=%(kov+)g*input"' % par)
Flow('ss+','ss','math "output=%(kov+)g*input"' % par)
# ------------------------------------------------------------

# migration
for h in ('o','x','t','m'):           # Imaging: o,x,t            
    locpar = par
    if(h=='o'): locpar['misc']='itype=o'
    if(h=='t'): locpar['misc']='itype=t nht=160 oht=-0.200 dht=0.0025'
    if(h=='x'): locpar['misc']='itype=x hsym=y nhx=40'
    if(h=='m'): locpar['misc']='itype=t nht=160 oht=-0.200 dht=0.0025'

    for k in ('p'):               # P-waves; C-waves
        for n in ('-','0','+'):   # velocity: +,0,-
            img = 'i' + h + k + n
            slo = 's'     + k + n
            dds = 'd'     + k + 's'
            ddr = 'd'     + k + 'r'
            cig = 'h' + h + k + n

            # images
            spmig.imagePW3(img,cig,slo,dds,ddr,locpar)

# ------------------------------------------------------------
# angle-gather overlays
# ------------------------------------------------------------
for n in ('-','0','+'):            # velocity: +,0,-
    kov = 'kov' + n

    tov = 'tov' + n
    tmv = 'tmv' + n
    Flow(tmv,None,
         '''
         spike n1=400 o1=0 d1=25 |
         math output="400*sqrt(1 + x1^2 * 0.0005^2*(1-%g^2) )" |
         window max1=4250 
         ''' % par[kov])
    Plot(tmv,tmv,igraph('min1=0 max1=4000 min2=%(zmin)g max2=%(zmax)g') % par)

    Flow('vet',None,'spike nsp=1 mag=1 n1=%(nz)d o1=%(oz)g d1=%(dz)g'% par)
    vet = 'vet' + n
    Flow(vet,'vet','math output="%s*input"' % str(2000/par[kov]) )
    Plot(vet,vet,igraph('transp=y plotcol=2 min2=0 max2=4000 min1=%(zmin)g max1=%(zmax)g') % par)
    Plot(tov,[tmv,vet],'Overlay')

    xov = 'xov' + n
    Flow(xov,None,
         '''
         spike n1=400 o1=-4 d1=0.02 |
         math output="400*sqrt( 1/%g^2 + x1^2 * (1/%g^2-1) )" |
         window min1=-2.1 max1=2.1 
         ''' % (par[kov],par[kov]) )
    Plot(xov,xov,igraph('min1=0 max1=2.5 min2=%(zmin)g max2=%(zmax)g') % par)

# CAG
for h in ('x','t'):
    for k in ('p'):
        for n in ('-','0','+'):
            if(n=='-'): tit='title="s<s\s60 \_0\^\s100 "'
            if(n=='0'): tit='title="s=s\s60 \_0\^\s100 "'
            if(n=='+'): tit='title="s>s\s60 \_0\^\s100 "'

            img = 'i' + h + k + n
            cig = 'h' + h + k + n
            ssk = 's' + h + k + n
            ssko= 'c' + h + k + n

            off = 'o' + h + k + n

            ang = 'a' + h + k + n

            slo = 's' + k + n
            vel = 'v' + k + n
            dip = 'd' + k + n

            # offset gather
            Flow(off,cig,'window | transp plane=12 | stack')

            if(h=='x'):
                Plot(off,
                     igrey('pclip=100 label1=z unit1=m label2=h unit2=m %s') % tit)
                Flow(ssk,off,'radon adj=y p0=-4 np=200 dp=0.04')
                Plot(ssk,
                     igrey('pclip=100  min2=0 max2=2.5  label1=z unit1=m label2="tan \F10 q\F3 "  %s') % tit)
                ovl = 'xov' + n

                Flow(ang,ssk,'tan2ang a0=0 na=150 da=0.45')
                Plot(ang,
                     igrey('pclip=100 label1=z unit1=m label2="\F10 q\F3 (\^o\_)"  %s') % tit)

            if(h=='t'):
                Plot(off,
                     igrey('pclip=100 label1=z unit1=m label2="\F10 t\F3" unit2=s %s') % tit)
                Flow(ssk,off,'radon adj=y p0=0  np=200 dp=50')
                Plot(ssk,
                     igrey('pclip=99.9 min2=0 max2=4000 label1=z unit1=m label2="\F10 n\F3 " unit2="m/s" %s') % tit)
                ovl = 'tov' + n

                Flow(vel,slo,'window n1=1 f1=200 | math output=1/input')
                Flow(dip,vel,'math output=0')
                Flow(ang,[ssk,vel,dip],
                     'tshift cos=y a0=0 na=150 da=0.45 velocity=${SOURCES[1]} dip=${SOURCES[2]}')
                Plot(ang,
                     igrey('pclip=99.9 label1=z unit1=m label2="\F10 q\F3 " unit2="\^o\_"  %s') % tit)

            Plot(ssko,[ssk,ovl],'Overlay')

#Result('xoffall',['hxp-','hxp0','hxp+'],'Movie')
#Result('toffall',['htp-','htp0','htp+'],'Movie')
#Result('xangall',['cxp-','cxp0','cxp+'],'Movie')
#Result('tangall',['ctp-','ctp0','ctp+'],'Movie')

#Result('xoff',['hxp-','hxp0','hxp+'],'SideBySideIso')
#Result('toff',['htp-','htp0','htp+'],'SideBySideIso')
#Result('xang',['cxp-','cxp0','cxp+'],'SideBySideIso')
#Result('tang',['ctp-','ctp0','ctp+'],'SideBySideIso')

Result('xoff',['oxp-','oxp0','oxp+'],'OverUnderIso')
Result('toff',['otp-','otp0','otp+'],'OverUnderIso')
Result('xssk',['cxp-','cxp0','cxp+'],'OverUnderIso')
Result('tssk',['ctp-','ctp0','ctp+'],'OverUnderIso')
Result('xang',['axp-','axp0','axp+'],'OverUnderIso')
Result('tang',['atp-','atp0','atp+'],'OverUnderIso')

pplot.p3x2('off','oxp-','oxp0','oxp+','otp-','otp0','otp+',0.30,0.30,-11,-13)
pplot.p3x2('ssk','cxp-','cxp0','cxp+','ctp-','ctp0','ctp+',0.30,0.30,-11,-13)
pplot.p3x2('ang','axp-','axp0','axp+','atp-','atp0','atp+',0.30,0.30,-11,-13)

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

End()