from rsf.proj import *

# Download data
Fetch('horizon.asc','hall')

# Convert format
Flow('data','horizon.asc',
     '''
     echo in=$SOURCE data_format=ascii_float n1=3 n2=57036 | 
     dd form=native | window n1=1 f1=-1 | add add=-65 | 
     put 
     n2=291 o2=35.031 d2=0.01 label2=y unit2=km 
     n1=196 o1=33.139 d1=0.01 label1=x unit1=km |
     costaper nw1=25 nw2=25 
     ''')

# Display
def plot(title):
    return '''
    grey color=j title="%s" 
    transp=y yreverse=n clip=14
    ''' % title
Result('data',plot('Horizon'))

# Cut three square holes (!!! CHANGE ME !!!)
cut = '''
cut n1=20 n2=20 f1=125 f2=150 | 
cut n1=20 n2=20 f1=150 f2=50  | 
cut n1=20 n2=20 f1=75  f2=175
'''

Flow('hole','data',cut)
Flow('mask','data',
     'math output=1 | %s | math output=1-input' % cut)
Plot('hole',plot('Input'))
Result('hole','Overlay')

# Fourier transform
forw = 'rtoc | fft3 axis=1 pad=1 | fft3 axis=2 pad=1'
back = 'fft3 axis=2 inv=y | fft3 axis=1 inv=y | real'

for data in ('data','hole'):
    fft = 'fft-'+data
    Flow(fft,data,forw)
    Result(fft,
           '''
           math output="abs(input)" | real | 
           grey allpos=y title="Fourier Transform (%s)" 
           screenratio=1
           ''' % data.capitalize())

# Create Fourier mask
Flow('fft-mask','fft-hole',
     '''
     real | math output="x1*x1+x2*x2" | mask min=50 |
     dd type=float | math output=1-input | 
     smooth rect1=5 rect2=5 repeat=3 | rtoc
     ''')
Result('fft-mask',
       '''
       real | 
       grey allpos=y title="Fourier Mask" screenratio=1
       ''')

# POCS iterations
niter=5 # !!! CHANGE ME !!!

data = 'hole'
plots = ['hole']
for iter in range(niter):
    old = data
    data = 'data%d' % iter

    # 1. Forward FFT
    # 2. Multiply by Fourier mask
    # 3. Inverse FFT
    # 4. Multiply by space mask
    # 5. Add data outside of hole
    Flow(data,[old,'fft-mask','mask','hole'],
         '''
         %s | mul ${SOURCES[1]} | 
         %s | mul ${SOURCES[2]} | 
         add ${SOURCES[3]}
         ''' % (forw,back))
    Plot(data,plot('Iteration %d' % (iter+1)))
    plots.append(data)
# Put frames in a movie
Plot('pocs',plots,'Movie',view=1)

# Last frame
Result('pocs',data,
       plot('POCS interpolation (%d iterations)' % niter))

End()