from rsf.proj import *
import math
import rsf.recipes.beg as beg

##################################################
## Velocity Model


# V_fast
v_background=4.0


# Angle between V_fast and x1 axes
angle_deg=112


# Percent anisotropy (V_fast/V_slow-1)*100
vfvs_perc_aniso=3


vf=v_background
vs=vf-((vfvs_perc_aniso/100.0)*vf)
wf=1.0/(vf*vf)
ws=1.0/(vs*vs)

deg2rad=math.pi/180.0
angle_rad=(angle_deg)*deg2rad
c=math.cos(angle_rad)
s=math.sin(angle_rad)


w1m=wf*c*c+ws*s*s
w12m=(wf-ws)*(s*c)
w2m=wf*s*s+ws*c*c
print w1m
print w12m
print w2m

## Or just set the W's manually:
w1m=0.0659
w12m=0.0014
w2m=0.0631
print w1m
print w12m
print w2m

## Flow('w1',None,'spike n1=376 n2=101 n3=101 mag=%g d2=0.1 d3=0.1' % (w1m))
## Flow('w2','w1','math output=%g'% (w2m))
## Flow('w12','w1','math output=%g'% (w12m))

##################################################
## Data

diffs='diffs.HH'

Fetch(diffs,'rpsea',beg.server)
Flow('diffs',diffs,'dd form=native')

Plot('diffs',
       '''
       window max1=2.9 max2=14.9 j3=2 j2=2 |
       byte gainpanel=all pclip=99 |
       grey3 frame1=200 frame2=125 frame3=125
       title="a) Fault Diffractions" 
       label2="x\_1\^" label3="x\_2\^"
       color=e
       ''')
Plot('diff_slice','diffs',
     '''
     window n1=1 f1=200 |
     transp |
     grey title="a) Fault Diffractions" 
     label2="x\_1\^" label3="x\_2\^"
     color=e
     ''')
##################################################
## Velocity Continuation--Migration

## t^2 Time stretch + CosFT in space + FFT in time
Flow('fft','diffs','t2stretch | pad end1=1024 | fft1 | fft3 axis=3 | fft3 axis=2')
Plot('fft',
       '''
       real |
       byte gainpanel=all |
       grey3 frame1=50 frame2=108 frame3=108
       pclip=99 grey title=FFT label2=k_x label3=k_y
       ''')


## W parameters for velocity continuation
w1=w1m
w12=w12m
w2=w2m


eps=0.0000000000001
## Velocity Continuation Phase Shift

shiftmig='input*exp(I*%g*(x3*x3*%g-2*x2*x3*%g+x2*x2*%g)/(16.0*((x1/2.0)+%g)*(%g*%g-%g*%g)))'% (math.pi,w1,w12,w2,eps,w12,w12,w1,w2)

Flow('vcfft','fft','math output="%s"' % (shiftmig))
Plot('vcfft',
       '''
       real |
       byte gainpanel=all |
       grey3 frame1=50 frame2=108 frame3=108
       pclip=99 grey title=VCFFT label2=k_x label3=k_y
       ''')
#Result('fft','fft vcfft','SideBySideAniso')
Flow('vc_mig','vcfft','fft3 axis=2 inv=y | fft3 axis=3 inv=y | fft1 inv=y | window n1=376 | t2stretch inv=y')
Plot('vc_mig',
       '''
       window max1=2.9 max2=14.9 j2=2 j3=2 |
       byte gainpanel=all pclip=99 |
       grey3 frame1=200 frame2=125 frame3=125
       title="b) Anisotropic Velocity Continuation" 
       label2="x\_1\^" label3="x\_2\^"
       color=e
       ''')
Plot('mig_slice','vc_mig',
     '''
     window n1=1 f1=200 |
     transp |
     grey title="b) Anisotropic Velocity Continuation"
     label2="x\_1\^" label1="x\_2\^"
     color=e
     ''')
#Result('images-mig','diffs vc_mig','SideBySideIso')


##################################################
## Velocity Continuation on Noisy Diffraction Data

# noise==100 --> SNR~0.8
noise=100
Flow('noise','diffs','math output="0.0" | noise mean=0 range=%g seed=314'%(noise))
Flow('noisy_diffs','diffs','noise mean=0 range=%g seed=314'%(noise))
Plot('noisy_diffs',
       '''
       window max1=2.9 max2=14.9 j2=2 j3=2 |
       byte gainpanel=all pclip=99 |
       grey3 frame1=200 frame2=125 frame3=125
       title="b) Noisy Fault Diffractions" 
       label2="x\_1\^" label3="x\_2\^"
       color=e flat=n 
       ''')
Plot('diff_slice_noise','noisy_diffs',
     '''
     window n1=1 f1=200 |
     transp |
     grey title="b) Fault Diffractions" 
     label2="x\_1\^" label1="x\_2\^"
     color=e
     wheretitle=top yll=1.5 wherexlabel=bottom
     ''')

Result('noisy_diffs','diffs noisy_diffs','SideBySideAniso')
## t^2 Time stretch + CosFT in space + FFT in time
Flow('noisy_fft','noisy_diffs','t2stretch | pad end1=1024 | fft1 | fft3 axis=3 | fft3 axis=2')
Plot('noisy_fft',
       '''
       real |
       byte gainpanel=all |
       grey3 frame1=50 frame2=108 frame3=108
       pclip=99 grey title=FFT label2=k_x label3=k_y
       ''')


## W parameters for velocity continuation
w1=w1m
w12=w12m
w2=w2m


eps=0.0000000000001
## Velocity Continuation Phase Shift
shiftmig='input*exp(I*%g*(x3*x3*%g-2*x2*x3*%g+x2*x2*%g)/(16.0*((x1/2.0)+%g)*(%g*%g-%g*%g)))'% (math.pi,w1,w12,w2,eps,w12,w12,w1,w2)

Flow('noisy_vcfft','noisy_fft','math output="%s"' % (shiftmig))
Plot('noisy_vcfft',
       '''
       real |
       byte gainpanel=all |
       grey3 frame1=50 frame2=108 frame3=108
       pclip=99 grey title=VCFFT label2=k_x label3=k_y
       ''')
#Result('fft','fft vcfft','SideBySideAniso')
Flow('noisy_vc_mig','noisy_vcfft','fft3 axis=2 inv=y | fft3 axis=3 inv=y | fft1 inv=y | window n1=376 | t2stretch inv=y')
Plot('vc_noise','noisy_vc_mig',
       '''
       window max1=2.9 max2=14.9 |
       byte gainpanel=all pclip=99 |
       grey3 frame1=200 frame2=250 frame3=250
       title="d) Noise Added Before Anisotropic Continuation" 
       label2="x\_1\^" label3="x\_2\^"
       color=e
       ''')
Plot('noise_slice','noisy_vc_mig',
     '''
     window n1=1 f1=200 |
     transp |
     grey title="d) Anisotropic Velocity Continuation" 
     label2="x\_1\^" label1="x\_2\^"
     color=e
     wheretitle=top yll=1.5 wherexlabel=bottom
     ''')

## Isotropic Comparison

w1=(w1m+w2m)/2.0
w2=w1
w12=0

shiftmig='input*exp(I*%g*(x3*x3*%g-2*x2*x3*%g+x2*x2*%g)/(16.0*((x1/2.0)+%g)*(%g*%g-%g*%g)))'% (math.pi,w1,w12,w2,eps,w12,w12,w1,w2)

Flow('vcfftiso','noisy_fft','math output="%s"' % (shiftmig))
Plot('vcfftiso',
       '''
       real |
       byte gainpanel=all |
       grey3 frame1=50 frame2=108 frame3=108
       pclip=99 grey title=VCFFT label2=k_x label3=k_y
       ''')
#Result('fft','fft vcfft','SideBySideAniso')
Flow('vc_iso','vcfftiso','fft3 axis=2 inv=y | fft3 axis=3 inv=y | fft1 inv=y | window n1=376 | t2stretch inv=y')
Plot('vc_iso',
       '''
       window max1=2.9 max2=14.9 j2=2 j3=2 |
       byte gainpanel=all pclip=99 |
       grey3 frame1=200 frame2=125 frame3=125
       title="c) Isotropic Velocity Continuation" 
       label2="x\_1\^" label3="x\_2\^"
       color=e
       ''')
Plot('iso_slice','vc_iso',
     '''
     window n1=1 f1=200 |
     transp |
     grey title="c) Isotropic Velocity Continuation" 
     label2="x\_1\^" label1="x\_2\^"
     color=e
     wheretitle=top yll=1.5 wherexlabel=bottom
     ''')

Result('images-mig-iso','diffs vc_iso vc_mig','TwoRows')


#Result('images-mig-all','diffs vc_mig vc_iso vc_noise','TwoRows')
## < faultmapsimple.jpg sfjpg2byte | sfdd type=float | sftransp > hargrove_fracs1.rsf

Fetch('hargrove_map.jpg','rpsea',beg.server)

Flow('fracmap','hargrove_map.jpg',
     '''
     jpg2byte | dd type=float |
     reverse which=2 | 
     put o1=0 o2=0 d1=0.025 d2=0.025 unit1=km unit2=km |
     reverse which=1 
     ''')
Plot('fracmap',
     '''
     grey title="a) Fault Map" 
     label2="x\_1\^" label1="x\_2\^"
     wheretitle=top yll=1.5 wherexlabel=bottom
     ''')
Result('images-mig-all','fracmap diff_slice_noise iso_slice noise_slice','TwoRows')
End()