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

Nt=1000
dt=0.004
t0=0

Nx=1000
dx=0.005
x0=0

# Make synthetic data
Flow('vel',None,'math n1=%d d1=%g o1=%g output="(1.5+0.5*x1)" ' % (Nt,dt,t0))

# From velocity to slowness
Flow('slo','vel','math output="1/input" ')

Flow('data','vel',
     '''
     noise seed=1999 rep=y |
     math output="input^3" |
     cut n1=80 | cut f1=999 |
     ricker1 frequency=10 |
     spray n=%d d=%g o=%g label=Offset unit=km |
     inmo velocity=$SOURCE half=n slowness=n |
     put label1=Time unit1=s |
     mutter half=n tp=0.1 v0=1.5 
     ''' % (Nx,dx,x0))
Result('data','grey title=" " ')
Flow('datacpx','data','rtoc')

#costaper nw1=200 

Ntau=Nt
dtau=dt
tau0=t0

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

Plot('slo','graph transp=y yreverse=y min2=0.1 max2=0.6994 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=" " ')

Flow('fftdatac','fftdata','window n1=100 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','vscan half=n slowness=y nv=%d dv=%g v0=%g type=power weight=n extend=0 mute=0 str=0 | math output=input*%g' % (Np,dp,p0,Nt*Nx*dx))
#Plot('smod','grey label1=Time unit1=s label2=Slowness unit2=s/km title="smod" scalebar=y')
#Result('smod','smod slo','Overlay')

# direct adjoint Radon / single integral (velmod)
#Flow('sdat','smod','velmod half=n slowness=y nh=%d dh=%g h0=%g extend=0 | math output=input*%g' % (Nx,dx,x0,Nt*Nx*dp))
#Result('sdat','grey label1=Time unit1=s label2=Offset unit2=km title="sdat" scalebar=y')

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

# direct ajoint Radon / single integral (nearest integer interpolation)
Flow('didat','dimod','adiradon2 nt=%d dt=%g t0=%g nx=%d dx=%g x0=%g | math output=input*%g' % (Nt,dt,t0,Nx,dx,x0,dp))
Result('didat','grey label1=Time unit1=s label2=Offset unit2=km title=" " ')

# 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','radon2 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/Nt))
Plot('fmod','real | grey label1=Time unit1=s label2=Slowness unit2=s/km title=" " ')
Result('fmod','fmod slo','Overlay')

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

dw=1/Nt/dt
Flow('fdatmid','fmod bfio.bin','radon2 ntau=%d dtau=%g tau0=%g np=%d dp=%g p0=%g fi=2 EL=0 N=32 EPSx1=9 EPSx2=9 EPSk1=9 EPSk2=9' % (100,dw,5*dw,Nx,dx,x0))
Flow('fdat','fdatmid','pad beg1=505 end1=395 | fft3 inv=y axis=1 pad=1 | math output="input*%g" ' % (2*dx))
Result('fdat','real | grey label1=Time unit1=s label2=Offset unit2=km title=" " ')

# direct Radon / double integral
# TEMPORARY !!!
Flow('diimod','fftdatac','diiradon2 ntau=%d dtau=%g tau0=%g np=%d dp=%g p0=%g | math output=input*%g' % (Ntau,dtau,tau0,Np,dp,p0,2*dx/Nt))
Plot('diimod','real | grey label1=Time unit1=s label2=Slowness unit2=s/km title="diimod" ')
Result('diimod','diimod slo','Overlay')

Flow('diff','fmod diimod','add scale=1,-1 ${SOURCES[1]}')
Result('diff','real | grey title=" " clip=4.49516 label1=Time unit1=s label2=Slowness unit2=s/km')


# compute inner product
Flow('fhu','fmod fmod fdat datacpx','sum input2=${SOURCES[1]} input3=${SOURCES[2]} input4=${SOURCES[3]}')

End()