from rsf.proj import *
import math

## Plot font and screen ratio, width, height for uniform ppt figs.
p1=0.7
sr=1.0
sw=7.0
sh=10.0
xll=2.0
fat=2

#####
##Synthetic CMP Parameters
nx=101
delx=0.0125*2
ox=-1.25
ny=nx
dely=delx
oy=ox
nt=1001
dt=0.004
fw=10.0
####

################################################################################################################
########  4 Reflection Events ################################################

# Make a layer: z = z0 + x*dx + y*dy

dx=0.4
dy=0.1
dt=.004
z0=0.8
##z0=0.75

cg = 1/math.sqrt(1+dx*dx+dy*dy)
ca = -dx*cg
cb = -dy*cg
d = z0*cg

mx = 0
my = 0

D = d - mx*ca - my*cb

v0 = 2.5
t0 = 2*D/v0
it0=int(t0/dt)
it1=it0
#print it1
Flow('spike1',None,
     '''
     spike n1=%d k1=%d nsp=1 | ricker1 frequency=%g |
     spray axis=2 n=%d o=%g d=%g |
     spray axis=3 n=%d o=%g d=%g
     ''' % (nt,it1,fw,ny,oy,dely,nx,ox,delx))

wx = 1-ca*ca
wy = 1-cb*cb
wxy = - ca*cb

Flow('vel1.asc',None,'echo %g %g %g n1=3 data_format=ascii_float in=$TARGET' % (wx,wy,wxy))
Flow('vel1','vel1.asc','dd form=native | scale dscale=%g | spray axis=1 n=%d' % (1/(v0*v0),nt),local=1)


Flow('cmp1','spike1 vel1','inmo3 velocity=${SOURCES[1]}',local=1)

# Make another layer: z = z0 + x*dx + y*dy

dx=0.4
dy=0.41
dt=.004
z0=1.5

cg = 1/math.sqrt(1+dx*dx+dy*dy)
ca = -dx*cg
cb = -dy*cg
d = z0*cg

mx = 0
my = 0

D = d - mx*ca - my*cb

v0 = 1.7
t0 = 2*D/v0
it0=int(t0/dt)
it2=it0
#print it2
Flow('spike2',None,
     '''
     spike n1=%d k1=%d nsp=1 | ricker1 frequency=%g |
     spray axis=2 n=%d o=%g d=%g |
     spray axis=3 n=%d o=%g d=%g
     ''' % (nt,it2,fw,ny,oy,dely,nx,ox,delx))

wx = 1-ca*ca
wy = 1-cb*cb
wxy = - ca*cb

Flow('vel2.asc',None,'echo %g %g %g n1=3 data_format=ascii_float in=$TARGET' % (wx,wy,wxy))
Flow('vel2','vel2.asc','dd form=native | scale dscale=%g | spray axis=1 n=%d' % (1/(v0*v0),nt),local=1)


Flow('cmp2','spike2 vel2','inmo3 velocity=${SOURCES[1]}',local=1)


# Make yet another layer: z = z0 + x*dx + y*dy

dx=0.2
dy=0.5
dt=.004
z0=2.5

cg = 1/math.sqrt(1+dx*dx+dy*dy)
ca = -dx*cg
cb = -dy*cg
d = z0*cg

mx = 0
my = 0

D = d - mx*ca - my*cb

v0 = 1.75
t0 = 2*D/v0
it0=int(t0/dt)
it3=it0
#print it3
Flow('spike3',None,
     '''
     spike n1=%d k1=%d nsp=1 | ricker1 frequency=%g |
     spray axis=2 n=%d o=%g d=%g |
     spray axis=3 n=%d o=%g d=%g
     ''' % (nt,it3,fw,ny,oy,dely,nx,ox,delx))

wx = 1-ca*ca
wy = 1-cb*cb
wxy = - ca*cb

Flow('vel3.asc',None,'echo %g %g %g n1=3 data_format=ascii_float in=$TARGET' % (wx,wy,wxy))
Flow('vel3','vel3.asc','dd form=native | scale dscale=%g | spray axis=1 n=%d' % (1/(v0*v0),nt),local=1)

Flow('cmp3','spike3 vel3','inmo3 velocity=${SOURCES[1]}',local=1)

# Make a 4th layer: z = z0 + x*dx + y*dy

dx=0.2
dy=0.1
dt=.004
z0=3.5

cg = 1/math.sqrt(1+dx*dx+dy*dy)
ca = -dx*cg
cb = -dy*cg
d = z0*cg

mx = 0
my = 0

D = d - mx*ca - my*cb

v0 = 2.0
t0 = 2*D/v0
it0=int(t0/dt)
it4=it0
#print it4
Flow('spike4',None,
     '''
     spike n1=%d k1=%d nsp=1 | ricker1 frequency=%g |
     spray axis=2 n=%d o=%g d=%g |
     spray axis=3 n=%d o=%g d=%g
     ''' % (nt,it4,fw,ny,oy,dely,nx,ox,delx))

wx = 1-ca*ca
wy = 1-cb*cb
wxy = - ca*cb

Flow('vel4.asc',None,'echo %g %g %g n1=3 data_format=ascii_float in=$TARGET' % (wx,wy,wxy))
Flow('vel4','vel4.asc','dd form=native | scale dscale=%g | spray axis=1 n=%d' % (1/(v0*v0),nt),local=1)

Flow('cmp4','spike4 vel4','inmo3 velocity=${SOURCES[1]}',local=1)

## Add events to create CMP.
Flow('cmp','cmp1 cmp2 cmp3 cmp4',
     '''
     math cmp2=${SOURCES[1]} cmp3=${SOURCES[2]} cmp4=${SOURCES[3]} output="input+cmp2+cmp3+cmp4" |
     noise seed=1 range=0.011
     ''',local=1)

################################################################################################################
##### Repeatedly used slice display

def slice(name,title,bias,barfmt):
    Result(name,
           '''
           grey title="%s" color=j allpos=n screenht=6 barwidth=0.2 screenratio=0.85
           crowd=0.6 scalebar=y colorbar=y xll=2 parallel2=n labelfat=%d titlefat=%d
           label1="Crossline offset" label2="Inline offset"
           format2=%s font=2 bias=%g formatbar=%s barlabelfat=%d
           ''' % (title,fat,fat,'%3.1f',bias,barfmt,fat), local=1)

def smoothslices(name,title):
    global it1
    global it2
    global it3
    global it4
    ishift=2
    it1=it1+ishift
    it2=it2+ishift
    it3=it3+ishift
    it4=it4+ishift
    name1=name+'slice1'
    name2=name+'slice2'
    name3=name+'slice3'
    name4=name+'slice4'
    Flow(name1,name,
         '''
         window n1=1 f1=%d squeeze=y |
         smooth rect1=3 rect2=3 |
         put label1="Crossline offset" unit1=km label2="Inline offset" unit2=km
         ''' % it1)
    Flow(name2,name,
         '''
         window n1=1 f1=%d squeeze=y |
         smooth rect1=3 rect2=3 |
         put label1="Crossline offset" unit1=km label2="Inline offset" unit2=km
         ''' % it2)
    Flow(name3,name,
         '''
         window n1=1 f1=%d squeeze=y |
         smooth rect1=3 rect2=3 |
         put label1="Crossline offset" unit1=km label2="Inline offset" unit2=km
         ''' % it3)
    Flow(name4,name,
         '''
         window n1=1 f1=%d squeeze=y |
         smooth rect1=3 rect2=3 |
         put label1="Crossline offset" unit1=km label2="Inline offset" unit2=km
         ''' % it4)
    title1='Event 1 '+title
    title2='Event 2 '+title
    title3='Event 3 '+title
    title4='Event 4 '+title
    slice(name1,title,0,'%3.1f')
    slice(name2,title,0,'%3.1f')
    slice(name3,title,0,'%3.1f')
    slice(name4,title,0,'%3.1f')
##     Result(name1,'grey title="%s" color=j allpos=n screenht=6 barwidth=0.2 screenratio=0.9 crowd=0.6 scalebar=y colorbar=y xll=2 parallel2=n labelfat=%d titlefat=%d label1="Crossline offset" label2="Inline offset"' % (title,fat,fat), local=1)
##     Result(name2,'grey title="%s" color=j allpos=n screenht=6 barwidth=0.2 screenratio=0.9 crowd=0.6 scalebar=y colorbar=y xll=2 parallel2=n labelfat=%d titlefat=%d label1="Crossline offset" label2="Inline offset"' % (title,fat,fat), local=1)
##     Result(name3,'grey title="%s" color=j allpos=n screenht=6 barwidth=0.2 screenratio=0.9 crowd=0.6 scalebar=y colorbar=y xll=2 parallel2=n labelfat=%d titlefat=%d label1="Crossline offset" label2="Inline offset"' % (title,fat,fat), local=1)
##     Result(name4,'grey title="%s" color=j allpos=n screenht=6 barwidth=0.2 screenratio=0.9 crowd=0.6 scalebar=y colorbar=y xll=2 parallel2=n labelfat=%d titlefat=%d label1="Crossline offset" label2="Inline offset"' % (title,fat,fat), local=1)
    #Result(name,[name1, name2, name3, name4],'TwoRows', local=1)

def slices(name,title):
    global it1
    global it2
    global it3
    global it4
    ishift=2
    it1=it1+ishift
    it2=it2+ishift
    it3=it3+ishift
    it4=it4+ishift
    name1=name+'slice1'
    name2=name+'slice2'
    name3=name+'slice3'
    name4=name+'slice4'
    Flow(name1,name,
         '''
         window n1=1 f1=%d squeeze=y |
         put label1="Crossline offset" unit1=km label2="Inline offset" unit2=km 
         ''' % it1)
    Flow(name2,name,
         '''
         window n1=1 f1=%d squeeze=y |
         put label1="Crossline offset" unit1=km label2="Inline offset" unit2=km
         ''' % it2)
    Flow(name3,name,
         '''
         window n1=1 f1=%d squeeze=y |
         put label1="Crossline offset" unit1=km label2="Inline offset" unit2=km
         ''' % it3)
    Flow(name4,name,
         '''
         window n1=1 f1=%d squeeze=y |
         put label1="Crossline offset" unit1=km label2="Inline offset" unit2=km
         ''' % it4)
    title1='Event 1 '+title
    title2='Event 2 '+title
    title3='Event 3 '+title
    title4='Event 4 '+title
    slice(name1,title,0,'%3.1f')
    slice(name2,title,0,'%3.1f')
    slice(name3,title,0,'%3.1f')
    slice(name4,title,0,'%3.1f')
##     Result(name1,'grey title="%s" color=j allpos=n screenht=6 barwidth=0.2 screenratio=0.9 crowd=0.6 scalebar=y colorbar=y xll=2 parallel2=n labelfat=%d titlefat=%d label1="Crossline offset" label2="Inline offset"' % (title,fat,fat), local=1)
##     Result(name2,'grey title="%s" color=j allpos=n screenht=6 barwidth=0.2 screenratio=0.9 crowd=0.6 scalebar=y colorbar=y xll=2 parallel2=n labelfat=%d titlefat=%d label1="Crossline offset" label2="Inline offset"' % (title,fat,fat), local=1)
##     Result(name3,'grey title="%s" color=j allpos=n screenht=6 barwidth=0.2 screenratio=0.9 crowd=0.6 scalebar=y colorbar=y xll=2 parallel2=n labelfat=%d titlefat=%d label1="Crossline offset" label2="Inline offset"' % (title,fat,fat), local=1)
##     Result(name4,'grey title="%s" color=j allpos=n screenht=6 barwidth=0.2 screenratio=0.9 crowd=0.6 scalebar=y colorbar=y xll=2 parallel2=n labelfat=%d titlefat=%d label1="Crossline offset" label2="Inline offset"' % (title,fat,fat), local=1)
    #Result(name,[name1, name2, name3, name4],'TwoRows', local=1)

def slowslices(name,title):
    global it1
    global it2
    global it3
    global it4
    ishift=2
    it1=it1+ishift
    it2=it2+ishift
    it3=it3+ishift
    it4=it4+ishift
    name1=name+'slice1'
    name2=name+'slice2'
    name3=name+'slice3'
    name4=name+'slice4'
    Flow(name1,name,
         '''
         window n1=1 f1=%d squeeze=y |
         put label1="Crossline offset" unit1=km label2="Inline offset" unit2=km 
         ''' % it1)
    Flow(name2,name,
         '''
         window n1=1 f1=%d squeeze=y |
         put label1="Crossline offset" unit1=km label2="Inline offset" unit2=km
         ''' % it2)
    Flow(name3,name,
         '''
         window n1=1 f1=%d squeeze=y |
         put label1="Crossline offset" unit1=km label2="Inline offset" unit2=km
         ''' % it3)
    Flow(name4,name,
         '''
         window n1=1 f1=%d squeeze=y |
         put label1="Crossline offset" unit1=km label2="Inline offset" unit2=km
         ''' % it4)
    title1='Event 1 '+title
    title2='Event 2 '+title
    title3='Event 3 '+title
    title4='Event 4 '+title
    slice(name1,title,0.10,'%3.3f')
    slice(name2,title,0.20,'%3.3f')
    slice(name3,title,0.18,'%3.3f')
    slice(name4,title,0.18,'%3.3f')
##     Result(name1,'grey title="%s" color=j allpos=n screenht=6 barwidth=0.2 screenratio=0.9 crowd=0.6 scalebar=y bias=0.10 xll=2 parallel2=n labelfat=%d titlefat=%d label1="Crossline offset" label2="Inline offset"' % (title,fat,fat), local=1)
##     Result(name2,'grey title="%s" color=j allpos=n screenht=6 barwidth=0.2 screenratio=0.9 crowd=0.6 scalebar=y bias=0.20 xll=2 parallel2=n labelfat=%d titlefat=%d label1="Crossline offset" label2="Inline offset"' % (title,fat,fat), local=1)
##     Result(name3,'grey title="%s" color=j allpos=n screenht=6 barwidth=0.2 screenratio=0.9 crowd=0.6 scalebar=y bias=0.18 xll=2 parallel2=n labelfat=%d titlefat=%d label1="Crossline offset" label2="Inline offset"' % (title,fat,fat), local=1)
##     Result(name4,'grey title="%s" color=j allpos=n screenht=6 barwidth=0.2 screenratio=0.9 crowd=0.6 scalebar=y bias=0.18 xll=2 parallel2=n labelfat=%d titlefat=%d label1="Crossline offset" label2="Inline offset"' % (title,fat,fat), local=1)
    #Result(name,[name1, name2, name3, name4],'TwoRows', local=1)        

################################################################################################################
##### Measure Local Slopes of CMP

patches=2
patch3=patches*patches*patches

### Break CMP into patches for parellel computation
Flow('pat','cmp',
     '''
     patch p=%d,%d,%d verb=y w=600,60,60 
     ''' % (patches,patches,patches))

## Sort the 3-D patches to 1-D list along 4-th axis
Flow('pat1','pat',
     '''
     put n4=%d n5=1 n6=1 
     ''' % patch3 )

## Measure raw slope fields.  Used for attribute analysis, not NMO.
tsmooth=1
xsmooth=1
ysmooth=1

##  Raw x-slope
Flow('patdipx','pat1',
     '''
     window squeeze=n |
     dip rect1=%d rect2=%d rect3=%d order=3 n4=1 
       ''' % (tsmooth,xsmooth,xsmooth),split=[4,patch3],reduce='cat axis=4')
##  Raw y-slope
Flow('patdipy','pat1',
     '''
     window squeeze=y |
     dip rect1=%d rect2=%d rect3=%d order=3 n4=0
     ''' % (tsmooth,ysmooth,ysmooth),split=[4,patch3],reduce='cat axis=4')

##  Rearrange patch lists back to proper dimensions
Flow('dipxtemp','patdipx',
     '''
     put n4=%d n5=%d n6=%d
     ''' % (patches,patches,patches),local=1)
Flow('dipytemp','patdipy',
     '''
     put n4=%d n5=%d n6=%d
     ''' % (patches,patches,patches),local=1)

##  Merge patches back into single volume.
Flow('px','dipxtemp','patch inv=y weight=y dim=3 n0=%d,%d,%d' % (nt,ny,nx),local=1)
Flow('py','dipytemp','patch inv=y weight=y dim=3 n0=%d,%d,%d' % (nt,ny,nx),local=1)
slices('px','Px')
slices('py','Py')

## Measure smoothed slope field.  Used for PNMO.
tsmooth=10
xsmooth=10
ysmooth=10

##  Smooth x-slope
Flow('patdipxsmooth','pat1',
     '''
     window squeeze=n |
     dip rect1=%d rect2=%d rect3=%d order=3 n4=1
       ''' % (tsmooth,xsmooth,xsmooth),split=[4,patch3],reduce='cat axis=4')
##  Smooth y-slope
Flow('patdipysmooth','pat1',
     '''
     window squeeze=y |
     dip rect1=%d rect2=%d rect3=%d order=3 n4=0
     ''' % (tsmooth,ysmooth,ysmooth),split=[4,patch3],reduce='cat axis=4')

##  Rearrange patch lists back to proper dimensions
Flow('dipxtempsmooth','patdipxsmooth',
     '''
     put n4=%d n5=%d n6=%d
     ''' % (patches,patches,patches),local=1)
Flow('dipytempsmooth','patdipysmooth',
     '''
     put n4=%d n5=%d n6=%d
     ''' % (patches,patches,patches),local=1)

##  Merge patches back into single volume.
Flow('pxsmooth','dipxtempsmooth','patch inv=y weight=y dim=3 n0=%d,%d,%d' % (nt,ny,nx),local=1)
Flow('pysmooth','dipytempsmooth','patch inv=y weight=y dim=3 n0=%d,%d,%d' % (nt,ny,nx),local=1)


slices('pxsmooth','Px for PNMO')
slices('pysmooth','Py for PNMO')

Result('cmp3d','cmp',
     '''
     byte gainpanel=all |
     grey3 point1=%g title="CMP" label2="Offset (km)" parallel2=n
     frame1=%d frame2=50 frame3=50 screenratio=%g screenwd=%g screenht=%g xll=%g labelfat=%d titlefat=%d
     format2=%s
     ''' % (p1,it2,sr,sw,sh,xll,fat,fat,'%3.1f'),local=1)

Result('pxsmooth3d','pxsmooth',
     '''byte gainpanel=all |
     grey3 point1=%g title="Inline slope" frame1=%d frame2=50 frame3=50 label2="Distance (km)"
     screenratio=%g screenwd=%g screenht=%g color=j parallel2=n xll=%g labelfat=%d titlefat=%d framelabelcol=7
     format2=%s
     ''' % (p1,it2,sr,sw,sh,xll,fat,fat,'%3.1f'),local=1)
Result('pysmooth3d','pysmooth',
     '''byte gainpanel=all |
     grey3 point1=%g title="Crossline slope" frame1=%d frame2=50 frame3=50 label2="Distance (km)"
     screenratio=%g screenwd=%g screenht=%g color=j parallel2=n xll=%g labelfat=%d titlefat=%d framelabelcol=7
     format2=%s
     ''' % (p1,it2,sr,sw,sh,xll,fat,fat,'%3.1f'),local=1)

###########################################################################################################
#### Predictive flattening

# Trace similarity (optional)

Flow('shift1','cmp','window f2=1')
Flow('shift2','cmp','window f3=1')

Flow('last1','cmp','window f2=-1 squeeze=n')
Flow('last2','cmp','window f3=-1 squeeze=n')

Flow('ref1','shift1 last1','cat axis=2 ${SOURCES[1]}')
Flow('ref2','shift2 last2','cat axis=3 ${SOURCES[1]}')

Flow('ref1s','ref1','add mode=p $SOURCE | stack axis=1 norm=n')
Flow('ref2s','ref2','add mode=p $SOURCE | stack axis=1 norm=n')

Flow('corr1','ref1 cmp','add mode=p ${SOURCES[1]} | stack axis=1 norm=n')
Flow('corr2','ref2 cmp','add mode=p ${SOURCES[1]} | stack axis=1 norm=n')

Flow('num','cmp','add mode=p $SOURCE | stack axis=1 norm=n')

Flow('cos1','corr1 num ref1s',
     '''
     math s1=${SOURCES[1]} s2=${SOURCES[2]} output="(s1*s2)/(input*input)"
     ''')
Flow('cos2','corr2 num ref2s',
     '''
     math s1=${SOURCES[1]} s2=${SOURCES[2]} output="(s1*s2)/(input*input)"
     ''')
Flow('cos','cos1 cos2',
     '''
     cat axis=3 ${SOURCES[1]} |
     smooth rect1=3 rect2=3 
     ''')

Flow('time2','cos',
     '''
     mul $SOURCE | stack axis=3 norm=n |
     put o1=0 d1=1 o2=0 d2=1 o3=0 d3=1 |
     eikonal vel=n zshot=50 yshot=50
     ''')
Result('time2',
       '''
       grey color=j scalebar=y allpos=y
       title="Minimum Time" transp=n yreverse=n
       ''')

# Time shifts

Flow('pxysmooth','pysmooth pxsmooth','cat axis=4 ${SOURCES[1]}')

Flow('seed','pxysmooth','window n2=1 n3=1 n4=1 | math output=x1')
Flow('pick3-old','pxysmooth seed cos',
     '''
     pwpaint3 seed=${SOURCES[1]} cost=${SOURCES[2]} ref2=50 ref3=50 |
     clip2 lower=0 upper=4
     ''')
Flow('pick3','pxysmooth seed time2',
     '''
     pwpaint2 seed=${SOURCES[1]} cost=${SOURCES[2]} |
     clip2 lower=0 upper=4
     ''')

Flow('t0.asc',None,'echo %g %g %g %g n1=4 data_format=ascii_float in=$TARGET' % (it1,it2,it3,it4))
Flow('t0','t0.asc','dd form=native | math output="(input-3)*0.004" ')

Plot('cmp',
     '''
     byte gainpanel=all |
     grey3 point1=%g title="CMP" label2="Offset (km)"
     frame1=%d frame2=50 frame3=50 screenratio=%g screenwd=%g screenht=%g
     ''' % (p1,it2,sr,sw,sh),local=1)

Plot('pick3','pick3 t0',
     '''
     contour3 point1=%g wanttitle=n wanraxis=n plotfat=3
     frame1=%d frame2=50 frame3=50 screenratio=%g screenwd=%g screenht=%g
     cfile=${SOURCES[1]}
     ''' % (p1,it2,sr,sw,sh),local=1)
Result('pick3','cmp pick3','Overlay')

# Flattening

Flow('flat3','cmp pick3','iwarp warp=${SOURCES[1]} eps=1')


Result('flat3',
     '''
     byte gainpanel=all |
     grey3 point1=%g title="CMP" label2="Offset (km)"
     frame1=%d frame2=50 frame3=50 screenratio=%g screenwd=%g screenht=%g
     ''' % (p1,it2,sr,sw,sh),local=1)

# T0(T) -> T(T0)

Flow('time','pick3','math output=x1 | iwarp warp=$SOURCE eps=1')

Flow('moveout','time','math output="input*input-x1*x1" ')

Result('moveout',
     '''
     byte clip=4 allpos=y |
     grey3 point1=%g title="CMP" label2="Offset (km)" color=j
     frame1=%d frame2=50 frame3=50 screenratio=%g screenwd=%g screenht=%g
     ''' % (p1,it2,sr,sw,sh),local=1)


################################################################################################################
#####   Prepare Geometry Header Files

## Absolute offset
Flow('offset','spike1','window n1=1 | math output="sqrt(x1*x1+x2*x2)" ',local=1)
Result('offset','grey title=Offset allpos=y color=j scalebar=y',local=1)

## Azimuth
Flow('azimuth','spike1','window n1=1 | math output="%g*(x2&x1)" ' % (180/math.pi),local=1)
Result('azimuth','grey title=Azimuth color=j scalebar=y',local=1)

## Header File
Flow('head','offset azimuth',
     'cat axis=3 ${SOURCES[1]} | transp plane=23 | transp plane=12 | put n3=1 n2=10201',local=1)


###################################################################################################################
#####   Conventional NMO

##  Velocity Scan
Flow('scan','cmp offset',
     'put n2=10201 n3=1 | vscan semblance=y offset=${SOURCES[1]} v0=1 nv=51 dv=0.02',local=1)
Result('scan','grey allpos=y bias=0.01 color=j title="Velocity Scan" unit2=km/s',local=1)

##  Isotropic Velocity Models

Flow('vrms_il',None,'math d1=%g n1=%d o1=0 output="1.83-0.05*x1+0.38*(x1-2)*(x1-2)" ' % (dt,nt),local=1)
Flow('vrms_xl',None,'math d1=%g n1=%d o1=0 output="1.83-0.05*x1+0.38*(x1-2)*(x1-2)" ' % (dt,nt),local=1)

##  Isotropic NMO
Flow('nmo06','cmp offset vrms_il',
     'put n2=10201 n3=1 | nmo offset=${SOURCES[1]} velocity=${SOURCES[2]} | put n2=101 n3=101',local=1)
Flow('nmo07','cmp offset vrms_xl',
     'put n2=10201 n3=1 | nmo offset=${SOURCES[1]} velocity=${SOURCES[2]} | put n2=101 n3=101',local=1)
Result('nmo063d','nmo06',
       '''
       byte gainpanel=all |
       grey3 point1=%g title="Circular NMO" label2="Distance (km)" frame1=%d frame2=50 frame3=50
       screenratio=%g screenwd=%g screenht=%g parallel2=n xll=%g labelfat=%d titlefat=%d
       format2=%s
       ''' % (p1,it2,sr,sw,sh,xll,fat,fat,'%3.1f'),local=1)
Plot('nmo073d','nmo07',
       '''
       byte gainpanel=all |
       grey3 point1=%g title="Circular NMO, crossline" label2="Distance (km)" frame1=%d frame2=50 frame3=50
       screenratio=%g screenwd=%g screenht=%g parallel2=n xll=%g labelfat=%d titlefat=%d
       ''' % (p1,it2,sr,sw,sh,xll,fat,fat),local=1)
#Result('nmo0','cmp3d nmo063d','SideBySideAniso',local=1)

###################################################################################################################
#####   Apply PNMO

Flow('PNMO','cmp pysmooth pxsmooth',
     '''
     pnmo3d dipx=${SOURCES[2]} dipy=${SOURCES[1]}
     ''')
Result('PNMO3d','PNMO',
     '''byte gainpanel=all |
     grey3 point1=%g title="Elliptical NMO" frame1=%d frame2=50 frame3=50 label2="Distance (km)"
     screenratio=%g screenwd=%g screenht=%g parallel2=n xll=%g labelfat=%d titlefat=%d
     format2=%s
     ''' % (p1,it2,sr,sw,sh,xll,fat,fat,'%3.1f'),local=1)
#Result('dips3d','cmp3d PNMO3d pxsmooth3d pysmooth3d','TwoRows')
## Result('PNMO3','pxsmooth pysmooth PNMO','SideBySideAniso',local=1)

###################################################################################################################
#####   Calculate Pxy and/or Pyx before PNMO

## Calculate Pxy
Flow('dpxdy','pxsmooth',
     '''
     math output="input" |
     transp plane=12 |
     deriv |
     transp plane=12
     ''')
Flow('dpxdt','pxsmooth','deriv')
Flow('pxy','pysmooth dpxdy dpxdt',
     '''
     math pxy=${SOURCES[1]} pxt=${SOURCES[2]} output="input*pxt+pxy"
     ''')
slices('pxy','P\_xy\^')

## Calculate Pyx
Flow('dpydx','pysmooth',
     '''
     math output="input" |
     transp plane=13 |
     deriv |
     transp plane=31 
     ''')
Flow('dpydt','pysmooth','deriv')
Flow('pyx','pxsmooth dpydx dpydt',
     '''
     math pyx=${SOURCES[1]} pyt=${SOURCES[2]} output="input*pyt+pyx"
     ''')
slices('pyx','P\_yx\^')

## Take their average or sum
pavg='(pxy+pyx)/1.0'
Flow('pcross','pxy pyx',
     '''
     math pxy=${SOURCES[0]} pyx=${SOURCES[1]} output="%s"
     ''' % (pavg),local=1)
slices('pcross','P\_xy\^+P\_yx\^/1')

## Test
Flow('pxy_test','pxsmooth',
     '''
     math output="input*x1/0.004" |
     transp plane=12 |
     deriv |
     transp plane=12
     ''')
slices('pxy_test','pxytest=W\_xy\^')

Flow('testvol',None,'spike n1=10 n2=20 nsp=2 k1=3,6 k2=3,17 d1=1 d2=1| spray axis=3 n=30 d=1 o=0')
Plot('testvol1','testvol','byte | grey3 frame1=5 frame2=5 frame3=5')
Plot('testvol2','testvol','transp plane=12 |  byte gainpanel=all | grey3 frame1=5 frame2=5 frame3=5')
Flow('testvol3','testvol','transp plane=31 ')
Plot('testvol3','byte gainpanel=all | grey3 frame1=5 frame2=5 frame3=5')
Result('testvol','testvol1 testvol2 testvol3','SideBySideAniso')
###################################################################################################################
#####   Apply PNMO to slope fields

Flow('pxnmod','px pxsmooth pysmooth','pnmo3d dipx=${SOURCES[1]} dipy=${SOURCES[2]}')
Flow('pynmod','py pxsmooth pysmooth','pnmo3d dipx=${SOURCES[1]} dipy=${SOURCES[2]}')
Flow('pxynmod','pxy pxsmooth pysmooth','pnmo3d dipx=${SOURCES[1]} dipy=${SOURCES[2]}')
Flow('pyxnmod','pyx pxsmooth pysmooth','pnmo3d dipx=${SOURCES[1]} dipy=${SOURCES[2]}')
Flow('pcrossnmod','pcross pxsmooth pysmooth','pnmo3d dipx=${SOURCES[1]} dipy=${SOURCES[2]}')


###################################################################################################################
#####   Calculate Alpha

## Formula for Alpha
formula='0.5*atan((2*(x3*x2)*(x1*(pxy*6.4)+(px*0.16)*(py*0.16)))/(((x1*(pxy*6.4))+(px*0.16)*(py*0.16))*(x3*x3-x2*x2)+x1*(x2*(px*0.16)-x3*(py*0.16))+0.0000000001))'

## Compute Alpha Volume
Flow('alpha','pxsmooth pysmooth pxy',
     '''
     math px=${SOURCES[0]} py=${SOURCES[1]} pxy=${SOURCES[2]} output="%s"
     ''' % (formula),local=1)


##  Apply PNMO to Alpha Volume
Flow('alphanmo','alpha pxsmooth pysmooth',
     'pnmo3d dipx=${SOURCES[1]} dipy=${SOURCES[2]}')
Flow('bar','alphanmo','bar')
Plot('alphanmo3d','alphanmo bar',
     '''
     byte |
     grey3 point1=%g title=Azimuth color=j scalebar=y bar=${SOURCES[1]} 
     frame1=%d frame2=90 frame3=90 label2=Distance unit2=km
     screenratio=%g screenwd=%g screenht=%g
     ''' % (p1,it2,sr,sw,sh),local=1)
#Result('Azimuth','cmp3d PNMO3d alphanmo3d','SideBySideAniso')

##  Display Time Slices of Alpha Volume
smoothslices('alphanmo','Azimuth from px,py,pxy')

###################################################################################################################
#####   Bin the data in offset and azimuth

##  Sort/Bin CMP x-y volume into offset-azimuth volume
Flow('transp','cmp','put n3=1 n2=10201 | transp',local=1)
Flow('bin','transp head',
     'bin head=${SOURCES[1]} xkey=0 ykey=1 x0=0 dx=0.025 nx=70 y0=-180 dy=4 ny=91 interp=2',local=1)
Flow('oacmp','bin','transp plane=23 | transp plane=12',local=1)
Result('oacmp','window n2=1 min2=0.75 | grey pclip=95 label2=Azimuth unit2="\^o\_" xll=2 title="Input CMP" screenratio=2 parallel2=n xll=%g labelfat=%d titlefat=%d format2=%s' % (xll,fat,fat,'%3.1f'),local=1)

##  Sort/Bin PNMO x-y volume into offset-azimuth volume
Flow('PNMOtransp','PNMO','put n3=1 n2=10201 | transp',local=1)
Flow('PNMObin','PNMOtransp head',
     'bin head=${SOURCES[1]} xkey=0 ykey=1 x0=0 dx=0.025 nx=70 y0=-180 dy=4 ny=91 interp=2',local=1)
Flow('oaPNMO','PNMObin','transp plane=23 | transp plane=12',local=1)
Result('oaPNMO','window n2=1 min2=0.75 | grey pclip=95 label2=Azimuth unit2="\^o\_" xll=2 title="After PNMO" screenratio=2 parallel2=n xll=%g labelfat=%d titlefat=%d format2=%s' % (xll,fat,fat,'%3.1f'),local=1)

##  Sort/Bin Alpha x-y volume into offset-azimuth volume
Flow('alphatransp','alphanmo','put n3=1 n2=10201 | transp',local=1)
Flow('alphabin','alphatransp head',
     'bin head=${SOURCES[1]} xkey=0 ykey=1 x0=0 dx=0.025 nx=70 y0=-180 dy=4 ny=91 interp=2',local=1)
Flow('oaalpha','alphabin','transp plane=23 | transp plane=12',local=1)
Plot('oaalpha','window n2=1 min2=0.5 | grey label2=Azimuth color=j scalebar=y unit2="\^o\_" ',local=1)
#Result('Aziview','oacmp oaPNMO','SideBySideAniso')

## Smooth Alpha along offset and azimuth coordinates.
Flow('oaalphasmooth','oaalpha','smooth rect1=2 rect2=5 rect3=1 repeat=1 ')
Plot('oaalphasmooth',
     '''
     window n2=1 min2=0.5 |
     grey label2=Azimuth color=j scalebar=y unit2="\^o\_"
     ''',local=1)
#Result('Azismooth','oaPNMO oaalpha oaalphasmooth','SideBySideAniso')

###################################################################################################################
#####   View single ring of angle estimates

## Results from event 1
Flow('alphawiggle1','oaalpha','window n1=1 f1=308 n2=1 f2=26 | put label1=Azimuth unit1=Degrees')
Plot('alphawiggle1','wiggle title=" " plotcol=3 wantaxis=n')
## Results from event 2
Flow('alphawiggle2','oaalpha','window n1=1 f1=450 n2=1 f2=21 | put label1=Azimuth unit1=Degrees')
Plot('alphawiggle2','wiggle title=" " wantaxis=n')
## Results from event 3
Flow('alphawiggle3','oaalpha','window n1=1 f1=853 n2=1 f2=20 | put label1=Azimuth unit1=Degrees')
Plot('alphawiggle3','wiggle title=" " plotcol=2 wantaxis=n')

## Isotropic Guide
Flow('aziaxi',None,'spike mag=1 n1=91 k1=0 l1=90 o1=-180 d1=4 label1=Azimuth unit1=Degrees')
Flow('iso1','aziaxi',
     '''
     math output="0.785398163*(x1+180)" |
     window n1=12 f1=0
     ''')
Flow('iso2','aziaxi',
     '''
     math output="0.785398163*(x1+90)"|
     window n1=22 f1=12
     ''')
Flow('iso3','aziaxi',
     '''
     math output="0.785398163*(x1)"|
     window n1=23 f1=34
     ''')
Flow('iso4','aziaxi',
     '''
     math output="0.785398163*(x1-90)"|
     window n1=22 f1=57
     ''')
Flow('iso5','aziaxi',
     '''
     math output="0.785398163*(x1-180)"|
     window n1=12 f1=79
     ''')
Flow('isowiggle','iso1 iso2 iso3 iso4 iso5',
     '''
     cat axis=1 ${SOURCES[1]} ${SOURCES[2]} ${SOURCES[3]} ${SOURCES[4]}   
     ''')

Plot('isowiggle','wiggle plotcol=5 label2=Alpha unit2=Radians xreverse=y title="Alpha versus Azimuth"')

Result('wiggles','alphawiggle2 isowiggle','Overlay')

###################################################################################################################
#####   Examine slope fields

##  Slope Magnitudes
slopemag='sqrt(px*0.16*px*0.16+py*0.16*py*0.16)'
Flow('pmag','pxsmooth pysmooth',
     '''
     math px=${SOURCES[0]} py=${SOURCES[1]} output="%s" |
     pnmo3d dipx=${SOURCES[0]} dipy=${SOURCES[1]}
     ''' % slopemag, local=1)
smoothslices('pmag','Slope Magnitude')


##  X-Slopes
smoothslices('pxnmod','Slope x')

##  Y-Slopes
smoothslices('pynmod','Slope y')

##  XY-Slopes
smoothslices('pxynmod','Slope xy')

##  YX-Slopes
smoothslices('pyxnmod','Slope yx')

##  Cross-Slopes
smoothslices('pcrossnmod','Slope Cross')


## p-hat dot product with x-hat
Flow('offsetvol','offset','spray axis=1 n=1001 | put d1=0.004 o1=0')

Flow('nmovel','pmag offsetvol',
     '''
     math l=${SOURCES[1]} output="((x1*input)/(l+0.05))" 
     ''')
slowslices('nmovel','NMO Slowness Squared')

dotprod0='acos((px*0.16/(P+0.0000001))*(x3/(X+0.0000001))+(py*0.16/(P+0.0000001))*(x2/(X+0.0000001)))'
dotprod1='((px*0.16)*(x3/(X+0.0000000001))+(py*0.16)*(x2/(X+0.0000000001)))'
dotprod2='(px*0.16)*(x3)+(py*0.16)*(x2)/(X*P+1)'

Flow('pdotx','pxnmod pynmod pmag offsetvol',
     '''
     math px=${SOURCES[0]} py=${SOURCES[1]} P=${SOURCES[2]} X=${SOURCES[3]} output="(%s*%g)"
     ''' % (dotprod2,(1)), local=1)
smoothslices('pdotx','Slope Azimuth Projection')

##  Sort/Bin pdotx x-y volume into offset-azimuth volume
Flow('pdotxtransp','pdotx','put n3=1 n2=10201 | transp',local=1)
Flow('pdotxbin','pdotxtransp head',
     'bin head=${SOURCES[1]} xkey=0 ykey=1 x0=0 dx=0.025 nx=70 y0=-180 dy=1 ny=361 interp=2',local=1)
Flow('oapdotx','pdotxbin','transp plane=23 | transp plane=12',local=1)
Result('oapdotx','window n2=1 min2=0.5 | grey allpos=y label2=Azimuth color=j scalebar=y unit2="\^o\_" ',local=1)


###################################################################################################################
#####   Calculate Wx Wy Wxy

##  Wxy
formWxy='(x1*pxy*6.4)+(0.16*0.16*py*px)'
Flow('Wxy','pxsmooth pysmooth dpydx',
     '''
     math px=${SOURCES[0]} py=${SOURCES[1]} pxy=${SOURCES[2]} output="%s"
     ''' % formWxy)
slices('Wxy','W\_xy\^')
Flow('Wxynmod','Wxy pxsmooth pysmooth','pnmo3d dipx=${SOURCES[1]} dipy=${SOURCES[2]}')
smoothslices('Wxynmod','W\_xy\^nmod')

##  Wx
formWx='(x1*(0.16*px)-Wxy*x2)/(x3+0.00001)'
Flow('Wx','pxsmooth pysmooth Wxy',
     '''
     math px=${SOURCES[0]} py=${SOURCES[1]} Wxy=${SOURCES[2]} output="%s"
     ''' % formWx)
slices('Wx','W\_x\^')
Flow('Wxnmod','Wx pxsmooth pysmooth','pnmo3d dipx=${SOURCES[1]} dipy=${SOURCES[2]}')
smoothslices('Wxnmod','W\_x\^nmod')

##  Wy
formWy='(x1*(0.16*py)-Wxy*x3)/(x2+0.00001)'
Flow('Wy','pxsmooth pysmooth Wxy',
     '''
     math px=${SOURCES[0]} py=${SOURCES[1]} Wxy=${SOURCES[2]} output="%s"
     ''' % formWy)
slices('Wy','W\_y\^')
Flow('Wynmod','Wy pxsmooth pysmooth','pnmo3d dipx=${SOURCES[1]} dipy=${SOURCES[2]}')
smoothslices('Wynmod','W\_y\^nmod')

##  Alpha using Wx Wy Wxy directly
wxyalpha='0.5*atan(2*Wxy/(Wx-Wy+0.0000000001))'
Flow('newalpha','Wx Wy Wxy pxsmooth pysmooth',
     '''
     math Wx=${SOURCES[0]} Wy=${SOURCES[1]} Wxy=${SOURCES[2]} output="%s" |
     pnmo3d dipx=${SOURCES[3]} dipy=${SOURCES[4]}
     '''% wxyalpha)
smoothslices('newalpha','Alpha from W\_x\^,W\_y\^,W\_xy\^')

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

## Result('debug1','pmagslice1 pdotxslice1 Wxynmodslice1 Wxnmodslice1 Wynmodslice1 newalphaslice1','TwoRows')

## Result('debug2','pmagslice2 pdotxslice2 Wxynmodslice2 Wxnmodslice2 Wynmodslice2 newalphaslice2','TwoRows')

## Result('debug3','pmagslice3 pdotxslice3 Wxynmodslice3 Wxnmodslice3 Wynmodslice3 newalphaslice3','TwoRows')
############################

###########################################   W Inversion
## Delta(T^2) Volume
## The transpose at the end prepares the volume for input into moveout.c, where axes={1:x,2:y,3:t}
Flow('deltaT','cmp pxsmooth pysmooth',
     '''
     math output="x1*x1" |
     pnmo3d dipx=${SOURCES[1]} dipy=${SOURCES[2]} |
     math output="input-(x1*x1)" | window n2=51 f2=25 n3=51 f3=25 |
     transp plane=23
     ''')
##  t \_\F10 D\^\F3
smoothslices('deltaT',' ')
Flow('dTtransp','deltaT','transp plane=13')
## Dummy W vector
Flow('W0',None,'spike n1=3')

## Output template
Flow('cmpT','cmp','window n2=51 f2=25 n3=51 f3=25 | transp plane=13 | window n3=1')

nt=1001
wlist=[]
W100list=[]
j=0
for i in range(nt):
    Winv = 'Winv_'+str(i)
    deltaT = 'deltaT_'+str(i)


    ##T0 slice
    Flow(deltaT,'dTtransp','window n3=1 f3=%d' % i)
    ## W extraction
    Flow(Winv,[deltaT,'cmpT','W0'],
         '''
         conjgrad nmow_adj in=${SOURCES[0]} gather=${SOURCES[1]} nw=3 adj=y mod=${SOURCES[2]} niter=10 | put n2=1
         ''')
    wlist.append(Winv)
    if j==100:
        W00='W_'+str(i)
        Flow(W00,wlist,'cat ${SOURCES[1:%d]} axis=2' % 100)
        W100list.append(W00)
        if i==100:
            W100list.append(Winv)
        wlist=[]
        j=0
    j=j+1

Flow('Winv',W100list,'cat ${SOURCES[1:%d]} axis=2 ' % 11)

Plot('Winv',
     '''
     transp |
     put label2=W unit2=" " |
     graph transp=y yreverse=y min2=-0.5 max2=0.5 min1=0 max1=4
     grid1=y grid2=n title="W Inversion Results" 
     ''' )
sr=2.0
sra=2.0
ll=1.5
Flow('Wxinv','Winv','window n1=1 f1=0 squeeze=y')
Result('Wxinv',
     '''
     put label1="Time t\_0\^" unit1=s d1=%g label2=W unit2="s\^2\_/km\^2\_"|
     graph transp=y yreverse=y min2=-0.5 max2=0.5 min1=0 max1=4 symbol=+ symbolsz=4
     grid1=y grid2=n title="W\_x\^ results" screenratio=%g xll=%g parallel2=n labelfat=%d titlefat=%d
     format2=%s
     ''' % (dt,sr,ll,fat,fat,'%3.1f'))
Flow('Wyinv','Winv','window n1=1 f1=1')
Result('Wyinv',
     '''
     put label1="Time t\_0\^" unit1=s d1=%g label2=W unit2="s\^2\_/km\^2\_"|
     graph transp=y yreverse=y min2=-0.5 max2=0.5 min1=0 max1=4 symbol=- symbolsz=4
     grid1=y grid2=n title="W\_y\^ results" screenratio=%g xll=%g parallel2=n labelfat=%d titlefat=%d
     format2=%s
     ''' % (dt,sr,ll,fat,fat,'%3.1f'))
Flow('Wxyinv','Winv','window n1=1 f1=2')
Result('Wxyinv',
     '''
     put label1="Time t\_0\^" unit1=s d1=%g label2=W unit2="s\^2\_/km\^2\_"|
     graph transp=y yreverse=y min2=-0.1 max2=0.1 min1=0 max1=4 symbol=o symbolsz=4
     grid1=y grid2=n title="W\_xy\^ results" screenratio=%g  xll=%g parallel2=n labelfat=%d titlefat=%d
     format2=%s
     ''' % (dt,sr,ll,fat,fat,'%3.1f'))

Flow('InvAlpha','Wxinv Wyinv Wxyinv',
     '''
     math wx=${SOURCES[0]} wy=${SOURCES[1]} wxy=${SOURCES[2]} output="%g*0.5*atan(2*wxy/(wx-wy+0.000000001))"
     ''' % (180/math.pi))
Result('InvAlpha',
     '''
     put label1="Time t\_0\^" unit1=s d1=%g label2="Alpha" unit2="\^o\_"|
     graph transp=y yreverse=y xreverse=n min1=0 max1=4 min2=-50 max2=50 symbol=* symbolsz=8
     grid1=y grid2=n title="Alpha results" wantaxis=y screenratio=%g xll=%g parallel2=n labelfat=%d titlefat=%d
     format2=%s
     ''' % (dt,sra,ll,fat,fat,'%3.1f'))
#Result('Woverlay','Wxinv Wyinv Wxyinv','Overlay')

#Result('AziW','oacmp oaPNMO Winv','SideBySideAniso')

#Result('3dW','cmp3d PNMO3d Winv','SideBySideAniso')


######### EigenValues and Vnmo #########

Flow('azimuthvol','azimuth','spray axis=1 n=1001 | put d1=%g o1=0 | math output="input/%g"' % (dt,(180/math.pi)))

Flow('beta','InvAlpha',
     '''
     window squeeze=y |
     spray axis=2 n=%d | spray axis=3 n=%d |
     put d2=%g o2=%g d3=%g o3=%g d1=%g |
     math output="input/%g"
     ''' % (ny,nx,dely,oy,delx,ox,dt,(180/math.pi)))

Flow('lambda1','Wxinv Wyinv Wxyinv',
     '''
     math wx=${SOURCES[0]} wy=${SOURCES[1]} wxy=${SOURCES[2]} output="0.5*(wx+wy+sqrt((wx-wy)*(wx-wy)+4*wxy*wxy))" |
     window squeeze=y |
     spray axis=2 n=%d | spray axis=3 n=%d |
     put d2=%g o2=%g d3=%g o3=%g d1=%g
     ''' % (ny,nx,dely,oy,delx,ox,dt))

Flow('lambda2','Wxinv Wyinv Wxyinv',
     '''
     math wx=${SOURCES[0]} wy=${SOURCES[1]} wxy=${SOURCES[2]} output="0.5*(wx+wy-sqrt((wx-wy)*(wx-wy)+4*wxy*wxy))" |
     window squeeze=y |
     spray axis=2 n=%d | spray axis=3 n=%d |
     put d2=%g o2=%g d3=%g o3=%g d1=%g
     ''' % (ny,nx,dely,oy,delx,ox,dt))

Flow('slowsqrd','lambda1 lambda2 azimuthvol beta',
     '''
     math l1=${SOURCES[0]} l2=${SOURCES[1]} a=${SOURCES[2]} b=${SOURCES[3]}
     output="((l1*cos(a-b)*cos(a-b)+l2*sin(a-b)*sin(a-b)))"
     ''')

slowslices('slowsqrd','NMO slowness squared')

End()