from rsf.proj import*
import sys
sys.path.append('../../temp/recipes')
import deblend2

## parameters definition
clip=0.9                #display percentage
fraction=0.5    #B=fraction*I
niter=40        #number of iterations
n1=1001         #temporal sampling number
n2=751          #spatial sampling number
padno=1024      #padding for seislet tranform
r1=20           #smoothing radius
r2=20           #smoothing radius
ddip=5          #changing dip map interval
fhi=50          #bandpass frequency value
mute1='cp'      #muting function in shaping
mute2='cp'      #muting function in shaping
mode='soft'     #thresholding type (hard thresholding turns to be very bad
thr=20          #thresholding level(percentage)

## module defining
def Grey(complex,other):
        Result(complex,'grey label2=Trace unit2="" label1=Time unit1="s" labelsz=11 title="" wherexlabel=b wheretitle=t %s'%other)

def Greyplot(complex,other):
        Plot(complex,'grey label2=Trace unit2="" label1=Time unit1="s" labelsz=11 title="" wherexlabel=b wheretitle=t %s'%other)

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

def Graphplot(complex,other):
        Plot(complex,'graph label1="" label2="" unit1="" unit2=""  labelsz=11 title="" wherexlabel=b wheretitle=t %s' %other)


##########################################
#    Make synthetic complex:complex1* & complex2
##########################################
##########################################
#    Make synthetic data:data1* & data2
##########################################
Flow('section1',None,
     '''
     spike n1=1001 |
     noise seed=2011 rep=y |
     math output="input^3" |
     cut n1=100 |
     ricker1 frequency=10 |
     spray axis=2 n=751 d=0.01 o=-3.5 label=Offset unit=km |
     nmostretch inv=y v0=4 half=n
     ''')
Flow('section2',None,
     '''
     spike n1=1001 |
     noise seed=2011 rep=y |
     math output="input^3" |
     cut n1=100 |
     ricker1 frequency=10 |
     spray axis=2 n=751 d=0.01 o=-3.5 label=Offset unit=km |
     nmostretch inv=y v0=0.5 half=n|
     mutter hyper=y tp=20 v0=1
     ''')
Flow('section3',None,
     '''
     spike n1=1001 |
     noise seed=2011 rep=y |
     math output="input^3" |
     cut n1=100 |
     ricker1 frequency=10 |
     spray axis=2 n=2000 d=0.01 o=-10 label=Offset unit=km |
     nmostretch inv=y v0=4 half=y | window f2=1249   
     ''')
Flow('complex1','section1 section2 section3','add scale=1,1,1 ${SOURCES[1]} ${SOURCES[2]} | scale axis=2 | put d2=1 o2=0')

Flow('usection1',None,
     '''
     spike n1=1001 |
     noise seed=2012 rep=y |
     math output="input^3" |
     cut n1=100 |
     ricker1 frequency=10 |
     spray axis=2 n=751 d=0.01 o=-3.5 label=Offset unit=km |
     nmostretch inv=y v0=4 half=n
     ''')

Flow('usection2',None,
     '''
     spike n1=1001 |
     noise seed=2012 rep=y |
     math output="input^3" |
     cut n1=100 |
     ricker1 frequency=10 |
     spray axis=2 n=751 d=0.01 o=-3.5 label=Offset unit=km |
     nmostretch inv=y v0=0.5 half=n|
     mutter hyper=y tp=20 v0=1
     ''')

Flow('usection3',None,
     '''
     spike n1=1001 |
     noise seed=2012 rep=y |
     math output="input^3" |
     cut n1=100 |
     ricker1 frequency=10 |
     spray axis=2 n=2000 d=0.01 o=-10 label=Offset unit=km |
     nmostretch inv=y v0=4 half=y | window f2=1249   
     ''')
Flow('complex2','usection1 usection2 usection3','add scale=1,1,1 ${SOURCES[1]} ${SOURCES[2]} | scale axis=2 | put d2=1 o2=0 ')

Grey('complex1',' clip=%g'%clip)
Grey('complex2',' clip=%g'%clip)
#############################################
#               Experiment
#############################################
## Apply dithering
# var=1 makes the dithering range larger, unit=ms
Flow('dither','complex1',
     '''
     window n1=1 |
     noise rep=y seed=122011 var=0.1 | math output="400*input"
     ''')
Flow('complexshottime1','complex1','window n1=1 | math output=6*1000*x1')
Flow('complexshottime2','complexshottime1 dither','add scale=1,1 ${SOURCES[1]}')

## Blend 
Flow('complexs','complex2 complex1 complexshottime1 complexshottime2','blend shot_time_in=${SOURCES[3]} shot_time_out=${SOURCES[2]} |add scale=1,1 ${SOURCES[1]}' )
Flow('ucomplexs','complex1 complex2 complexshottime1 complexshottime2','blend shot_time_in=${SOURCES[2]} shot_time_out=${SOURCES[3]} |add scale=1,1 ${SOURCES[1]}' )

Grey('complexs',' clip=%g'%clip)
Grey('ucomplexs',' clip=%g'%clip)
Flow('complexsfft','complexs','math output="input/2"')
Flow('ucomplexsfft','ucomplexs','math output="input/2"')
Flow('complexsslet','complexs','math output="input/2"')
Flow('ucomplexsslet','ucomplexs','math output="input/2"')
Flow('complexsfxdecon','complexs','math output="input/2"')
Flow('ucomplexsfxdecon','ucomplexs','math output="input/2"')

## fk transform and filtering
Flow('complexsfka','complexs','fft1 | fft3 axis=2 pad=1 | cabs')
Flow('complexsfkr','complexs','fft1 | fft3 axis=2 pad=1 | real')
Flow('complexsfki','complexs','fft1 | fft3 axis=2 pad=1 | imag')
Flow('complexsfk','complexs','fft1 | fft3 axis=2 pad=1')
Flow('complexsfkr_filt','complexsfkr','mutter half=n t0=0 slope0=40 x0=0 ')
Flow('complexsfki_filt','complexsfki','mutter half=n t0=0 slope0=40 x0=0 ')

Grey('complexsfka',' label1=Frequency unit1=Hz label2=Wavenumber color=j')
Grey('complexsfkr',' label1=Frequency unit1=Hz label2=Wavenumber color=j')
Grey('complexsfki',' label1=Frequency unit1=Hz label2=Wavenumber color=j')
Grey('complexsfkr_filt',' label1=Frequency unit1=Hz label2=Wavenumber color=j')
Grey('complexsfki_filt',' label1=Frequency unit1=Hz label2=Wavenumber color=j')

## fk inverse transform -> recon and difference
Flow('complexs_re','complexsfkr_filt complexsfki_filt','cmplx ${SOURCES[1]} | fft3 axis=2 inv=y | fft1 inv=y')
Flow('complexs_redif','complexs complexs_re','add scale=1,-1 ${SOURCES[1]}')
Grey('complexs_re','')
Grey('complexs_redif','')

#deblend using fft
deblend2.deblendfft('complex1',
          'complex2',
          'complexs',
          'ucomplexs',
          'complexsfft',
          'ucomplexsfft',
          'complexdeblendedfft1',
          'complexdeblendedfft2',
          'complexshottime1',
          'complexshottime2',
           mute1,
           mute2,
           n1,
           n2,
           niter,
           mode,
           7,
           clip,
           fraction)

#deblend using seislet
deblend2.deblendslet('complex1',
          'complex2',
          'complexs',
          'ucomplexs',
          'complexsslet',
          'ucomplexsslet',
          'complexdeblendedslet1',
          'complexdeblendedslet2',
          'complexshottime1',
          'complexshottime2',
           mute1,
           mute2,
           n1,
           n2,
           r1,
           r2,
           padno,
           niter,
           ddip,
           mode,
           thr,
           clip,
           fhi,
           fraction)

#deblend using fxdecon
deblend2.deblendfxdecon('complex1',
          'complex2',
          'complexs',
          'ucomplexs',
          'complexsfxdecon',
          'ucomplexsfxdecon',
          'complexdeblendedfxdecon1',
          'complexdeblendedfxdecon2',
          'complexshottime1',
          'complexshottime2',
           mute1,
           mute2,
           n1,
           n2,
           niter,
           clip,
           fraction)

## Ploting difference, error, deblended sections
Flow('complexdifffft1','complexs complexdeblendedfft1','add scale=1,-1 ${SOURCES[1]}')
Flow('complexdifffft2','ucomplexs complexdeblendedfft2','add scale=1,-1 ${SOURCES[1]}')
Flow('complexdiffslet1','complexs complexdeblendedslet1','add scale=1,-1 ${SOURCES[1]}')
Flow('complexdiffslet2','ucomplexs complexdeblendedslet2','add scale=1,-1 ${SOURCES[1]}')
Flow('complexdifffxdecon1','complexs complexdeblendedfxdecon1','add scale=1,-1 ${SOURCES[1]}')
Flow('complexdifffxdecon2','ucomplexs complexdeblendedfxdecon2','add scale=1,-1 ${SOURCES[1]}')

Flow('complexerrorslet1','complex1 complexdeblendedslet1','add scale=1,-1 ${SOURCES[1]}')
Flow('complexerrorslet2','complex2 complexdeblendedslet2','add scale=1,-1 ${SOURCES[1]}')
Flow('complexerrorfxdecon1','complex1 complexdeblendedfxdecon1','add scale=1,-1 ${SOURCES[1]}')
Flow('complexerrorfxdecon2','complex2 complexdeblendedfxdecon2','add scale=1,-1 ${SOURCES[1]}')
Flow('complexerrorfft1','complex1 complexdeblendedfft1','add scale=1,-1 ${SOURCES[1]}')
Flow('complexerrorfft2','complex2 complexdeblendedfft2','add scale=1,-1 ${SOURCES[1]}')




Grey('complexdifffft1','title="" clip=%g'%clip)
Grey('complexdifffft2','title="" clip=%g'%clip)
Grey('complexdiffslet1','title="" clip=%g'%clip)
Grey('complexdiffslet2','title="" clip=%g'%clip)
Grey('complexdifffxdecon1','title="" clip=%g'%clip)
Grey('complexdifffxdecon2','title="" clip=%g'%clip)

Grey('complexerrorfft1','title="" clip=%g'%clip)
Grey('complexerrorfft2','title="" clip=%g'%clip)
Grey('complexerrorslet1','title="" clip=%g'%clip)
Grey('complexerrorslet2','title="" clip=%g'%clip)
Grey('complexerrorfxdecon1','title="" clip=%g'%clip)
Grey('complexerrorfxdecon2','title="" clip=%g'%clip)


#Grey('complexdeblendedslet1','title="Deblended 1 (Seislet)" clip=%g'%clip)
#Grey('complexdeblendedslet2','title="Deblended 1 (Seislet)" clip=%g'%clip)

Grey('complexdeblendedfft1',' clip=%g'%clip)
Grey('complexdeblendedfft2',' clip=%g'%clip)
Grey('complexdeblendedslet1',' clip=%g'%clip)
Grey('complexdeblendedslet2',' clip=%g'%clip)
Grey('complexdeblendedfxdecon1',' clip=%g'%clip)
Grey('complexdeblendedfxdecon2',' clip=%g'%clip)

## Ploting
Flow('complexsnrsa','complexsfft-snrsa complexsslet-snrsa complexsfxdecon-snrsa','cat axis=2 ${SOURCES[1:3]}')
Flow('complexsnrsb','ucomplexsfft-snrsb ucomplexsslet-snrsb ucomplexsfxdecon-snrsb','cat axis=2 ${SOURCES[1:3]}')

Graph('complexsnrsa','dash=0,1,0 title=""  symbol="o+*" symbolsz=8 label1="Iteration no. #" label2="SNR" unit2="dB"  min1=0 max1=%g min2=0 max2=18 d1=1'%niter)
Graph('complexsnrsb','dash=0,1,0 title=""  symbol="o+*" symbolsz=8 label1="Iteration no. #" label2="SNR" unit2="dB"  min1=0 max1=%g min2=0 max2=21 d1=1'%niter)

Greyplot('complexs',' title="Iter # = %g"'%(0))
Greyplot('ucomplexs',' title="Iter # = %g"'%(0))

Greyplot('complexs',' title="Iter # = %g"'%(0))
Greyplot('ucomplexs',' title="Iter # = %g"'%(0))

deblendslets1=['complexs']
deblendslets2=['ucomplexs']


End()