from rsf.proj import*
import math

## module definition
def Grey(data,other):
        Result(data,'grey label2=Trace unit2=""  labelsz=10 labelfat=4 font=2 titlesz=10 titlefat=4  label1=Time unit1="s" title="" wherexlabel=b wheretitle=t color=b clip=0.003 %s '%other)

def Greyplot(data,other):
        Plot(data,'grey label2=Trace unit2=""  labelsz=10 labelfat=4 font=2 titlesz=10 titlefat=4  label1=Time unit1="s" title="" wherexlabel=b wheretitle=t color=b %s'%other)

def Graph(data,other):
        Result(data,'graph label1="" label2="" unit1=""  labelsz=10 labelfat=4 font=2 titlesz=10 titlefat=4  unit2=""  title="" wherexlabel=b wheretitle=t %s' %other)

def Graphplot(data,other):
        Plot(data,'graph label1="" label2="" unit1=""  labelsz=10 labelfat=4 font=2 titlesz=10 titlefat=4  unit2=""  title="" wherexlabel=b wheretitle=t %s' %other)

def sqrt(a):
        return math.sqrt(a)

## parameters definition 
zeroperc=30
ifperc=1
thr=15
mode="soft"
ddip=50
padno=256
r1=10
r2=10
fhi=60
niter=34

#Create the well-known sigmoid model
#############################################################################
Flow('sigmoid',None,
     '''
     sigmoid d1=.004 n1=200 d2=.008 n2=256 |
     smooth rect1=3 diff1=1 | smooth rect1=3 |
     put label2=Distance | put d2=1
     ''')

## zero trace (remove certain percent of traces)
Flow('sigmoid-mask','sigmoid','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=200 |dd type=float'%(100-zeroperc))
Flow('sigmoid-mask2','sigmoid-mask','math output="1-input"')
Flow('sigmoid-zero','sigmoid sigmoid-mask','add mode=p ${SOURCES[1]}')

Grey('sigmoid','')
Grey('sigmoid-mask','color=j')
Grey('sigmoid-mask','color=j')
Grey('sigmoid-zero','')

## 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'

## define forward and backward transform strings
forwseis1 = 'seislet adj=y inv=y dip=${SOURCES[3]} eps=0.1 type=b '
backseis1 = 'seislet adj=n inv=y dip=${SOURCES[3]} eps=0.1 type=b '

forwseis2 = 'seislet adj=y inv=y dip=${SOURCES[4]} eps=0.1 type=b '
backseis2 = 'seislet adj=n inv=y dip=${SOURCES[4]} eps=0.1 type=b '

## compute the initial snr (SNR=10log(sum(s^2)/sum(n^2))
Flow('diff0','sigmoid sigmoid-zero','add scale=1,-1 ${SOURCES[1]}')
Flow('snr0',['sigmoid','diff0'],'snr2 noise=${SOURCES[1]}')

## FK
Flow('sigmoidfk','sigmoid','%s | cabs | window f2=200'%forw)
Grey('sigmoidfk','color=j')
Flow('sigmoidfktest','sigmoidfk','threshold1 type=hard thr=12')


##1## FFT POCS iterations (N=2, soft thresholding)  
sig="sigmoid-zero"
Greyplot(sig,'title="Iteration 0"')
data = sig
datas = [sig]
snrs_pocs=['snr0']
for iter in range(niter):
    old = data
    data = 'data-1pocs%d' % (iter+1)
    snr  ='snr-1pocs%d'%(iter+1)
    diff ='diff-1pocs%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,'sigmoid-mask2',sig],
         '''
         %s | threshold1 ifperc=%d type=%s thr=%g ifverb=1 | 
         %s | mul ${SOURCES[1]} | 
         add ${SOURCES[2]}
         ''' % (forw,ifperc,mode,thr,back))
    Flow(diff,['sigmoid',data],'add scale=1,-1 ${SOURCES[1]}')
    Flow(snr,['sigmoid',diff],'snr2 noise=${SOURCES[1]}')
    Greyplot(data,'title="Iteration %d"' % (iter+1))
    datas.append(data)
    snrs_pocs.append(snr)
Plot('pocs',datas,'Movie')

##2## FFT FPOCS 
sig="sigmoid-zero"
Greyplot(sig,'title="Iteration 0"')
data = sig
datas = [sig]
snrs_fpocs=['snr0']
old1 = sig
t1=1
for iter in range(niter):
    t2=t1
    t1=(1+sqrt(1+4*t2*t2))/2
    frac=(t2-1)/t1
    old2 = old1
    old1 = data
    data = 'data-1fpocs%d' % (iter+1)
    snr  ='snr-1fpocs%d'%(iter+1)
    diff ='diff-1fpocs%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,old2,'sigmoid-mask2',sig],
         '''
         add scale=%g,%g ${SOURCES[1]} | 
         %s | threshold1 ifperc=%d type=%s thr=%g | %s |
         mul ${SOURCES[2]} | 
         add scale=1,1 ${SOURCES[3]}
         ''' % (1+frac,-frac,forw,ifperc,mode,thr,back))
    Flow(diff,['sigmoid',data],'add scale=1,-1 ${SOURCES[1]}')
    Flow(snr,['sigmoid',diff],'snr2 noise=${SOURCES[1]}')
    Greyplot(data,'title="Iteration %d"' % (iter+1))
    datas.append(data)
    snrs_fpocs.append(snr)
Plot('fpocs',datas,'Movie')


thr=25
##3## Seislet POCS iterations (soft thresholding)  
sig="sigmoid-zero"
Greyplot(sig,'title="Iteration 0"')
data = sig
datas = [sig]
snrs_pocs_seis=['snr0']
for iter in range(niter):
    old = data
    if iter % ddip ==0 :
        dip='dip-pocs%g'%int(iter/ddip)
        Flow(dip,'sigmoid',
            '''
            bandpass fhi=%g | pad n2=%g |                       
            dip rect1=%g rect2=%g
            '''%(fhi,padno,r1,r2))
    data = 'data-1pocs-seis%d' % (iter+1)
    snr  ='snr-1pocs-seis%d'%(iter+1)
    diff ='diff-1pocs-seis%d'%(iter+1)
    # Seislet based POCS
    Flow(data,[old,'sigmoid-mask2',sig,dip],
         '''
         %s | threshold1 ifperc=%d type=%s thr=%g ifverb=1 | 
         %s | mul ${SOURCES[1]} | 
         add scale=1,1 ${SOURCES[2]} 
         ''' % (forwseis1,ifperc,mode,thr,backseis1))
    Flow(diff,['sigmoid',data],'add scale=1,-1 ${SOURCES[1]}')
    Flow(snr,['sigmoid',diff],'snr2 noise=${SOURCES[1]}')
    Greyplot(data,'title="Iteration %d"' % (iter+1))
    datas.append(data)
    snrs_pocs_seis.append(snr)
Plot('pocs_seis',datas,'Movie')

##4## Seislet FPOCS 
sig="sigmoid-zero"
Greyplot(sig,'title="Iteration 0"')
data = sig
datas = [sig]
snrs_fpocs_seis=['snr0']
old1 = sig
t1=1
for iter in range(niter):
    t2=t1
    t1=(1+sqrt(1+4*t2*t2))/2
    frac=(t2-1)/t1

    old2 = old1
    old1 = data
    if iter % ddip ==0 :
        dip='dip-fpocs%g'%int(iter/ddip)
        Flow(dip,'sigmoid',
             '''
             bandpass fhi=%g | pad n2=%g | 
             dip rect1=%g rect2=%g
             '''%(fhi,padno,r1,r2))

    data = 'data-1fpocs-seis%d' % (iter+1)
    snr  ='snr-1fpocs-seis%d'%(iter+1)
    diff ='diff-1fpocs-seis%d'%(iter+1)

    # Seislet based Fast POCS
    Flow(data,[old1,old2,'sigmoid-mask2',sig, dip],
         '''
         add scale=%g,%g ${SOURCES[1]} | 
         %s | threshold1 ifperc=%d type=%s thr=%g | %s |
         mul ${SOURCES[2]} | 
         add scale=1,1 ${SOURCES[3]}
         ''' % (1+frac,-frac,forwseis2,ifperc,mode,thr,backseis2))

    Flow(diff,['sigmoid',data],'add scale=1,-1 ${SOURCES[1]}')
    Flow(snr,['sigmoid',diff],'snr2 noise=${SOURCES[1]}')
    Greyplot(data,'title="Iteration %d"' % (iter+1))
    datas.append(data)
    snrs_fpocs_seis.append(snr)
Plot('fpocs-seis',datas,'Movie')

## ploting convergence diagram
Flow('SNR-POCS',snrs_pocs,'cat axis=1 ${SOURCES[1:%d]}'%len(snrs_pocs))
Flow('SNR-FPOCS',snrs_fpocs,'cat axis=1 ${SOURCES[1:%d]}'%len(snrs_pocs))
Flow('SNR-POCS-seis',snrs_pocs_seis,'cat axis=1 ${SOURCES[1:%d]}'%len(snrs_pocs))
Flow('SNR-FPOCS-seis',snrs_fpocs_seis,'cat axis=1 ${SOURCES[1:%d]}'%len(snrs_pocs))
Flow('SNRs','SNR-POCS SNR-FPOCS SNR-POCS-seis SNR-FPOCS-seis','cat axis=2 ${SOURCES[1:4]}')

Graph('SNRs','label1="Iteration no. #" label2=SNR unit2=dB symbol="-o*." plotcol="5,2,3,4" symbolsz=10 max2=19 title="SNR Comparison"')

Flow('data-pocs-fft','data-1pocs%d'%niter,'cp')
Flow('data-fpocs-fft','data-1fpocs%d'%niter,'cp')
Flow('data-pocs-seis','data-1pocs-seis%d'%niter,'cp')
Flow('data-fpocs-seis','data-1fpocs-seis%d'%niter,'cp')

Flow('diff-pocs-fft','diff-1pocs%d'%niter,'cp')
Flow('diff-fpocs-fft','diff-1fpocs%d'%niter,'cp')
Flow('diff-pocs-seis','diff-1pocs-seis%d'%niter,'cp')
Flow('diff-fpocs-seis','diff-1fpocs-seis%d'%niter,'cp')

Flow('fk','sigmoid','%s | cabs | window f1=200'%forw)
Flow('fk-zero','sigmoid-zero','%s | cabs | window f1=200'%forw)
Flow('fk-pocs-fft','data-pocs-fft','%s | cabs | window f1=200'%forw)
Flow('fk-fpocs-fft','data-fpocs-fft','%s | cabs | window f1=200'%forw)
Flow('fk-pocs-seis','data-pocs-seis','%s | cabs | window f1=200'%forw)
Flow('fk-fpocs-seis','data-fpocs-seis','%s | cabs | window f1=200'%forw)

Grey('data-pocs-fft','')
Grey('data-fpocs-fft','')
Grey('data-pocs-seis','')
Grey('data-fpocs-seis','')

Grey('diff-pocs-fft','')
Grey('diff-fpocs-fft','')
Grey('diff-pocs-seis','')
Grey('diff-fpocs-seis','')

Grey('fk','color=j label1=Frequency unit1=Hz label2="Normalized wavenumber" clip=1.3')
Grey('fk-zero','color=j label1=Frequency unit1=Hz label2="Normalized wavenumber" clip=1.3')
Grey('fk-pocs-fft','color=j label1=Frequency unit1=Hz label2="Normalized wavenumber" clip=1.3')
Grey('fk-fpocs-fft','color=j label1=Frequency unit1=Hz label2="Normalized wavenumber" clip=1.3')
Grey('fk-pocs-seis','color=j label1=Frequency unit1=Hz label2="Normalized wavenumber" clip=1.3')
Grey('fk-fpocs-seis','color=j label1=Frequency unit1=Hz label2="Normalized wavenumber" clip=1.3')

Flow('data-z','sigmoid','window min1=0.4 max1=0.6 min2=150 max2=200')
Flow('data-zero-z','sigmoid-zero','window min1=0.4 max1=0.6 min2=150 max2=200')
Flow('data-pocs-fft-z','data-pocs-fft','window min1=0.4 max1=0.6 min2=150 max2=200')
Flow('data-fpocs-fft-z','data-fpocs-fft','window min1=0.4 max1=0.6 min2=150 max2=200')
Flow('data-pocs-seis-z','data-pocs-seis','window min1=0.4 max1=0.6 min2=150 max2=200')
Flow('data-fpocs-seis-z','data-fpocs-seis','window min1=0.4 max1=0.6 min2=150 max2=200')

Grey('data-z','')
Grey('data-zero-z','')
Grey('data-pocs-fft-z','')
Grey('data-fpocs-fft-z','')
Grey('data-pocs-seis-z','')
Grey('data-fpocs-seis-z','')

## Creating framebox
x=280
y=1.03
w=90
w1=0.5

Flow('frame.asc',None,'echo %s n1=10 data_format=ascii_float in=$TARGET'% \
        ' '.join(map(str,(x,y,x+w,y,x+w,y+w1,x,y+w1,x,y))))
Plot('frame','frame.asc',
        '''
        dd type=complex form=native |
        graph min1=0 max1=471 min2=0 max2=2.048 pad=n plotfat=15 plotcol=2 
        wantaxis=n wanttitle=n yreverse=y 
        ''')
Result('sigmoid-0','Fig/sigmoid.vpl frame','Overlay')
Result('sigmoid-zero-0','Fig/sigmoid-zero.vpl frame','Overlay')
Result('data-pocs-fft-0','Fig/data-pocs-fft.vpl frame','Overlay')
Result('data-fpocs-fft-0','Fig/data-fpocs-fft.vpl frame','Overlay')
Result('data-pocs-seis-0','Fig/data-pocs-seis.vpl frame','Overlay')
Result('data-fpocs-seis-0','Fig/data-fpocs-seis.vpl frame','Overlay')

# adding reference trace
Flow('trace.asc',None,
         '''
         echo %d 0 %d 1 n1=4 in=$TARGET data_format=ascii_float
         ''' % (175,175)) #81 is the trace number
Plot('trace','trace.asc',
         '''
         dd form=native type=complex | 
         graph min1=0 max1=255 min2=0 title="" wantaxis=n scalebar=n pad=n plotfat=8 dash=2
         ''') #250 is the number of traces

Result('sigmoid-1','Fig/sigmoid.vpl frame trace','Overlay')
Result('sigmoid-zero-1','Fig/sigmoid-zero.vpl frame trace','Overlay')
Result('data-pocs-fft-1','Fig/data-pocs-fft.vpl frame trace','Overlay')
Result('data-fpocs-fft-1','Fig/data-fpocs-fft.vpl frame trace','Overlay')
Result('data-pocs-seis-1','Fig/data-pocs-seis.vpl frame trace','Overlay')
Result('data-fpocs-seis-1','Fig/data-fpocs-seis.vpl frame trace','Overlay')

## Similarity
Flow('simi-pocs-fft','sigmoid data-pocs-fft','similarity other=${SOURCES[1]} niter=20 rect1=5 rect2=5')
Flow('simi-fpocs-fft','sigmoid data-fpocs-fft','similarity other=${SOURCES[1]} niter=20 rect1=5 rect2=5')
Flow('simi-pocs-seis','sigmoid data-pocs-seis','similarity other=${SOURCES[1]} niter=20 rect1=5 rect2=5')
Flow('simi-fpocs-seis','sigmoid data-fpocs-seis','similarity other=${SOURCES[1]} niter=20 rect1=5 rect2=5')
Grey('simi-pocs-fft','color=j clip=2 allpos=y scalebar=y minval=0.4 maxval=1')
Grey('simi-fpocs-fft','color=j clip=2 allpos=y scalebar=y minval=0.4 maxval=1')
Grey('simi-pocs-seis','color=j clip=2 allpos=y scalebar=y minval=0.4 maxval=1')
Grey('simi-fpocs-seis','color=j clip=2 allpos=y scalebar=y minval=0.4 maxval=1')

Flow('simi-diff-pocs-fft','sigmoid diff-pocs-fft','similarity other=${SOURCES[1]} niter=20 rect1=5 rect2=5')
Flow('simi-diff-fpocs-fft','sigmoid diff-fpocs-fft','similarity other=${SOURCES[1]} niter=20 rect1=5 rect2=5')
Flow('simi-diff-pocs-seis','sigmoid diff-pocs-seis','similarity other=${SOURCES[1]} niter=20 rect1=5 rect2=5')
Flow('simi-diff-fpocs-seis','sigmoid diff-fpocs-seis','similarity other=${SOURCES[1]} niter=20 rect1=5 rect2=5')
Grey('simi-diff-pocs-fft','color=j clip=1 allpos=y scalebar=y minval=0.0 maxval=1')
Grey('simi-diff-fpocs-fft','color=j clip=1 allpos=y scalebar=y minval=0.0 maxval=1')
Grey('simi-diff-pocs-seis','color=j clip=1 allpos=y scalebar=y minval=0.0 maxval=1')
Grey('simi-diff-fpocs-seis','color=j clip=1 allpos=y scalebar=y minval=0.0 maxval=1')

Flow('trace-comp','sigmoid data-fpocs-fft data-fpocs-seis','cat axis=3 ${SOURCES[1:3]} | window n2=1 f2=175')
Graph('trace-comp','label2=Amplitude unit2= label1=Time unit1=s screenratio=0.6 dash=0,3,5 plotcol="7,3,5" symbolsz=10 max2=0.006 min2=-0.005 title="Amplitude Comparison"')

Flow('true','trace-comp','window n2=1 ')
Flow('tfft','trace-comp','window n2=1 f2=1 ')
Flow('tseis','trace-comp','window n2=1 f2=2 ')
Flow('tsimi1','true tfft','similarity other=${SOURCES[1]} niter=20 rect1=5')
Flow('tsimi2','true tseis','similarity other=${SOURCES[1]} niter=20 rect1=5')
Flow('tsimi-comp-0','tsimi1 tsimi2','cat axis=2 ${SOURCES[1]}')


Graph('tsimi-comp-0','label2=Similarity unit2= label1=Time unit1=s screenratio=0.6 dash=3,5 plotcol="3,5" plotfat=10 symbolsz=10 max2=1.1 min2=0 title="Local Similarity Comparison"')

Plot('Seislet',None,
        '''
        box x0=7.5 y0=8.2 label="Seislet" xt=-0.5 yt=-0.5 length=2
        ''')

Plot('FK',None,
        '''
        box x0=13 y0=4.8 label="FK" xt=-0.5 yt=-0.5 length=1 
        ''')
Result('tsimi-comp','Fig/tsimi-comp-0.vpl FK Seislet','Overlay')



End()