from rsf.proj import*

Fetch('sean.HH','bp')

Flow('sean','sean.HH',
     'dd form=native | window n3=1 f3=3 n1=500 | scale axis=2 ')

mute='mutter hyper=y half=y t0=1.8 x0=0 slope0=0.005 | mutter hyper=n half=y t0=2.1 x0=0 slope0=0.002'

## module definition
def Grey(data,other):
        Result(data,'grey label2=Trace unit2="" label1=Time unit1="s" title="" wherexlabel=b screenratio=1.4 wheretitle=t color=g clip=0.30 %s'%other)

def Greyplot(data,other):
        Plot(data,'grey label2=Trace unit2="" label1=Time unit1="s" title="" wherexlabel=b screenratio=1.4 wheretitle=t color=g clip=0.30 %s'%other)

def Graph(data,other):
        Result(data,'graph label1="" label2="" unit1="" unit2=""  title="" wherexlabel=b wheretitle=t %s' %other)

def Graphplot(data,other):
        Plot(data,'graph label1="" label2="" unit1="" unit2=""  title="" wherexlabel=b wheretitle=t %s' %other)
## parameters definition
zeroperc=30

Flow('sean-testmute','sean','noise var=0.01 | mutter hyper=y half=y t0=1.8 x0=0 slope0=0.005 ')
Grey('sean-testmute','')

## zero trace (remove certain percent of traces)
Flow('sean-mask','sean','window n1=1 | noise type=y range=0.5 mean=0.5 rep=y seed=201414| threshold1 ifperc=1 type=hard thr=%d | mask min=0.000000001 | spray axis=1 n=500 |dd type=float'%(100-zeroperc))
Flow('sean-mask2','sean-mask','math output="1-input"')
Flow('sean-zero','sean sean-mask','add mode=p ${SOURCES[1]}')

Grey('sean','clip=0.30')
Grey('sean-mask','color=j')
Grey('sean-mask','color=j')
Grey('sean-zero','clip=0.30')

## define forward and backward transform strings
forw2 = 'rtoc | fft3 axis=1 pad=2 |fft3 axis=2 pad=2'
back2 = 'fft3 axis=2 pad=2 inv=y | fft3 axis=1 pad=2 inv=y |real'

## compute the initial snr (SNR=10log(sum(s^2)/sum(n^2))
Flow('diff0','sean sean-zero','add scale=1,-1 ${SOURCES[1]}')
Flow('snr0',['sean','diff0'],'snr2 noise=${SOURCES[1]}')
sig="sean-zero"
Greyplot(sig,'title="Iteration 0"')
niter=40
thr=24
percthr=19

## POCS (N=2) hard -constant
data = sig
datas = [sig]
snrs1=['snr0']

for iter in range(niter):
    old = data
    data = 'data1-%d' % (iter+1)
    snr  ='snr1-%d'%(iter+1)
    diff ='diff1-%d'%(iter+1)
    # 1. Forward FFT
    # 2. Threshold in FK domain
    # 3. Inverse FFT
    # 4. Multiply by space mask
    # 5. Add data outside of hole
    Flow(data,[old,'sean-mask',sig,'sean-mask2'],
         '''
         mul ${SOURCES[1]} | 
         add scale=-1,1,1 ${SOURCES[2]} ${SOURCES[0]}  | 
         %s | threshold1 type=hard  thr=%g ifverb=1 ifperc=0 | 
         %s  | mul ${SOURCES[3]} | add scale=1,1 ${SOURCES[2]} |%s
         ''' % (forw2,thr,back2,mute))
    Flow(diff,['sean',data],'add scale=1,-1 ${SOURCES[1]}')
    Flow(snr,['sean',diff],'snr2 noise=${SOURCES[1]}')
    Greyplot(data,'title="Iteration %d"' % (iter+1))
    datas.append(data)
    snrs1.append(snr)
Plot('pocs1',datas,'Movie')


## POCS  (N=2) soft -constant
data = sig
datas = [sig]
snrs2=['snr0']
for iter in range(niter):
    old = data
    data = 'data2-%d' % (iter+1)
    snr  ='snr2-%d'%(iter+1)
    diff ='diff2-%d'%(iter+1)
    # 1. Forward FFT
    # 2. Threshold in FK domain
    # 3. Inverse FFT
    # 4. Multiply by space mask
    # 5. Add data outside of hole
    Flow(data,[old,'sean-mask',sig,'sean-mask2'],
         '''
         mul ${SOURCES[1]} | 
         add scale=-1,1,1 ${SOURCES[2]} ${SOURCES[0]}  | 
         %s | threshold1 type=soft  thr=%g ifverb=1 ifperc=0 | 
         %s  | mul ${SOURCES[3]} | add scale=1,1 ${SOURCES[2]} |%s
         ''' % (forw2,thr,back2,mute))
    Flow(diff,['sean',data],'add scale=1,-1 ${SOURCES[1]}')
    Flow(snr,['sean',diff],'snr2 noise=${SOURCES[1]}')
    Greyplot(data,'title="Iteration %d"' % (iter+1))
    datas.append(data)
    snrs2.append(snr)
Plot('pocs2',datas,'Movie')

## POCS  (N=2) half thresholding -constant
data = sig
datas = [sig]
snrs3=['snr0']
for iter in range(niter):
    old = data
    data = 'data3-%d' % (iter+1)
    snr  ='snr3-%d'%(iter+1)
    diff ='diff3-%d'%(iter+1)
    # 1. Forward FFT
    # 2. Threshold in FK domain
    # 3. Inverse FFT
    # 4. Multiply by space mask
    # 5. Add data outside of hole
    Flow(data,[old,'sean-mask',sig,'sean-mask2'],
         '''
         mul ${SOURCES[1]} | 
         add scale=-1,1,1 ${SOURCES[2]} ${SOURCES[0]}  | 
         %s | halfthr  thr=%g ifverb=1 ifperc=0 | 
         %s | mul ${SOURCES[3]} | add scale=1,1 ${SOURCES[2]} |%s
         ''' % (forw2,thr,back2,mute))
    Flow(diff,['sean',data],'add scale=1,-1 ${SOURCES[1]}')
    Flow(snr,['sean',diff],'snr2 noise=${SOURCES[1]}')
    Greyplot(data,'title="Iteration %d"' % (iter+1))
    datas.append(data)
    snrs3.append(snr)
Plot('pocs3',datas,'Movie')




## POCS (N=2) hard - percentile
data = sig
datas = [sig]
snrs4=['snr0']

for iter in range(niter):
    old = data
    data = 'data4-%d' % (iter+1)
    snr  ='snr4-%d'%(iter+1)
    diff ='diff4-%d'%(iter+1)
    # 1. Forward FFT
    # 2. Threshold in FK domain
    # 3. Inverse FFT
    # 4. Multiply by space mask
    # 5. Add data outside of hole
    Flow(data,[old,'sean-mask',sig,'sean-mask2'],
         '''
         mul ${SOURCES[1]} | 
         add scale=-1,1,1 ${SOURCES[2]} ${SOURCES[0]}  | 
         %s | threshold1 type=hard  thr=%g ifverb=1 ifperc=1 | 
         %s  | mul ${SOURCES[3]} | add scale=1,1 ${SOURCES[2]} |%s
         ''' % (forw2,percthr,back2,mute))
    Flow(diff,['sean',data],'add scale=1,-1 ${SOURCES[1]}')
    Flow(snr,['sean',diff],'snr2 noise=${SOURCES[1]}')
    Greyplot(data,'title="Iteration %d"' % (iter+1))
    datas.append(data)
    snrs4.append(snr)
Plot('pocs4',datas,'Movie')


## POCS  (N=2) soft -percentile
data = sig
datas = [sig]
snrs5=['snr0']
for iter in range(niter):
    old = data
    data = 'data5-%d' % (iter+1)
    snr  ='snr5-%d'%(iter+1)
    diff ='diff5-%d'%(iter+1)
    # 1. Forward FFT
    # 2. Threshold in FK domain
    # 3. Inverse FFT
    # 4. Multiply by space mask
    # 5. Add data outside of hole
    Flow(data,[old,'sean-mask',sig,'sean-mask2'],
         '''
         mul ${SOURCES[1]} | 
         add scale=-1,1,1 ${SOURCES[2]} ${SOURCES[0]}  | 
         %s | threshold1 type=soft  thr=%g ifverb=1 ifperc=1 | 
         %s  | mul ${SOURCES[3]} | add scale=1,1 ${SOURCES[2]} | %s
         ''' % (forw2,percthr,back2,mute))
    Flow(diff,['sean',data],'add scale=1,-1 ${SOURCES[1]}')
    Flow(snr,['sean',diff],'snr2 noise=${SOURCES[1]}')
    Greyplot(data,'title="Iteration %d"' % (iter+1))
    datas.append(data)
    snrs5.append(snr)
Plot('pocs5',datas,'Movie')

## POCS  (N=2) half thresholding -percentile
data = sig
datas = [sig]
snrs6=['snr0']
for iter in range(niter):
    old = data
    data = 'data6-%d' % (iter+1)
    snr  ='snr6-%d'%(iter+1)
    diff ='diff6-%d'%(iter+1)
    # 1. Forward FFT
    # 2. Threshold in FK domain
    # 3. Inverse FFT
    # 4. Multiply by space mask
    # 5. Add data outside of hole
    Flow(data,[old,'sean-mask',sig,'sean-mask2'],
         '''
         mul ${SOURCES[1]} | 
         add scale=-1,1,1 ${SOURCES[2]} ${SOURCES[0]}  | 
         %s | halfthr  thr=%g ifverb=1 ifperc=1 | 
         %s | mul ${SOURCES[3]} | add scale=1,1 ${SOURCES[2]} |%s
         ''' % (forw2,percthr,back2,mute))
    Flow(diff,['sean',data],'add scale=1,-1 ${SOURCES[1]}')
    Flow(snr,['sean',diff],'snr2 noise=${SOURCES[1]}')
    Greyplot(data,'title="Iteration %d"' % (iter+1))
    datas.append(data)
    snrs6.append(snr)
Plot('pocs6',datas,'Movie')

## ploting convergence diagram 
Flow('SNR1',snrs1,'cat axis=1 ${SOURCES[1:%d]}'%len(snrs1))
Flow('SNR2',snrs2,'cat axis=1 ${SOURCES[1:%d]}'%len(snrs2))
Flow('SNR3',snrs3,'cat axis=1 ${SOURCES[1:%d]}'%len(snrs3))
Flow('SNR4',snrs4,'cat axis=1 ${SOURCES[1:%d]}'%len(snrs4))
Flow('SNR5',snrs5,'cat axis=1 ${SOURCES[1:%d]}'%len(snrs5))
Flow('SNR6',snrs6,'cat axis=1 ${SOURCES[1:%d]}'%len(snrs6))

Flow('sean-SNRs','SNR2 SNR3 SNR5 SNR6','cat axis=2 ${SOURCES[1:4]}')
#Graph('SNRs','label1="Iteration no. #" label2=SNR unit2=dB symbol="-*ohsf" symbolsz=10')
Graph('sean-SNRs','label1="Iteration no. #" label2=SNR unit2=dB symbol="*osp" symbolsz=10')

## Result comparison

## Plot result
Grey('data2-%d'%niter,'')
Grey('data3-%d'%niter,'')
Grey('data5-%d'%niter,'')
Grey('data6-%d'%niter,'')

Result('sean-data2','Fig/data2-%d.vpl'%niter,'Overlay')
Result('sean-data3','Fig/data3-%d.vpl'%niter,'Overlay')
Result('sean-data5','Fig/data5-%d.vpl'%niter,'Overlay')
Result('sean-data6','Fig/data6-%d.vpl'%niter,'Overlay')

## Plot difference
Greyplot('diff2-%d'%niter,'clip=0.3')
Greyplot('diff3-%d'%niter,'clip=0.3')
Greyplot('diff5-%d'%niter,'clip=0.3')
Greyplot('diff6-%d'%niter,'clip=0.3')
Result('sean-diff2','diff2-%d'%niter,'Overlay')
Result('sean-diff3','diff3-%d'%niter,'Overlay')
Result('sean-diff5','diff5-%d'%niter,'Overlay')
Result('sean-diff6','diff6-%d'%niter,'Overlay')

## FK domain ploting (without log)
Flow('seanfk','sean','%s | cabs'%forw2)
Flow('sean-zerofk','sean-zero','%s | cabs'%forw2)
Flow('data2-%d'%niter+'fk','data2-%d'%niter,'%s|cabs'%forw2)
Flow('diff2-%d'%niter+'fk','diff2-%d'%niter,'%s|cabs'%forw2)
Greyplot('seanfk','label1="Frequency" unit1=Hz label2="Normalized wavenumber" uni2= color=j')
Greyplot('sean-zerofk','label1="Frequency" unit1=Hz label2="Normalized wavenumber" uni2= color=j')
Greyplot('data2-%d'%niter+'fk','label1="Frequency" unit1=Hz label2="Normalized wavenumber" uni2= color=j')
Greyplot('diff2-%d'%niter+'fk','label1="Frequency" unit1=Hz label2="Normalized wavenumber" uni2= color=j')
Result('compfk','seanfk sean-zerofk '+'data2-%d'%niter+'fk'+' diff2-%d'%niter+'fk','SideBySideIso')


## FK domain ploting (with log)
Flow('seanfklog','sean','%s | cabs | math output="log(input)"'%forw2)
Flow('sean-zerofklog','sean-zero','%s | cabs |math output="log(input)"'%forw2)
Flow('data2-%d'%niter+'fklog','data2-%d'%niter,'%s|cabs |math output="log(input)"'%forw2)
Flow('diff2-%d'%niter+'fklog','diff2-%d'%niter,'%s|cabs |math output="log(input)"'%forw2)
Greyplot('seanfklog','label1="Frequency" unit1=Hz label2="Normalized wavenumber" uni2= color=j')
Greyplot('sean-zerofklog','label1="Frequency" unit1=Hz label2="Normalized wavenumber" uni2= color=j')
Greyplot('data2-%d'%niter+'fklog','label1="Frequency" unit1=Hz label2="Normalized wavenumber" uni2= color=j')
Greyplot('diff2-%d'%niter+'fklog','label1="Frequency" unit1=Hz label2="Normalized wavenumber" uni2= color=j')
Result('compfklog','seanfklog sean-zerofklog '+'data2-%d'%niter+'fklog'+' diff2-%d'%niter+'fklog','SideBySideIso')

End()