from rsf.proj import *

def cubeplot(title,clip='',extra=''):
    return '''
    byte gainpanel=all %s |
    grey3 frame1=32 frame2=256 frame3=10 flat=y point1=0.7 point2=0.7
    label1=Offset unit1=km label2="Midpoint wavenumber" unit2=1/km
    label3="Frequency" unit3=Hz
    title="%s" %s wanttitle=n labelfat=4 font=2 titlefat=4
    ''' % (clip,title,extra)

Fetch('french.asc','french')

Flow('french','french.asc',
     '''
     dd form=native | transp | scale dscale=0.0005 |
     put d1=0.10265 d2=0.10265
     label1=North-South label2=West-East unit1=km unit2=km
     ''')

Flow('french1','french.asc','dd form=native | transp | scale dscale=2')
Flow('refl','french1',
     '''
     remap1 n1=161 o1=0 d1=51.325 | transp |
     remap1 n1=161 o1=0 d1=51.325 | transp
     ''')

Flow('slice','french',
     '''
     window n1=1 f1=30 | put d1=0.025 |
     remap1 n1=256 o1=0 d1=0.008 |
     unif2 n1=256 d1=0.004 v00=1,2
     ''')
Result('slice','refl',
       '''
       unif3 v00=1.5,2.5 n1=401 d1=10.265 | 
       put d1=0.003849375 d2=0.012720496894409938 d3=0.012720496894409938 |
       transp plane=23 |
       byte allpos=y bias=1 | window n3=1 f3=60 |
       grey color=I wanttitle=n label1="Depth (km)" label2=Inline unit2=km
       labelfat=4 font=2 titlefat=4 flat=n
       ''')

Flow('cup','slice',
     '''
     deriv | bandpass flo=10 fhi=50 |
     transp memsize=1000| bandpass fhi=50 | transp
     ''')

Flow('data','cup',
     '''
     halfint inv=y |
     preconstkirch zero=y inv=y h0=0 dh=0.008 nh=64 vel=1.5 |
     window
     ''')

Result('data',
       cubeplot('Data','clip=3.1713','label1=Time unit1=s \
       label2="Midpoint" unit2=km label3="Half offset" unit3=km \
       frame1=150 frame2=125 frame3=25') )

Flow('dnmo','data',
     'transp plane=23 | nmostretch v0=1.5 | transp plane=23')

Flow('dfftd','dnmo',
     'window f1=40 | logstretch | fft1 | transp | transp plane=23')

# F-K domain
Flow('dfk','dfftd','fft3 axis=1')

# Test fkoclet
Flow('dinput','dfk','transp memsize=1000')
Result('dinput',
       'real | transp plane=13 |'
       + cubeplot('Input','','label1=Frequency unit1=Hz \
       label3="Half offset" label2="Midpoint wavenumber" unit2=1/km unit3=km \
       frame1=20 frame3=25 frame2=257') )

Flow('dtran','dinput','fkoclet verb=y type=b adj=n inv=n')
Result('dtran',
       'put d1=1 | real | transp plane=13 |'
       + cubeplot('Forward transform','','label1=Frequency \
       unit1=Hz label3=Scale unit3="" label2="Midpoint wavenumber" unit2=1/km \
       frame1=20 frame3=0 frame2=257') )

# Test Fourier
Flow('dfourier','dinput','fft3 axis=1 | window f1=64')
Result('dfourier',
       'put d1=1 o1=0 | real | transp plane=13 |'
       + cubeplot('Forward transform','','label1=Frequency \
       unit1=Hz label3=Scale unit3="" label2="Midpoint wavenumber" unit2=1/km \
       frame1=20 frame3=0 frame2=257') )
Flow('dcfourier','dfourier',
     '''
     math output="abs(input)" | real |
     put n1=3571712 d1=1 label1=Scale unit1= n2=1 n2=1 n3=1 |
     scale axis=1 | sort memsize=2000 | math output="log(input)"
     ''')

Flow('dctran','dtran',
     '''
     math output="abs(input)" | real | 
     put n1=3571712 d1=1 label1=Scale unit1= n2=1 n2=1 n3=1 |
     scale axis=1 | sort memsize=2000 | math output="log(input)"
     ''')

compare = '''
       cat axis=2 ${SOURCES[1:2]} |
       window n1=71434 | put d1=2.7997778096330277e-04 |
       graph  wanttitle=n dash=0,1 plotfat=5 labelfat=3 gridfat=3 
       label2="a\_\s75 n"  unit2=dB
       label1="Percentage of total coefficient number (%)" wheretitle=b 
       grid=n pad=n labelfat=4 font=2 titlefat=4 gridfat=4
       '''
box = '''
      box x0=%g y0=%g label="%s" xt=%g yt=%g lab_fat=4 font=2 titlefat=4
      '''
Plot('ccom','dctran dcfourier',compare )
Plot('ocseislet',None,box % (6.0,2.65,"OC-seislet transform",-1,-0.5))
Plot('fourier',None,box % (8.2,4.25,"Fourier transform",1,0.5))
Result('compare','ccom ocseislet fourier','Overlay')

######################
# Noise test
######################
Flow('noise','data','noise seed=2008 var=0.285472')
Result('noise',
       cubeplot('Noisy data','clip=3.1713','label1=Time unit1=s \
       label2="Midpoint" unit2=km label3="Half offset" unit3=km \
       frame1=150 frame2=125 frame3=25') )

Flow('nmo','noise',
     'transp plane=23 | nmostretch v0=1.5 | transp plane=23')

Flow('ndwt','nmo',
     '''
     transp plane=13 |
     dwt type=b inv=y unit=y |
     transp plane=13 
     ''')

Flow('ndthr','ndwt','threshold pclip=1.7')

Flow('nidwt','ndthr',
     '''
     transp plane=13 memsize=1000 |
     dwt type=b inv=y unit=y adj=y |
     transp plane=13 memsize=1000 |
     transp plane=23 memsize=1000 |
     nmostretch v0=1.5 inv=y |
     transp plane=23 memsize=1000
     ''')
Result('nidwt',
       cubeplot('Denoised data','clip=3.1713','label1=Time unit1=s \
       label2="Midpoint" unit2=km label3="Half offset" \
       unit3=km frame1=150 frame2=125 frame3=25') )

Flow('diff1','noise nidwt','add scale=1,-1 ${SOURCES[1]}')
Result('diff1',
       cubeplot('Denoised data','clip=3.1713','label1=Time unit1=s \
       label2="Midpoint" unit2=km label3="Half offset" \
       unit3=km frame1=150 frame2=125 frame3=25') )

Flow('nlevel1','data nidwt',
     '''
     add scale=1,-1 ${SOURCES[1]} |
     math output="input^2" |
     stack axis=3 | stack axis=2 | stack axis=1
     ''')
Flow('slevel','data',
     '''
     math output="input^2" |
     stack axis=3 | stack axis=2 | stack axis=1
     ''')
Flow('snr1','slevel nlevel1',
     'math d=${SOURCES[0]} n=${SOURCES[1]} output="10*log(d/n)"')

Flow('fftd','nmo',
     'window f1=40 | logstretch | fft1 | transp | transp plane=23')

# F-K domain
Flow('fk','fftd','fft3 axis=1')

# Test fkoclet
Flow('input','fk','transp memsize=1000')
Result('input',
       'real | transp plane=13 |'
       + cubeplot('Input','','label1=Frequency unit1=Hz \
       label3="Half offset" label2="Midpoint wavenumber" unit2=1/km unit3=km \
       frame1=20 frame3=25 frame2=257') )

Flow('tran','input','fkoclet verb=y type=b adj=n inv=n')
Result('tran',
       'put d1=1 | real | transp plane=13 |'
       + cubeplot('Forward transform','','label1=Frequency \
       unit1=Hz label3=Scale unit3="" label2="Midpoint wavenumber" unit2=1/km \
       frame1=20 frame3=0 frame2=257') )

# Retrun to X-T
Flow('inv-tran','tran',
     '''
     transp | fft3 axis=1 inv=y | transp plane=23 |
     transp | fft1 inv=y | logstretch inv=y | pad beg1=40 |
     transp plane=23 | nmostretch v0=1.5 inv=y | transp plane=23
     ''')
Result('inv-tran',
       'put d3=1 |'
       + cubeplot('Inverse','','label1=Time unit1=s \
       label2="Midpoint" unit2=km label3=Scale unit3= \
       frame1=150 frame2=125 frame3=0') )

# Thresholding
Flow('thr','tran','threshold pclip=2')

# Inverse seislet transform
Flow('ithr','thr','fkoclet inv=y adj=y inv=y verb=y')
Result('ithr',
       'put d1=1 | real | transp plane=13 |'
       + cubeplot('Forward transform','','label1=Frequency unit1=Hz \
       label3="Half offset" unit3="km" label2="Midpoint wavenumber"\
       unit2=1/km frame1=20 frame3=0 frame2=257') )
# Retrun to X-T
Flow('inver','ithr',
     '''
     transp memsize=1000 | fft3 axis=1 inv=y | transp plane=23 memsize=1000 |
     transp memsize=1000 | fft1 inv=y | logstretch inv=y | pad beg1=40 |
     transp plane=23 memsize=1000 | nmostretch v0=1.5 inv=y |
     transp plane=23 memsize=1000
     ''')
Result('inver',
       cubeplot('Denoised data','clip=3.1713','label1=Time unit1=s \
       label3="Half offset" unit3=km label2="Midpoint"\
       unit2=km frame1=150 frame2=125 frame3=25') )
Flow('diff2','noise inver','add scale=1,-1 ${SOURCES[1]}')
Result('diff2',
       cubeplot('Denoised data','clip=3.1713','label1=Time unit1=s \
       label2="Midpoint" unit2=km label3="Half offset" \
       unit3=km frame1=150 frame2=125 frame3=25') )

Flow('nlevel2','data inver',
     '''
     add scale=1,-1 ${SOURCES[1]} |
     math output="input^2" |
     stack axis=3 | stack axis=2 | stack axis=1
     ''')
Flow('snr2','slevel nlevel2',
     'math d=${SOURCES[0]} n=${SOURCES[1]} output="10*log(d/n)"')

Result('snr','snr1 snr2',
       '''
       cat axis=1 ${SOURCES[1]} |
       graph symbol=X symbolsz=10 min2=0 max2=50 min1=-0.5 max1=1.5
       label2=SNR unit2=dB label1=Methods unit1= wanttitle=n
       ''')

End()