from rsf.proj import *
import math

def plot(title):
    return '''
    wiggle title="%s" poly=y yreverse=y transp=y
    ''' % (title)

Flow('data',None,
     '''
     spike n1=64 n2=32 d2=1 o2=1 d1=1 o1=1 
     label2=Trace label1=Samples unit1= unit2=
     nsp=3 k1=8,20,32 k2=4 l2=28 p2=2,1,0 
     ''')
Result('data',plot('Data: 3 sloping reflectors--Use dt=dx=1'))

zap = {'flat':   'v1=-2.5 v2=-2 v3=2 v4=2.5',
       'middle': 'v1=0.5  v2=0.75 v3=1.25 v4=1.5 pass=n',
       'steep':  'v1=0    v2=0.25 v3=0.75 v4=1   pass=n'}

Flow('fft','data','fft1 | fft3')

Result('fft',
       '''
       math output="abs(input)" | real | 
       grey title="F-K Spectrum" allpos=y
       ''')

for event in zap.keys():
    Flow(event,'fft',
         '''
         dipfilter %s | 
         fft3 inv=y | fft1 inv=y
         ''' % zap[event])
    Result(event,plot('zap ' + event))

# Linear moveout to unalias steep event

Flow('fft2','data',
     '''
     fft1 | math output="input*exp(I*%g*x2*x1)" | fft3
     ''' % (2*2*math.pi))

Result('fft2',
       '''
       math output="abs(input)" | real | 
       grey title="F-K Spectrum" allpos=y
       ''')
Flow('steep2','fft2',
     '''
     dipfilter %s | 
     fft3 inv=y | 
     math output="input*exp(-I*%g*x2*x1)" |
     fft1 inv=y
     ''' % (zap['flat'],2*2*math.pi))
Result('steep2',plot('zap steep'))

End()