from rsf.proj import*
import math
wf = 2*math.pi
dx = 1
n1=257
Flow('plane',None,
     '''
     spike n1=512 n2=256 d2=1 o2=0 label2=Trace unit2=
     nsp=3 k1=64,160,286 p2=0.5,1,0 mag=0.5,0.5,1 |
     ricker1 frequency=20 |
     noise seed=2008 var=0.000000000000000001
     ''')

def wiggle(title):
    return '''
           window j2=4 |
           wiggle transp=y yreverse=y poly=y clip=0.15
           title="%s" wheretitle=b wherexlabel=t color=j
           ''' % title

Result('plane',wiggle(''))

# Fourier transform in time
Flow('fft','plane','fft1')
Result('fft','real | window max1=60 | grey title="Frequency domain" clip=1.5')

# Take a frequency slice
Flow('fft1','fft','window n1=1 min1=5')
Result('fft1','real | graph title="Frequency Slice at 5 Hz" label2= unit2= unit1=')


##########
# dip filtering
##########
foursr1=[] #real part
foursr2=[]
foursr3=[]
foursr4=[]
foursi1=[] #imag part
foursi2=[]
foursi3=[]
foursi4=[]

for i in range(n1):
        fourr0='fourr0%d'%(i+1)
        fourr1='fourr1%d'%(i+1)
        fourr2='fourr2%d'%(i+1)
        #fourr3='fourr3%d'%(i+1)
        #fourr4='fourr4%d'%(i+1)
        fr='fr%d'%(i+1)

        Flow(fr,'fft','window n1=1 f1=%d |real'%i)
        Flow(fourr0,fr,'emd | pad n2=10 ')
        Flow(fourr1,fourr0,'window n2=1')
        Flow(fourr2,fourr0,'window f2=1 n2=1')
        #Flow(fourr3,fourr0,'window f2=2 n2=1')
        #Flow(fourr4,fourr0,'window f2=3 n2=1')
        foursr1.append(fourr1)
        foursr2.append(fourr2)
        #foursr3.append(fourr3)
        #foursr4.append(fourr4)

for i in range(n1):
        fouri0='fouri0%d'%(i+1)
        fouri1='fouri1%d'%(i+1)
        fouri2='fouri2%d'%(i+1)
        #fouri3='fouri3%d'%(i+1)
        #fouri4='fouri4%d'%(i+1)
        fi='fi%d'%(i+1)

        Flow(fi,'fft','window n1=1 f1=%d |imag'%i)
        Flow(fouri0,fi,'emd | pad n2=10 ')
        Flow(fouri1,fouri0,'window n2=1')
        Flow(fouri2,fouri0,'window f2=1 n2=1')
        #Flow(fouri3,fouri0,'window f2=2 n2=1')
        #Flow(fouri4,fouri0,'window f2=3 n2=1')
        foursi1.append(fouri1)
        foursi2.append(fouri2)
        #foursi3.append(fouri3)
        #foursi4.append(fouri4)

Flow('fftemdr1',foursr1,'cat axis=2 ${SOURCES[1:%d]}'%len(foursr1))
Flow('fftemdr2',foursr2,'cat axis=2 ${SOURCES[1:%d]}'%len(foursr2))
#Flow('fftemdr3',foursr3,'cat axis=2 ${SOURCES[1:%d]}'%len(foursr3))
#Flow('fftemdr4',foursr4,'cat axis=2 ${SOURCES[1:%d]}'%len(foursr4))

Flow('fftemdi1',foursi1,'cat axis=2 ${SOURCES[1:%d]}'%len(foursi1))
Flow('fftemdi2',foursi2,'cat axis=2 ${SOURCES[1:%d]}'%len(foursr2))
#Flow('fftemdi3',foursi3,'cat axis=2 ${SOURCES[1:%d]}'%len(foursr3))
#Flow('fftemdi4',foursi4,'cat axis=2 ${SOURCES[1:%d]}'%len(foursr4))

Flow('fftemd1','fftemdr1 fftemdi1','cmplx ${SOURCES[1]} | transp')
Flow('fftemd2','fftemdr2 fftemdi2','cmplx ${SOURCES[1]} | transp')
#Flow('fftemd3','fftemdr3 fftemdi3','cmplx ${SOURCES[1]} | transp')
#Flow('fftemd4','fftemdr4 fftemdi4','cmplx ${SOURCES[1]} | transp')


Flow('planeemd1','fftemd1','fft1 inv=y | put unit2= ')
Flow('planeemd2','fftemd2','fft1 inv=y | put unit2= ')
#Flow('planeemd3','fftemd3','fft1 inv=y')
#Flow('planeemd4','fftemd4','fft1 inv=y')
#Flow('res','plane planeemd1 planeemd2 planeemd3 planeemd4','add scale=1,-1,-1,-1,-1 ${SOURCES[1:5]}')
Flow('planeemd3','plane planeemd1 planeemd2','add scale=1,-1,-1 ${SOURCES[1:3]}')

Result('planeemd1',wiggle(''))
Result('planeemd2',wiggle(''))
Result('planeemd3',wiggle(''))
#Result('planeemd4',wiggle(''))
#Result('res',wiggle('residual'))


End()