from rsfproj import *
from math import *
import sys
import string, spmig, adcig

# ------------------------------------------------------------
par = {
    'nx':800, 'dx':10,    'ox':0,
    'nz':500, 'dz':5,     'oz':0,
    'nt':4000,'dt':0.004, 'ot':0,
    'nh':800, 'dh':10,    'oh':-1000,
    'ns':1,   'ds':20,    'os':1000, 'fs':0, 'js':1,
    'nw':600, 'jw':1,
    'velp':3200,
    'vels':800,
    'nhx':100,
    'nht':200, 'oht':-0.125, 'dht':0.00125,
    'na':200, 'oa':-10, 'da':0.5
    }
par['ow']=0.
par['dw']=1./(par['nt']*par['dt'])
par['vpvs']= par['velp'] / par['vels']

par['amin'] = par['oa']
par['amax'] = par['oa'] + (par['na']-1)*par['da']
# ------------------------------------------------------------
def igrey(custom):
    return '''
    grey labelrot=n title="" pclip=100 wantaxis=y
    grid=n gridcol=0
    labelsz=6
    %s
    ''' % custom
def igraph(custom):
    return '''
    graph labelrot=n title="" transp=y yreverse=y
    wantaxis=n wantlabels=n
    min2=%g max2=%g min1=0 max1=2500
    %s
    ''' % (par['amin'],par['amax'],custom)
# ------------------------------------------------------------
# CIG to slant-stack
def tcig2ssk():
    return '''
    slant adj=y p0=1000 np=600 dp=10 |
    put label2=v
    '''
def hcig2ssk():
    return '''
    slant adj=y p0=-0.5 np=300 dp=0.01 |
    put label2=tan
    '''
# slant-stack to angle
def tssk2ang():
    return '''
    pp2pstsic a0=%(oa)g na=%(na)d da=%(da)g
    velocity=${SOURCES[1]}
    vpvs=${SOURCES[2]}
    dip=${SOURCES[3]} |
    put label2=ang
    ''' % par
def hssk2ang():
    return '''
    pp2psang2
    vpvs=${SOURCES[1]}
    dip=${SOURCES[2]} |
    tan2ang a0=%(oa)g na=%(na)d da=%(da)g |
    put label2=ang
    ''' % par
# ------------------------------------------------------------

ISCR = ' screenratio=0.5 screenht=7'
OSCR = ' screenratio=1.5 screenht=10'

ilabel = 'label1="z" label2="x"              unit1=m unit2=m'

olabelx = 'label1="z" label2="h"             unit1=m unit2=m'
olabelt = 'label1=z unit1=m label2="\F10 t\F3" unit2=s'

slabelx = 'label1=z unit1=m label2="tan \F10 q\F3" unit2="."'
slabelt = 'label1=z unit1=m label2="\F10 n\F3" unit2=m/s'

alabel = 'label1=z unit1=m label2="\F10 q\F3" unit2="\^o\_"'

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

deg  = (0,15,30,45,60)
ndeg = len(deg)

# ------------------------------------------------------------
ppan = string.join(map(lambda a: str(a),deg),',')
psan = string.join(map(lambda a: str(
    atan(tan( ( asin(sin(a*pi/180)/par['vpvs']) + a*pi/180 )/2 ))/pi*180
    ),deg),',')

tppan = string.join(map(lambda a: str(
    tan(a*pi/180)
    ),deg),',')
tpsan = string.join(map(lambda a: str(
    tan( ( asin(sin(a*pi/180)/par['vpvs']) + a*pi/180 )/2 )
    ),deg),',')

dpth = string.join(map(lambda a: str(
    ( 300+1000*tan(a*pi/180) ) / 5
    ),deg),',')

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

Flow('dip',None,
     'spike n1=%d d1=%g o1=%g n2=%d nsp=%d k2=%s mag=%s' %
     (par['nx'],
      par['dx'],
      par['ox'],
      ndeg,
      ndeg,
      string.join(map(str,range(1,1+ndeg)),','),
      string.join(map(lambda a: str(tan(a*pi/180)),deg),',')
      )
     )
Flow('ref','dip','math output="300+input*x1" ')
Plot('ref',
     '''
     graph yreverse=y min2=0 max2=2500 min1=0 max1=5000
     title="" wantaxis=n
     ''' + ISCR)

# wavelet
Flow('wav',None,
     '''
     spike nsp=1 mag=1 k1=1
     n1=%(nt)d d1=%(dt)g o1=0
     ''' % par)

# velocity/slowness
Flow('vel',None,
     '''
     spike nsp=1 mag=1
     n1=%(nz)d d1=%(dz)g o1=%(oz)g
     n2=%(nx)d d2=%(dx)g o2=%(ox)g |
     transp |
     spray axis=2 n=1 o=0 d=1 |
     put label1=x label2=y label3=z
     ''' % par)
Flow('slop','vel','math output=input/%(velp)d' % par)
Flow('slos','vel','math output=input/%(vels)d' % par)

Flow('vp','slop','window n1=1 min1=1000 | math output=%(velp)d' % par)
Flow('vs','slos','window n1=1 min1=1000 | math output=%(vels)d' % par)

# vpvs ratio
Flow('vpvs',None, 'spike nsp=1 mag=%(vpvs)g n1=%(nz)d o1=%(oz)g d1=%(dz)g' % par)
Flow('vpvp',None, 'spike nsp=1 mag=1        n1=%(nz)d o1=%(oz)g d1=%(dz)g' % par)

# dip angle
Flow('dipa',None,
     '''
     spike nsp=5 mag=%s
     n1=%d o1=%g d1=%g
     k1=0,76,141,221,351 l1=75,140,220,350,500
     ''' % (tppan,par['nz'],par['oz'],par['dz']))

# PP reflection angles
Flow('ppovl',None,
     '''
     spike nsp=5 mag=%s n1=5 k1=1,2,3,4,5 |
     spray axis=2 n=%d o=%g d=%g |
     transp
     ''' % (ppan,par['nz'],par['oz'],par['dz']))
Plot('ppovl',igraph(''+OSCR))

# PS reflection angles
Flow('psovl',None,
     '''
     spike nsp=5 mag=%s n1=5 k1=1,2,3,4,5 |
     spray axis=2 n=%d o=%g d=%g |
     transp
     ''' % (psan,par['nz'],par['oz'],par['dz']))
Plot('psovl',igraph(''+OSCR))
# ------------------------------------------------------------

for case in ('pp','ps'):

    # data
    Flow(case,'ref dip',
         '''
         kirmod vel=%d vel2=%d dip=${SOURCES[1]}
         nt=%d dt=%g t0=%g     freq=15
         nh=%d dh=%g h0=%g
         ns=%d ds=%g s0=%g |
         tpow tpow=1 |
         window f2=100 | pad beg2=100 |
         put label1=t label2=h
         ''' %
         ( par['velp'],
          (par['velp'],par['vels'])[case=='ps'],
           par['nt'],par['dt'],par['ot'],
           par['nh'],par['dh'],par['oh'],
           par['ns'],par['ds'],par['os'])
         )
    Result(case,'window max1=8 |'
           + igrey('label1="t" label2="h" unit1=s unit2=m screenratio=0.5 screenht=7'))

    # migration
    sou = case + 'sou'
    rec = case + 'rec'
    spmig.wflds(sou,rec,'wav',case,par)
#    Result(rec,'window | real |' + igrey(''))

    for TYP in ('h','t'): # imaging condition
        # imaging condition parameters
        if(TYP=='h'):
            par['misc']='itype=x jcx=50 nhx=%(nhx)d nhz=1 hsym=y' % par
        else:
            par['misc']='itype=t jcx=50 nht=%(nht)d oht=%(oht)g dht=%(dht)g' % par

        # migration
        img = case + 'img' + '-'+TYP
        cig = case + 'cig' + '-'+TYP
        if(case=='pp'):
            spmig.imagePW3(img,cig,'slop',       sou,rec,par)
        else:
            spmig.imageCW3(img,cig,'slop','slos',sou,rec,par)
        Plot(img,'window max1=5000 | transp |'+ igrey(ilabel+ISCR))
        Result(img,[img,'ref'],'Overlay')

        # offset CIGs
        off = case + 'off' + '-'+TYP
        Flow(off,cig,'window squeeze=y n1=1 min1=%(os)g' % par)
        if(TYP=='h'):
            Result(off,igrey(olabelx + OSCR))
        else:
            Result(off,igrey(olabelt + OSCR))

        # angle CIGs
        ang = case + 'ang' + '-'+TYP
        ssk = case + 'ssk' + '-'+TYP
        ovl = case + 'ovl'

        if(TYP=='h'):
            Flow  (ssk,off,hcig2ssk())
            Result(ssk,igrey(slabelx+OSCR))

            if(case=='pp'):
                Flow(ang,[ssk,'vpvp','dipa'],hssk2ang())
            if(case=='ps'):
                Flow(ang,[ssk,'vpvs','dipa'],hssk2ang())

            Plot(  ang,igrey(alabel+OSCR))
            Result(ang,[ang,ovl],'Overlay')

        if(TYP=='t'):
            Flow  (ssk,off,tcig2ssk())
            Result(ssk,igrey(slabelt+OSCR))

            if(case=='pp'):
                Flow(ang,[ssk,'vp','vpvp','dipa'],tssk2ang())
            if(case=='ps'):
                Flow(ang,[ssk,'vp','vpvs','dipa'],tssk2ang())

            Plot(  ang, igrey(alabel+OSCR))
            Result(ang,[ang,ovl],'Overlay')

# ------------------------------------------------------------
## 
 # NEW RULES
 ##
# ------------------------------------------------------------
Flow(  'psang-x','psoff-h dipa vpvs',
       adcig.cig2ssk(300,-0.5,0.01) + '|' +
       adcig.xsk2ang(200,-10,0.5))
Plot(  'psang-x', igrey(alabel+OSCR))
Result('psang-x',['psang-x','psovl'],'Overlay')
# ------------------------------------------------------------
Flow(  'psang-u','psoff-t dipa vpvs vp',
       adcig.cig2ssk(600,1000,10) + '|' +
       adcig.tsk2ang(200,-10,0.5))
Plot(  'psang-u', igrey(alabel+OSCR))
Result('psang-u',['psang-u','psovl'],'Overlay')
# ------------------------------------------------------------

Flow('ssk',None,
     '''
     spike nsp=5 mag=1,1,1,1,1
     n1=%(nz)d d1=%(dz)g o1=%(oz)g
     n2=400    d2=10     o2=1000
     k1=60,113,175,260,406
     k2=28,35,55,98,190 | 
     bandpass fhi=0.03 | smooth rect2=7 repeat=1
     ''' % par)
Plot('cor',['ssk','vp','vpvs','dipa'],
     '''
     pp2pstsic a0=%(oa)g na=%(na)d da=%(da)g
     velocity=${SOURCES[1]}
     vpvs=${SOURCES[2]}
     dip=${SOURCES[3]}
     ''' % par + '|' + igrey(alabel+OSCR))

Result('ssk',igrey(slabelt+OSCR))
Result('cor',['cor','psovl'],'Overlay')
# ------------------------------------------------------------

End()