from rsf.proj import *
from math import *
import sys
import string
import rsf.recipes.spmig as spmig
import rsf.recipes.adcig as adcig

# ------------------------------------------------------------
par = {
    'nx':800, 'dx':10,    'ox':0,
    'nz':500, 'dz':5,     'oz':0,
    'nt':1000,'dt':0.004, 'ot':0,  'kt':50,
    'nh':800, 'dh':10,    'oh':-4000,
    'ns':25,   'ds':80,    'os':3000, 'fs':0, 'js':1,
    'nw':200, 'jw':1,
    'velp':3000,
    'vels':3000,
    'nhx':40,
    'nht':200, 'oht':-0.125, 'dht':0.00125,
    'na':281, 'oa':-70, '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=y gridcol=0
    labelsz=6
    %s
    ''' % custom
def igrey3(custom):
    return '''
    grey3 labelrot=n title="" pclip=100 wantaxis=y 
    grid=n gridcol=0
    labelsz=7
    %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=-3 np=601 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 '''
    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 gainpanel=a'

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\_"'

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


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

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

# wavelet
Flow('wav',None,
     '''
     spike nsp=1 mag=1 k1=%(kt)d
     n1=%(nt)d d1=%(dt)g o1=0 | ricker1 frequency=10
     ''' % 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('vp','slop','window n1=1 min1=1000 | math output=%(velp)d' % par)


# vpvs ratio

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=0
     n1=%d o1=%g d1=%g
     k1=0,76,141,221,351 l1=75,140,220,350,500
     ''' % (par['nz'],par['oz'],par['dz']))



###################################################################################


Flow('modl',None,
     '''
     spike n1=%d o1=%g d1=%g n2=3
     nsp=3 k2=1,2,3 mag=1000,1500,2000
     ''' % (par['nx'],par['ox'],par['dx']))
Flow('refl',None,
     '''
     spike n1=%d o1=%g d1=%g n2=3
     nsp=3 k2=1,2,3 mag=1,1,1
     '''%(par['nx'],par['ox'],par['dx']))



####################################################################################




# -------------------------------------Kirmod data
Flow('pp','modl refl',
     '''
     kirmod vel=%g refl=${SOURCES[1]}
     nt=%d dt=%g t0=%g     freq=10
     nh=%d dh=%g h0=%g
     ns=%d ds=%g s0=%g type=c twod=y|
     put label1=t label2=h | pad beg1=%d | put o1=0 
     ''' %
     ( par['velp'],
       par['nt'],par['dt'],par['ot'],
       par['nh'],par['dh'],par['oh'],
       par['ns'],par['ds'],par['os'], par['kt'])
     )
Result('pp','window min2=-3000 max2=3000 | byte  | '
       + igrey3('label1="Time" label2="Offset" pclip=100 label3="Shot position" frame1=300 frame2=200 frame3=10 point2=0.5 point3=0.55 point1=0.5 unit1=s unit2=m unit3=m'))


######################################################## wave data



#######################################################

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


####################################################

par['misc']='itype=x jcx=1 nhx=%(nhx)d nhz=1 hsym=y' % par
img = 'pp' + 'img' + '-'+ 'h'
cig = 'pp' + 'cig' + '-'+ 'h'

Plot('vel','slop',
     '''
     math output="1/input" | window | transp plane=12 |
     grey  label1=z unit1=m label2=x unit2=m   wanttitle=n screenratio=0.3  screenht=5
     ''')

Flow('sx','vel',
     '''
     window | window n2=1 | transp plane=12 | math output="x2" |
     pad end1=1 axis=1 | put o1=0 o2=0 | put o1=0 d1=2 o2=0 d2=%(dx)g | cut max2=3080 | cut min2=5080 | window j2=8 
     | cut max2=600
     '''% par )

Plot('sx',
     '''
     dd type=complex |
     graph wantaxis=n min1=80 max1=8080 max2=0 min2=-2500 symbol=* plotcol=1  plotfat=10 wanttitle=n  screenratio=0.3  screenht=5
     ''')

Result('vel','vel sx','Overlay')

spmig.imagePW3(img,cig,'slop',       sou,rec,par)

Result(img,'window min1=2000 max1=6000 | transp |'+ igrey(ilabel+ISCR))

off = 'pp' + 'off' + '-'+ 'h'
Flow(off,cig,'window squeeze=y min1=2000 max1=6000 | transp plane=13 | transp plane=12' )
Result(off,'transp plane=23| window min2=2200 max2=5800| byte gainpanel=all pclip=100| '
       + igrey3('pclip=100 label1="z" label2="x" label3="h" frame1=300 frame2=180 frame3=40 point2=0.7 point1=0.6 unit1=m unit2=m unit3=m'))

ang = 'pp' + 'ang' + '-'+ 'h'
angtr = 'pp' + 'angtr' + '-'+ 'h'
ssk = 'pp' + 'ssk' + '-'+ 'h'
ovl = 'pp' + 'ovl'

Flow  (ssk,off,hcig2ssk())
Plot(ssk,igrey(slabelx+OSCR),view=1)

Flow(ang,[ssk,'vpvp','dipa'],hssk2ang())
Flow(angtr,ang,'transp plane=23')
Result(  angtr,'window min2=2200 max2=5800|byte gainpanel=all pclip=100| '
         + igrey3('pclip=100 label1="z" label2="x" label3="\F10 q\F3" frame1=290 frame2=180 frame3=110 point2=0.7 point1=0.6 unit1=m unit3="\^o\_" unit2=m'))


Flow('stacks',ang,'noise var=1e-13 | window min3=2200 max3=5800 | stack | spray axis=2 n=281')

Flow('simil',[ang,'stacks'],'noise var=1e-13 | window min3=2200 max3=5800 | similarity other=${SOURCES[1]} rect1=10 rect2=3 niter=50')

Flow('similtr','simil','transp plane=23')

Flow('bar','similtr','bar gainpanel=all pclip=100')
Result('similtr','similtr bar',
       'smooth rect1=10 rect2=5 rect3=5| byte gainpanel=all pclip=100 | '
       + igrey3('scalebar=y maxval=1 bar=${SOURCES[1]} pclip=99.5 color=j bias=-1 label1="z" label2="x" label3="\F10 q\F3" frame1=290 frame2=180 frame3=110 point2=0.7 point1=0.6 unit1=m unit3="\^o\_" unit2=m'))
#------------------------------------------------------------

Flow('thsimil','simil','threshold pclip=30 ')
Flow('scale','thsimil','stack | spray axis=2 n=281')
Flow('simistack',[ang, 'thsimil', 'scale'],
     '''
     noise var=1e-13 | window min3=2200 max3=5800 |
     math x=${SOURCES[1]} y=${SOURCES[2]} output="input*x/(y+0.0000001)" | stack
     ''')

wigg='window j2=5 | wiggle transp=y yreverse=y poly=y  label1="z" unit1=m unit2=m label2="x" title=""'
grey='grey pclip=100 wanttitle=n label1=z unit1=m label2=x unit2=m'

Result('simistack',
       '''
       %s
       ''' % grey)
Flow('stack',ang,'window min3=2200 max3=5800| stack| window ')

Result('stack',
       '''
       %s
       ''' % grey)


End()