from rsf.proj import *
from rsf.prog import RSFROOT

Nt=1000
dt=0.004
t0=0

Nx=128
dx=0.08

# define slowness
Flow('slo',None,'math n1=%d d1=%g o1=%g output="1/(1.5+0.5*x1)" ' %(Nt,dt,t0))

# Make synthetic data
Flow('slo1',None,'math n1=%d d1=%g o1=%g output="1/(1.5+0.5*x1)^2" | spray axis=2 n=2 | pad n2=3' % (Nt,dt,t0))
Flow('data-4','slo1',
     '''
     window n2=1 |
     noise seed=2011 rep=y |
     math output="input^3" |
     cut n1=80 | cut f1=999 |
     ricker1 frequency=10 |
     spray axis=2 n=%d d=%g o=%g label="Inline Offset" unit=km |
     spray axis=3 n=%d d=%g o=%g label="Crossline Offset" unit=km |
     inmo3 velocity=$SOURCE half=n |
     put label1=Time unit1=s |
     mutter half=n tp=0.8 v0=1.5 |
     transp plane=23 | mutter half=n tp=0.8 v0=1.5 | transp plane=23
     ''' % (Nx,dx,-Nx*dx/2,Nx,dx,-Nx*dx/2))
Result('data-4','byte gainpanel=all | grey3 title=" " frame1=%d frame2=%d frame3=%d flat=n point2=0.7' % (Nt/2,Nx/2,Nx/2))
Flow('datacpx','data-4','rtoc')

Flow('offset','data-4','window n1=1 | math output="sqrt(x1*x1+x2*x2)" | put n1=%d n2=1' % (Nx*Nx))

Ntau=Nt
dtau=dt
tau0=t0

Np=128
p0=0.1
dp=(0.7-p0)/Np

Plot('slo','graph transp=y yreverse=y min2=0.1 max2=0.6953 plotfat=0.2 plotcol=4 pad=n wanttitle=n wantaxis=n')



# Apply FFT in time
Flow('fftdata','datacpx','fft3 axis=1 pad=1 | window n1=500 f1=500')
#Result('fftreal','fftdata','real | grey title="real" ')
#Result('fftimag','fftdata','imag | grey title="imag" ')
#Result('fftabs','fftdata','math output="abs(input)" | real | grey title="abs" ')

Flow('fftdatac','fftdata','window n1=110 f1=5')
#Result('fftrealc','fftdatac','real | grey title="realc" ')
#Result('fftimagc','fftdatac','imag | grey title="imagc" ')
#Result('fftabsc','fftdatac','math output="abs(input)" | real | grey title="absc" ')


# direct Radon / single integral (vscan)
#Flow('smod','data-4 offset','put n2=%d n3=1 | vscan offset=${SOURCES[1]} half=n slowness=y nv=%d dv=%g v0=%g extend=0 mute=0 str=0 weight=n type=power | math output=input*%g' % (Nx*Nx,Np,dp,p0,Nt*Nx*Nx*dx*dx))
#Plot('smod','grey label1=Time unit1=s label2=Slowness unit2=s/km title="smod" scalebar=y')
#Result('smod','smod slo','Overlay')


# direct Radon / single integral (nearest integer interpolation)
Flow('dimod','data-4','diradon32 ntau=%d dtau=%g tau0=%g np=%d dp=%g p0=%g | math output=input*%g' % (Ntau,dtau,tau0,Np,dp,p0,dx*dx))
Plot('dimod','grey label1=Time unit1=s label2=Slowness unit2=s/km title=" " ')
Result('dimod-4','dimod slo','Overlay')


# Fast Radon
Flow('bfio.bin',os.path.join(RSFROOT,'include','bfio.bin'),'/bin/cp $SOURCE $TARGET',stdin=0,stdout=-1)
Flow('fmod','fftdatac bfio.bin','radon32 ntau=%d dtau=%g tau0=%g np=%d dp=%g p0=%g fi=1 EL=0 N=64 EPSx1=5 EPSx2=5 EPSk1=5 EPSk2=5 | math output=input*%g' % (Ntau,dtau,tau0,Np,dp,p0,2*dx*dx/Nt))
Plot('fmod','real | grey label1=Time unit1=s label2=Slowness unit2=s/km title=" " ')
Result('fmod-4','fmod slo','Overlay')

#Flow('fmod1','fftdatac bfio.bin','radon32 ntau=%d dtau=%g tau0=%g np=%d dp=%g p0=%g fi=1 EL=0 N=32 EPSx1=9 EPSx2=9 EPSk1=9 EPSk2=9 | math output=input*%g' % (Ntau,dtau,tau0,Np,dp,p0,2*dx*dx/Nt))
#Plot('fmod1','real | grey label1=Time unit1=s label2=Slowness unit2=s/km title="fmod1" scalebar=y')
#Result('fmod1','fmod1 slo','Overlay')

End()