from rsf.proj import *
import math
from rsf.recipes.beg import server as private
from rsf.prog import RSFROOT

segy = 'mcelroy.sgy'

Fetch(segy,'chevron',private)

Flow('mac tmac mac.asc mac.bin', segy,
     '''
     segyread tfile=${TARGETS[1]} hfile=${TARGETS[2]} bfile=${TARGETS[3]} 
     ep=228 
     ''')

# From mac.asc:
#LINE# (IL) BYTES 9-12 (fldr)
#TRACE# (XL) BYTES 13-16 (tracf)
#X COOR. BYTES 73-76  (sx)  
#Y COOR. BYTES 77-80  (sy)
#INLINE SPACING:   55 FEET
#CROSSLINE SPACING:   55 FEET
#CMP DATUM CORRECTION - BYTES 225-228 
#3D AZIMUTH - BYTES 229-232

Flow('fhead','tmac','dd type=float')


# Window traces for inline and crossline 
for case in ('fldr','tracf'):
    Flow(case,'fhead','headermath output=%s | mask min=125 max=175' % case)
Flow('mask','fldr tracf','add mode=p ${SOURCES[1]}',local=1)

Flow('win','mac mask','headerwindow mask=${SOURCES[1]}',local=1)
Flow('twin','fhead mask','headerwindow mask=${SOURCES[1]}',local=1)


# Find offsets
Flow('hx','twin','headermath output="(gx-sx)/1000" | put d1=1 o1=0 d2=1 o2=0',local=1)
Flow('hy','twin','headermath output="(gy-sy)/1000" | put d1=1 o1=0 d2=1 o2=0',local=1)


# Unrotated coordinates
Flow('hxy','hx hy','cat axis=1 ${SOURCES[1]}',local=1)
Result('hxy','dd type=complex | graph symbol=x wanttitle=n plotcol=6',local=1)


# Rotate coordinates
az = 13.5
az = az*math.pi/180
cos = math.cos(az)
sin = math.sin(az)
Flow('rx','hx hy','math x=${SOURCES[0]} y=${SOURCES[1]} output="x*%g+y*%g" ' % (cos,sin),local=1)
Flow('ry','hx hy','math x=${SOURCES[0]} y=${SOURCES[1]} output="-1*%g*x+y*%g" ' % (sin,cos),local=1)
Flow('rxy','rx ry','cat axis=1 ${SOURCES[1]}',local=1)
Result('rxy','dd type=complex | graph symbol=x wanttitle=n plotcol=6',local=1)


# Supergather in x and y offsets
Flow('win2','win','transp memsize=1000')
Flow('win1','win2','window min2=0.8 max2=1.6')
Flow('bin','win1 rxy',
     'bin head=${SOURCES[1]} interp=2 xkey=0 nx=321 dx=0.025 x0=-4 ykey=1 ny=321 dy=0.025 y0=-4')



Nt=800
dt=0.002
t0=0

Nx=297
dx=0.025
x0=-3.7

Ny=Nx
dy=dx
y0=x0


Flow('super','bin',
     '''
     transp plane=23 memsize=500 | transp plane=12 memsize=500 |
     fft1 | fft3 axis=2 | fft3 axis=3 |
     dipfilter v1=-16.5 v2=-7.5 v3=7.5 v4=16.5 taper=2 pass=n dim=3 |
     fft3 axis=3 inv=y | fft3 axis=2 inv=y | fft1 inv=y |
     bandpass fhi=45 flo=8 | sfwindow min3=-3.7 max3=3.7 min2=-3.7 max2=3.7 |
     put label2="Inline Offset" unit2=km label3="Crossline Offset" unit3=km |
     window f1=0 n1=400
     ''')
Result('super','byte gainpanel=all | grey3 pclip=70 title=" " frame1=%d frame2=%d frame3=%d flat=n point2=0.7 label2="Inline Offset" unit2="km" label3="Crossline Offset" unit3="km"' % (89,Nx/2,Ny/2))


#v1=-8.0 v2=-3.5 v3=3.5 v4=8.0

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

Flow('data','super','pad beg1=400')
Result('data','byte gainpanel=all | grey3 pclip=70 title=" " frame1=%d frame2=%d frame3=%d' % (489,Nx/2,Nx/2))

# compute data^2
Flow('data_sq','data','math output="input^2" ')
#Result('data_sq','byte gainpanel=all | grey3 pclip=70 title="data\^2" frame1=%d frame2=%d frame3=%d' % (489,Nx/2,Ny/2))


# fft
Flow('fftdata','data','rtoc | fft3 axis=1 pad=1 | window n1=%d f1=%d' % (Nt/2,Nt/2))
#Result('fftreal','fftdata','real | grey title="real" ')
Flow('fftdatac','fftdata','window n1=105 f1=15')
#Result('fftrealc','fftdatac','real | grey title="realc" ')

Flow('fftdata_sq','data_sq','rtoc | fft3 axis=1 pad=1 | window n1=%d f1=%d' % (Nt/2,Nt/2))
#Result('fftreal_sq','fftdata_sq','real | grey title="real\^2" ')
Flow('fftdatac_sq','fftdata_sq','window n1=160 f1=0')
#Result('fftrealc_sq','fftdatac_sq','real | grey title="realc\^2" ')


Ntau=400
dtau=dt
tau0=0.8

Np=200
p0=-0.008
dp=(0.008-p0)/Np

Nq=Np
q0=p0
dq=dp

Ns=1
s0=0.0
ds=0.1

# prepare to apply fast butterfly
Flow('bfio.bin',os.path.join(RSFROOT,'include','bfio.bin'),'/bin/cp $SOURCE $TARGET',stdin=0,stdout=-1)

# compute integral of d(t,x,y)
Flow('fmod','fftdatac bfio.bin','radon34 ntau=%d dtau=%g tau0=%g np=%d dp=%g p0=%g nq=%d dq=%g q0=%g ns=%d ds=%g s0=%g fi=2 EL=0 N=16 EPSx1=7 EPSx2=5 EPSx3=5 EPSk1=7 EPSk2=5 EPSk3=5 | real | math output=input*%g' % (Ntau,dtau,tau0,Np,dp,p0,Nq,dq,q0,Ns,ds,s0,2*dx*dy/Nt))
#Result('fmod','byte gainpanel=all | grey3 title="fmod" label1=Time unit1=s label2=Wcos unit2="s\^2\_/km\^2\_" label3=Wsin unit3="s\^2\_/km\^2\_" frame1=%d frame2=%d frame3=%d' % (Ntau/2,Np/2,Nq/2))

# compute integral of d(t,x,y)^2
Flow('fmod_sq','fftdatac_sq bfio.bin','radon34 ntau=%d dtau=%g tau0=%g np=%d dp=%g p0=%g nq=%d dq=%g q0=%g ns=%d ds=%g s0=%g fi=2 EL=0 N=16 EPSx1=9 EPSx2=5 EPSx3=5 EPSk1=9 EPSk2=5 EPSk3=5 | real | math output=input*%g' % (Ntau,dtau,tau0,Np,dp,p0,Nq,dq,q0,Ns,ds,s0,2*dx*dy/Nt))
#Result('fmod_sq','byte gainpanel=all | grey3 title="fmod\^2\_" label1=Time unit1=s label2=Wcos unit2="s\^2\_/km\^2\_" label3=Wsin unit3="s\^2\_/km\^2\_" frame1=%d frame2=%d frame3=%d' % (Ntau/2,Np/2,Nq/2))

# compute semblance
Flow('semb','fmod fmod_sq','math output="input^2" | divn den=${SOURCES[1]} rect1=3 rect2=3 rect3=3 | math output="input/%g" ' %(Nx*dx*Ny*dy))
Result('semb','byte gainpanel=all allpos=y | grey3 title=" " color=j label1=Time unit1=s label2="Wcos" unit2="s\^2\_/km\^2\_" label3="Wsin" unit3="s\^2\_/km\^2\_" frame1=%d frame2=%d frame3=%d' % (85,100,100))

# Slow Radon compute integral of d(t,x,y)
#Flow('smod','super','diradon34 ntau=%d dtau=%g tau0=%g np=%d dp=%g p0=%g nq=%d dq=%g q0=%g ns=%d ds=%g s0=%g fi=2 | math output=input*%g' % (Ntau,dtau,tau0,Np,dp,p0,Nq,dq,q0,Ns,ds,s0,dx*dy))
#Result('smod','byte gainpanel=all | grey3 title="NMOsmod" label1=Time unit1=s label2=Wcos unit2="s\^2\_/km\^2\_" label3=Wsin unit3="s\^2\_/km\^2\_" frame1=%d frame2=%d frame3=%d' % (Ntau/2,Np/2,Nq/2))


# (80,107,100) 0.960 0.00056  0.0006
# (81,107,100) 0.962 0.00056  0.0006
# (82,107,100) 0.964 0.00056  0.0006
# (83,107,100) 0.966 0.00056  0.0006
# (84,108,100) 0.968 0.00064  0.0006
# (85,110,100) 0.970 0.00080  0.0008
# (86,112,100) 0.972 0.00096  0.0010
# (87,114,100) 0.974 0.00112  0.0011
# (88,114,100) 0.976 0.00112  0.0011
# (89,114,100) 0.978 0.00112  0.0011
# (90,114,100) 0.980 0.00112  0.0011


# pick Vcos-2 and Vsin-2
Flow('NMOpik','semb','scale axis=3 | pick3 vel0=0.0 smooth=y rect1=10')
Flow('pik1','NMOpik','window f4=0 n4=1')
Flow('pik2','NMOpik','window f4=1 n4=1')
Result('pik1','graph transp=y yreverse=y min2=-0.008 max2=0.008 plotfat=5 plotcol=3 pad=n title="Wcos" label1=Time unit1=s wantaxis=y wanttitle=y')
Result('pik2','graph transp=y yreverse=y min2=-0.008 max2=0.008 plotfat=5 plotcol=3 pad=n title="Wsin" label1=Time unit1=s wantaxis=y wanttitle=y')


# RMO using picked vel
Flow('Vavg0',None,'math n1=%d d1=%g o1=%g output="0.0" ' % (Ntau,dtau,tau0))

Flow('Vcos1',None,'spike n1=%d d1=%g o1=%g k1=81 l1=85 mag=0.0006' % (Ntau,dtau,tau0))
Flow('Vcos2',None,'spike n1=%d d1=%g o1=%g k1=86 l1=86 mag=0.0008' % (Ntau,dtau,tau0))
Flow('Vcos3',None,'spike n1=%d d1=%g o1=%g k1=87 l1=87 mag=0.0010' % (Ntau,dtau,tau0))
Flow('Vcos4',None,'spike n1=%d d1=%g o1=%g k1=88 l1=91 mag=0.0011' % (Ntau,dtau,tau0))
Flow('Vcos0','Vcos1 Vcos2 Vcos3 Vcos4',
     '''
     math Vcos2=${SOURCES[1]} Vcos3=${SOURCES[2]} Vcos4=${SOURCES[3]}
     output="input+Vcos2+Vcos3+Vcos4"
     ''')

#Flow('Vcos0',None,'spike n1=%d d1=%g o1=%g k1=81 l1=91 mag=0.0008' % (Ntau,dtau,tau0))
Result('Vcos0','graph transp=y yreverse=y min2=-0.008 max2=0.008 plotfat=5 plotcol=3 pad=n title="Wcos" label1=Time unit1=s wantaxis=y wanttitle=y')
Flow('RMOpikvel','Vavg0 Vcos0 Vavg0','cat axis=2 ${SOURCES[1:3]}')


#Flow('RMOpikvel','Vavg0 pik1 pik2','cat axis=2 ${SOURCES[1:3]}')
Flow('RMOdata','super RMOpikvel','nmo3_ort velocity=${SOURCES[1]} half=n extend=0 mute=0')
Result('RMOdata','byte gainpanel=all | grey3 pclip=70 title="RMOdata" frame1=%d frame2=%d frame3=%d' % (89,Nx/2,Ny/2))

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

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

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

# Supergather in off and az 
Flow('super1','super','put n2=88209 n3=1 | window squeeze=y | transp memsize=500')
Flow('RMOdata1','RMOdata','put n2=88209 n3=1 | window squeeze=y | transp memsize=500')
Flow('offset1','offset','put n2=88209 n1=1')
Flow('azimuth1','azimuth','put n2=88209 n1=1')
Flow('offaz','offset1 azimuth1','cat axis=1 ${SOURCES[1]}',local=1)

for cased in ('super1','RMOdata1'):
    Flow(cased+'_bin3',[cased,'offaz'],
         'bin head=${SOURCES[1]} interp=2 xkey=0 ykey=1 x0=0 dx=0.2 nx=25 y0=-180 dy=8 ny=45 interp=2',local=1)
    Flow(cased+'_oa',cased+'_bin3','transp plane=23 | transp plane=12',local=1)

Result('super1-aniso','super1_oa',
       '''
       window min1=0.8 max1=1.2 min2=3.6 max2=4.0 | stack |
       wiggle transp=y yreverse=y poly=y pclip=97 color=j parallel2=n
       title="Isotropic NMO" label1=Time unit1=s label2=Azimuth unit2="\^o\_"
       ''')

Result('RMOdata1-aniso','RMOdata1_oa',
       '''
       window min1=0.8 max1=1.2 min2=3.6 max2=4.0 | stack |
       wiggle transp=y yreverse=y poly=y pclip=97 color=j parallel2=n
       title="Residual NMO" label1=Time unit1=s label2=Azimuth unit2="\^o\_"
       ''')

End()