from rsf.proj import*
import math

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

def Wig(data,other):
        Result(data,'window j2=4 | wiggle poly=y transp=y yreverse=y label2=Trace unit2="" label1=Time unit1="s" title="" wherexlabel=b wheretitle=t color=b screenratio=1.4 %s'%other)

def Greyplot(data,other):
        Plot(data,'grey label2=Trace unit2="" label1=Time unit1="s" title="" wherexlabel=b wheretitle=t color=b screenratio=1.4 %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
thr=10

#Create the well-known complex model
#############################################################################
Flow('complex',None,
     '''
     spike n1=512 n2=256 d2=0.1 o2=0 label2=Trace unit2=
     nsp=4 k1=64,100,160,200 p2=1.5,-0.3,0,0.5 mag=1,1,1,1 |
     ricker1 frequency=20 |
     noise seed=2008 var=0 | put d2=1
     ''')

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

## define forward and backward transform strings
forw = 'rtoc | fft3 axis=1 pad=2 |fft3 axis=2 pad=2'
back = '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','complex complex-zero','add scale=1,-1 ${SOURCES[1]}')
Flow('snr0',['complex','diff0'],'snr2 noise=${SOURCES[1]}')

##1 POCS-form shaping (N=2)  d_{n+1} = d_{obs} + (I-M)T[d_{n}], where T=A^{-1}T_{\lambda}A and T_{\lambda} is soft thresholding
sig="complex-zero"
Greyplot(sig,'title="Iteration 0"')
niter=40 # 
data = sig
datas = [sig]
snrs1=['snr0']
for iter in range(niter):
    old = data
    data = 'data-shape%d' % (iter+1)
    snr  ='snr-shape%d'%(iter+1)
    diff ='diff-shape%d'%(iter+1)
    # 1. Forward FFT
    # 2. Threshold in FK domain
    # 3. Inverse FFT
    # 4. Multiply by space mask 2
    # 5. Add data outside of hole
    Flow(data,[old,'complex-mask2',sig],
         '''
         %s | threshold1 type=soft thr=%g | 
         %s | mul ${SOURCES[1]} | 
         add ${SOURCES[2]}
         ''' % (forw,thr,back))
    Flow(diff,['complex',data],'add scale=1,-1 ${SOURCES[1]}')
    Flow(snr,['complex',diff],'snr2 noise=${SOURCES[1]}')
    Greyplot(data,'title="Iteration %d"' % (iter+1))
    datas.append(data)
    snrs1.append(snr)
Plot('shape',datas,'Movie')

## 2 faster-shaping
sig="complex-zero"
Greyplot(sig,'title="Iteration 0"')
data = sig
datas = [sig]
snrs2=['snr0']
old1 = sig
t1=1
for iter in range(niter):
    #frac=1.0/2+(iter/(niter-1))/2.0
    frac = 1.0/2
    #frac=0
    #print frac
    old1 = data
    data = 'data-fshape%d' % (iter+1)
    snr  ='snr-fshape%d'%(iter+1)
    diff ='diff-fshape%d'%(iter+1)
    # 1. multiply by space mask 1
    # 2. add data from last iteration and observed data with scale (-1,1,1)
    # 3. Inverse FFT
    # 4. thresholding in FK domain
    # 5. Forward FFT
    Flow(data,[old1,'complex-mask2',sig],
         ''' 
         %s | threshold1 ifperc=1 type=soft thr=%g | %s |
         mul ${SOURCES[1]} | 
         add scale=1,1 ${SOURCES[2]} |
         add scale=%g,%g ${SOURCES[0]}
         ''' % (forw,thr,back,1+frac,-frac))
    Flow(diff,['complex',data],'add scale=1,-1 ${SOURCES[1]}')
    Flow(snr,['complex',diff],'snr2 noise=${SOURCES[1]}')
    Greyplot(data,'title="Iteration %d"' % (iter+1))
    datas.append(data)
    snrs2.append(snr)
Plot('fshape',datas,'Movie')

## creating data for ploting
Flow('complex-recon-o','data-fshape%d'%(40),'cp')
Flow('complex-recon','data-fshape%d'%(20),'cp')


Flow('complex-error-o','complex-recon-o complex','add scale=1,-1 ${SOURCES[1]}')
Flow('complex-error','complex-recon complex','add scale=1,-1 ${SOURCES[1]}')


Flow('complex-fk','complex','rtoc | fft3 axis=1 pad=2 | fft3 axis=2 pad=2 |cabs')
Flow('complex-zerofk','complex-zero','rtoc | fft3 axis=1 pad=2 | fft3 axis=2 pad=2 |cabs')
Flow('complex-reconfk-o','complex-recon-o','rtoc | fft3 axis=1 pad=2 | fft3 axis=2 pad=2 |cabs ')
Flow('complex-reconfk','complex-recon','rtoc | fft3 axis=1 pad=2 | fft3 axis=2 pad=2 |cabs ')
Flow('complex-errorfk-o','complex-error-o','rtoc | fft3 axis=1 pad=2 | fft3 axis=2 pad=2 |cabs')
Flow('complex-errorfk','complex-error','rtoc | fft3 axis=1 pad=2 | fft3 axis=2 pad=2 |cabs')

## ploting figure for writing a paper
Wig('complex','clip=0.158')
Wig('complex-zero','clip=0.158')
Wig('complex-recon','clip=0.158')
Wig('complex-recon-o','clip=0.158')
Wig('complex-error','clip=0.158')
Wig('complex-error-o','clip=0.158')
Grey('complex-fk','color=j label1=Frequency unit1=Hz label2=Wavenumber allpos=y clip=31')
Grey('complex-zerofk','color=j label1=Frequency unit1=Hz label2=Wavenumber allpos=y clip=31')
Grey('complex-reconfk','color=j label1=Frequency unit1=Hz label2=Wavenumber allpos=y clip=31')
Grey('complex-errorfk','color=j label1=Frequency unit1=Hz label2=Wavenumber allpos=y clip=62')
Grey('complex-reconfk-o','color=j label1=Frequency unit1=Hz label2=Wavenumber allpos=y clip=31')
Grey('complex-errorfk-o','color=j label1=Frequency unit1=Hz label2=Wavenumber allpos=y clip=62')

Grey('complex-mask','color=j')
Grey('complex-mask2','color=j')

## ploting convergence diagram (dashed -> pocs,solid -> pocs)
Flow('SNR1',snrs1,'cat axis=1 ${SOURCES[1:%d]}'%len(snrs1))
Flow('SNR2',snrs2,'cat axis=1 ${SOURCES[1:%d]}'%len(snrs2))

Flow('SNRs','SNR1 SNR2','cat axis=2 ${SOURCES[1]}')
Graph('SNRs','label1="Iteration no. #" label2=SNR unit2=dB dash=0,1')

Graph('SNR2','label1="Iteration no. #" symbol="*" symbolsz=10 label2=SNR unit2=dB dash=1')
End()