from rsf.proj import*
from rsf.recipes import deblend2

## parameters definition
clip=0.17       #display percentage
fraction=0.5    #B=fraction*I
niter=40        #number of iterations
n1=512          #temporal sampling number
n2=256          #spatial sampling number
padno=256       #padding for seislet tranform
r1=10           #smoothing radius
r2=10           #smoothing radius
ddip=5          #changing dip map interval
fhi=60          #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=8           #thresholding level(percentage)

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

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

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

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

## Generating synthetic data
Flow('data1',None,
     '''
     spike n1=512 n2=256 d2=1 o2=0 label2=Trace unit2=
     nsp=4 k1=64,160,200 p2=1.5,0,0.5 mag=1,1,1 |
     ricker1 frequency=20 |
     noise seed=2008 var=0
     ''')
Grey('data1',' clip=%g'%clip)
Flow('data2',None,
     '''
     spike n1=512 n2=256 d2=1 o2=0 label2=Trace unit2=
     nsp=4 k1=74,150,205 p2=1.3,0.2,1.0 mag=1,1,1 |
     ricker1 frequency=20 |
     noise seed=2008 var=0
     ''')
Grey('data2',' clip=%g'%clip)

## Apply dithering
# var=1 makes the dithering range larger, unit=ms
Flow('dither','data1',
     '''
     window n1=1 |
     noise rep=y seed=122011 var=0.1 | math output="1000*input"
     ''')
Flow('shottime1','data1','window n1=1 | math output=3*1000*x1') # +&- 3ms
Flow('shottime2','shottime1 dither','add scale=1,1 ${SOURCES[1]}')

## Blend 
Flow('datatemp','data1 data2','cat axis=2 ${SOURCES[1]}')
Flow('shottimetemp','shottime1 shottime2','cat axis=2 ${SOURCES[1]}')
Flow('datastemp','datatemp shottimetemp','blend shot_time_in=${SOURCES[1]} shot_time_out=${SOURCES[1]}')
#Flow('datas','datastemp','window n2=%g'%n2)
#Flow('udatas','datastemp','window n2=%g f2=%g'%(n2,n2-1))

Flow('datas','data2 data1 shottime1 shottime2','blend shot_time_in=${SOURCES[3]} shot_time_out=${SOURCES[2]} |add scale=1,1 ${SOURCES[1]}' )
Flow('udatas','data1 data2 shottime1 shottime2','blend shot_time_in=${SOURCES[2]} shot_time_out=${SOURCES[3]} |add scale=1,1 ${SOURCES[1]}' )

Grey('datas','  clip=%g'%clip)
Grey('udatas','  clip=%g'%clip)
Flow('datasfft','datas','math output="input/2"')
Flow('udatasfft','udatas','math output="input/2"')
Flow('datasslet','datas','math output="input/2"')
Flow('udatasslet','udatas','math output="input/2"')
Flow('datasfxdecon','datas','math output="input/2"')
Flow('udatasfxdecon','udatas','math output="input/2"')

## fk transform and filtering
Flow('datasfka','datas','fft1 | fft3 axis=2 pad=1 | cabs')
Flow('datasfkr','datas','fft1 | fft3 axis=2 pad=1 | real')
Flow('datasfki','datas','fft1 | fft3 axis=2 pad=1 | imag')
Flow('datasfk','datas','fft1 | fft3 axis=2 pad=1')
Flow('datasfkr_filt','datasfkr','mutter half=n t0=0 slope0=40 x0=0 ')
Flow('datasfki_filt','datasfki','mutter half=n t0=0 slope0=40 x0=0 ')

Grey('datasfka',' label1=Frequency unit1=Hz label2=Wavenumber color=j')
Grey('datasfkr',' label1=Frequency unit1=Hz label2=Wavenumber color=j')
Grey('datasfki',' label1=Frequency unit1=Hz label2=Wavenumber color=j')
Grey('datasfkr_filt',' label1=Frequency unit1=Hz label2=Wavenumber color=j')
Grey('datasfki_filt',' label1=Frequency unit1=Hz label2=Wavenumber color=j')

## fk inverse transform -> recon and difference
Flow('datas_re','datasfkr_filt datasfki_filt','cmplx ${SOURCES[1]} | fft3 axis=2 inv=y | fft1 inv=y')
Flow('datas_redif','datas datas_re','add scale=1,-1 ${SOURCES[1]}')
Grey('datas_re','')
Grey('datas_redif','')

#deblend using fft
deblend2.deblendfft('data1',
          'data2',
          'datas',
          'udatas',
          'datasfft',
          'udatasfft',
          'deblendedfft1',
          'deblendedfft2',
          'shottime1',
          'shottime2',
           mute1,
           mute2,
           n1,
           n2,
           niter,
           mode,
           thr,
           clip,
           fraction)

#deblend using seislet
deblend2.deblendslet('data1',
          'data2',
          'datas',
          'udatas',
          'datasslet',
          'udatasslet',
          'deblendedslet1',
          'deblendedslet2',
          'shottime1',
          'shottime2',
           mute1,
           mute2,
           n1,
           n2,
           r1,
           r2,
           padno,
           niter,
           ddip,
           mode,
           thr,
           clip,
           fhi,
           fraction)

#deblend using fxdecon
deblend2.deblendfxdecon('data1',
          'data2',
          'datas',
          'udatas',
          'datasfxdecon',
          'udatasfxdecon',
          'deblendedfxdecon1',
          'deblendedfxdecon2',
          'shottime1',
          'shottime2',
           mute1,
           mute2,
           n1,
           n2,
           niter,
           clip,
           fraction)

## Ploting difference, error, deblended sections
Flow('difffft1','datas deblendedfft1','add scale=1,-1 ${SOURCES[1]}')
Flow('difffft2','udatas deblendedfft2','add scale=1,-1 ${SOURCES[1]}')
Flow('diffslet1','datas deblendedslet1','add scale=1,-1 ${SOURCES[1]}')
Flow('diffslet2','udatas deblendedslet2','add scale=1,-1 ${SOURCES[1]}')
Flow('difffxdecon1','datas deblendedfxdecon1','add scale=1,-1 ${SOURCES[1]}')
Flow('difffxdecon2','udatas deblendedfxdecon2','add scale=1,-1 ${SOURCES[1]}')
Flow('errorfft1','data1 deblendedfft1','add scale=1,-1 ${SOURCES[1]}')
Flow('errorfft2','data2 deblendedfft2','add scale=1,-1 ${SOURCES[1]}')
Flow('errorslet1','data1 deblendedslet1','add scale=1,-1 ${SOURCES[1]}')
Flow('errorslet2','data2 deblendedslet2','add scale=1,-1 ${SOURCES[1]}')
Flow('errorfxdecon1','data1 deblendedfxdecon1','add scale=1,-1 ${SOURCES[1]}')
Flow('errorfxdecon2','data2 deblendedfxdecon2','add scale=1,-1 ${SOURCES[1]}')

Grey('difffft1','title="" clip=%g'%clip)
Grey('difffft2','title="" clip=%g'%clip)
Grey('diffslet1','title="" clip=%g'%clip)
Grey('diffslet2','title="" clip=%g'%clip)
Grey('difffxdecon1','title="" clip=%g'%clip)
Grey('difffxdecon2','title="" clip=%g'%clip)
Grey('errorfft1','title="" clip=%g'%clip)
Grey('errorfft2','title="" clip=%g'%clip)
Grey('errorslet1','title="" clip=%g '%clip)
Grey('errorslet2','title="" clip=%g '%clip)
Grey('errorfxdecon1','title="" clip=%g'%clip)
Grey('errorfxdecon2','title="" clip=%g'%clip)

#Grey('deblendedfft1','title="Deblended 1 (fft)"clip=%g'%clip)
#Grey('deblendedfft2','title="Deblended 1 (fft)"clip=%g'%clip)
#Grey('deblendedslet1','title="Deblended 1 (Seislet)" clip=%g'%clip)
#Grey('deblendedslet2','title="Deblended 1 (Seislet)" clip=%g'%clip)
Grey('deblendedfft1',' clip=%g'%clip)
Grey('deblendedfft2',' clip=%g'%clip)
Grey('deblendedslet1',' clip=%g'%clip)
Grey('deblendedslet2',' clip=%g'%clip)
Grey('deblendedfxdecon1',' clip=%g'%clip)
Grey('deblendedfxdecon2',' clip=%g'%clip)

## Ploting
Flow('snrsa','datasfft-snrsa datasslet-snrsa datasfxdecon-snrsa','cat axis=2 ${SOURCES[1:3]}')
Flow('snrsb','udatasfft-snrsb udatasslet-snrsb udatasfxdecon-snrsb','cat axis=2 ${SOURCES[1:3]}')

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


#Grey('datasfftsigtemp1','scalebar=y')
#Flow('comp1','udatasfft shottime1 shottime2 datasfft','blend shot_time_in=${SOURCES[2]} #shot_time_out=${SOURCES[1]} | add scale=1,1 ${SOURCES[3]}')
#Flow('comp2','datasfft shottime2 shottime1 udatasfft','blend shot_time_in=${SOURCES[2]} #shot_time_out=${SOURCES[1]} | add scale=1,1 ${SOURCES[3]}')
#Flow('comp','comp1 comp2','cat axis=2 ${SOURCES[1]}')
#Grey('comp','scalebar=y')
#Flow('diff','datasfftsigtemp1 comp','add scale=1,-1 ${SOURCES[1]}')
#Grey('diff','scalebar=y')

Greyplot('datas',' title="Iter # = %g"'%(0))
Greyplot('udatas',' title="Iter # = %g"'%(0))
deblendffts1=['datas']
deblendffts2=['udatas']
deblendslets1=['datas']
deblendslets2=['udatas']
deblendfxdecons1=['datas']
deblendfxdecons2=['udatas']

for i in range(niter):
        deblendfft1='datasfft%gp%g'%(i+1,i+1)
        deblendfft2='udatasfft%gp%g'%(i+1,i+1)
        deblendslet1='datasslet%gp%g'%(i+1,i+1)
        deblendslet2='udatasslet%gp%g'%(i+1,i+1)
        deblendfxdecon1='datasfxdecon%gp%g'%(i+1,i+1)
        deblendfxdecon2='udatasfxdecon%gp%g'%(i+1,i+1)
        deblendffts1.append(deblendfft1)
        deblendffts2.append(deblendfft2)
        deblendslets1.append(deblendslet1)
        deblendslets2.append(deblendslet2)
        deblendfxdecons1.append(deblendfxdecon1)
        deblendfxdecons2.append(deblendfxdecon2)

        Greyplot(deblendfft1,' title="Iter # = %g"'%(i+1))
        Greyplot(deblendfft2,' title="Iter # = %g"'%(i+1))
        Greyplot(deblendslet1,' title="Iter # = %g"'%(i+1))
        Greyplot(deblendslet2,' title="Iter # = %g"'%(i+1))
        Greyplot(deblendfxdecon1,' title="Iter # = %g"'%(i+1))
        Greyplot(deblendfxdecon2,' title="Iter # = %g"'%(i+1))
Plot('deblendfft1',deblendffts1,'Movie')
Plot('deblendfft2',deblendffts2,'Movie')
Plot('deblendslet1',deblendslets1,'Movie')
Plot('deblendslet2',deblendslets2,'Movie')
Plot('deblendfxdecon1',deblendfxdecons1,'Movie')
Plot('deblendfxdecon2',deblendfxdecons2,'Movie')


End()