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('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
     ''')

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

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

# Zero out random traces
Flow('cmask','cdata',
     'window n1=1 | noise rep=y type=n seed=2008 | mask min=0.3')
Flow('czero','cdata cmask','headercut mask=${SOURCES[1]}')

Flow('ccmask','cmask',
     '''
     dd type=float | 
     spray axis=1 n=256 d=0.004 o=0
     ''')
Flow('scmask','cmask',
     '''
     dd type=float | math output=1-input |
     spray axis=1 n=256 d=0.004 o=0
     ''')

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

Flow('cnmo','czero',
     'transp plane=23 | nmostretch v0=1.5 | transp plane=23')

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

Flow('cndthr','cndwt','threshold pclip=0.5')
Plot('cndthr',
     'put d1=1 | transp plane=23 |'
     +cubeplot('Missing data','','label1=Time unit1=s label2="Midpoint" \
     unit2=km label3="Half offset" unit3=km frame1=150 frame2=200 frame3=0'))

Flow('cnidwt','cndthr',
     '''
     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
     ''')

Flow('cfftd','cnmo',
     'window f1=40 | logstretch | fft1 | transp | transp plane=23')

# F-K domain
Flow('cfk','cfftd','fft3 axis=1')

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

Flow('ctran','cinput','fkoclet verb=y type=b adj=n inv=n')
Result('ctran',
       '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('cinv-tran','ctran',
     '''
     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
     ''')


# Thresholding
Flow('cthr','ctran','threshold pclip=0.5')

# Inverse seislet transform
Flow('cithr','cthr','fkoclet inv=y adj=y inv=y verb=y')
Result('cithr',
       'real | transp plane=13 |'
       + cubeplot('Input','','label1=Frequency unit1=Hz \
       label3="Half offset" label2="Midpoint wavenumber" unit2=1/km unit3=km \
       frame1=20 frame3=0 frame2=257') )

#########################
# Fourier POCS
#########################
Flow('tczero','czero','transp plane=23 memsize=1000')
Flow('tccmask','ccmask','transp plane=23 memsize=1000')

# 3-D Fourier POCS 
fniter=100
fforward = '''
fft1 | fft3 | fft3 axis=3
'''
fbackward = '''
fft3 axis=3 inv=y | fft3 inv=y | fft1 inv=y
'''
fdata = 'czero'
fplots = ['czero']
for iter in range(fniter):
    fold = fdata
    fdata = 'fdata%d' % iter

    # 1. Forward OC-seislet
    # 2. Thresholding
    # 3. Inverse OC-seislet
    # 4. Multiply by space mask
    # 5. Add data outside of hole
    Flow(fdata,[fold,'scmask','czero'],
         fforward +
         '''
         | threshold pclip=%g |
         ''' % (1.)
         + fbackward +
         '''
         | add mode=p ${SOURCES[1]} | 
         add ${SOURCES[2]}
         ''')

Flow('cffinal',fdata,
     fforward +
     '''
     | threshold pclip=85. |
     '''
     + fbackward
     )

# Last frame
Result('four3pocs','cffinal',
       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') )

#########################
# OC-seislet POCS
#########################
pniter=50
forward = '''
window f1=40 | logstretch | fft1 |
transp memsize=1000 | transp plane=23 memsize=1000 |
fft3 axis=1 | transp memsize=1000 |
fkoclet adj=n inv=n dwt=n verb=n type=b
'''
backward = '''
fkoclet dwt=n adj=y inv=y verb=n type=b |
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
'''
ftdata = 'cnmo'
ftplots = ['cnmo']
for iter in range(pniter):
    ftold = ftdata
    ftdata = 'ftdata%d' % iter

    # 1. Forward OC-seislet
    # 2. Thresholding
    # 3. Inverse OC-seislet
    # 4. Multiply by space mask
    # 5. Add data outside of hole
    Flow(ftdata,[ftold,'scmask','cnmo'],
         forward +
         '''
         | threshold pclip=%g |
         ''' % (1.)
         + backward +
         '''
         | add mode=p ${SOURCES[1]} | 
         add ${SOURCES[2]}
         ''')

Flow('cfinal',ftdata,
     forward +
     '''
     | threshold pclip=85. |
     '''
     + backward
     )

# Last frame
Flow('cpocs','cfinal',
     '''
     transp plane=23 memsize=1000 |
     nmostretch v0=1.5 inv=y |
     transp plane=23 memsize=1000
     ''')
Result('cpocs',
       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') )

End()