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

clip=0.0005             #display percentage
padno=1024      #padding for seislet tranform
fhi=60          #bandpass frequency value
r1=20           #smoothing radius
r2=20           #smoothing radius
thr1=50         #thresholding percentage
thr2=18         #thresholding percentage
n1=3001         #temporal sampling number
n2=560          #spatial sampling number
win = 560       #spatial processing window for fxdecon
niter=30        #iteration numbers
ddip=5          #changing dip map interval
fraction=0.5    # B=fraction*I
mute1=   '''mutter t0=3.2 slope0=-0.006 |
         mutter x0=250 slope0=0.006  |   
         cut f1=2001'''
mute2=   '''mutter t0=3.2 slope0=-0.006 |
         mutter x0=250 slope0=0.006  |   
         cut f1=2001'''
mode = 'soft'

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

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

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

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

## Fetch data from server
for file in [
          'Z700_SimSou_754-UNBLENDED-',
          'Z700_SimSou_763-UNBLENDED-',
         # 'Z700_SimSou_754-BLENDED-Close-Synch-1000dither',
         # 'Z700_SimSou_763-BLENDED-Close-Synch-1000dither',
             ]:
         Fetch(file+'.segy','fairfield',server)
         if file == 'Z700_SimSou_754-UNBLENDED-':
                file1='unblended754'
         elif file=='Z700_SimSou_763-UNBLENDED-':
                file1='unblended763'
         #elif file=='Z700_SimSou_754-BLENDED-Close-Synch-1000dither':
         #      file1='blended754'
         #elif file=='Z700_SimSou_763-BLENDED-Close-Synch-1000dither':
         #      file1='blended763'
         Flow([file1,'t'+file1,file1+'.asc',file1+'.bin'],
                file+'.segy',
                '''
                segyread
                tfile=${TARGETS[1]}
                hfile=${TARGETS[2]}
                bfile=${TARGETS[3]} | scale axis=12
                ''')
         Grey(file1,' title=%s clip=%g'%('',clip))


## extract shottime#.rsf from data header file tfile#
for no in (754, 763):
        Flow('shottime%g'%no,'tunblended%g'%no,
        '''
        dd type=float
        | headermath output="cdpx"
        ''')
## Blend 
Flow('blended754','unblended763 unblended754 shottime754 shottime763','blend shot_time_in=${SOURCES[3]} shot_time_out=${SOURCES[2]} |add scale=1,1 ${SOURCES[1]}' )
Flow('blended763','unblended754 unblended763 shottime754 shottime763','blend shot_time_in=${SOURCES[2]} shot_time_out=${SOURCES[3]} |add scale=1,1 ${SOURCES[1]}' )

Grey('blended754',' title=%s clip=%g'%('',clip))
Grey('blended763',' title=%s clip=%g'%('',clip))


## mutte the data
Flow('orig','unblended754','mutter t0=3.2 slope0=-0.006 |mutter x0=250 slope0=0.006  |cut f1=2001')
Flow('uorig','unblended763','mutter t0=3.2 slope0=-0.006 |mutter x0=250 slope0=0.006  |cut f1=2001')

Flow('shots','blended754','mutter t0=3.2 slope0=-0.006 |mutter x0=250 slope0=0.006  |cut f1=2001 | transp | mf nfw=3 |transp')
Flow('ushots','blended763','mutter t0=3.2 slope0=-0.006 |mutter x0=250 slope0=0.006  |cut f1=2001 |transp | mf nfw=3 |transp')

Result('orig','grey title=s1 clip=%g max1=8'%clip)
Result('uorig','grey title=s2 clip=%g max1=8'%clip)
Result('shots','grey title=d1 clip=%g max1=8'%clip)
Result('ushots','grey title=d2 clip=%g max1=8'%clip)

Result('mute','shots blended754','add scale=-1,1 ${SOURCES[1]} | grey title=mute1 clip=%g max1=8'%clip)
Result('umute','ushots blended763','add scale=-1,1 ${SOURCES[1]} | grey title=mute2 clip=%g max1=8'%clip)

Flow('shotsfft','shots','math output="input/2"')
Flow('ushotsfft','ushots','math output="input/2"')
Flow('shotsslet','shots','math output="input/2"')
Flow('ushotsslet','ushots','math output="input/2"')
Flow('shotsfxdecon','shots','math output="input/2"')
Flow('ushotsfxdecon','ushots','math output="input/2"')

## fk transform and filtering
Flow('shotsfka','shots','fft1 | fft3 axis=2 pad=1 | cabs')
Flow('ushotsfka','ushots','fft1 | fft3 axis=2 pad=1 | cabs')
Flow('shot1fka','unblended754','fft1 | fft3 axis=2 pad=1 | cabs')
Flow('shot2fka','unblended763','fft1 | fft3 axis=2 pad=1 | cabs')
Flow('shotsfka-mute','shots','fft1 | fft3 axis=2 pad=1 |cabs| mutter half=n slope0=60 | cut max1=5')
Flow('ushotsfka-mute','ushots','fft1 | fft3 axis=2 pad=1 |cabs')

Flow('shotsfkr','shots','fft1 | fft3 axis=2 pad=1 | real')
Flow('shotsfki','shots','fft1 | fft3 axis=2 pad=1 | imag')
Flow('shotsfk','shots','fft1 | fft3 axis=2 pad=1')
Flow('shotsfkr_filt','shotsfkr','mutter half=n t0=0 slope0=40 x0=0 ')
Flow('shotsfki_filt','shotsfki','mutter half=n t0=0 slope0=40 x0=0 ')

Grey('ushotsfka',' label1=Frequency unit1=Hz label2=Wavenumber color=j')
Grey('shot1fka',' label1=Frequency unit1=Hz label2=Wavenumber color=j')
Grey('shot2fka',' label1=Frequency unit1=Hz label2=Wavenumber color=j')
Grey('shotsfka-mute',' label1=Frequency unit1=Hz label2=Wavenumber color=j')
Grey('ushotsfka-mute',' label1=Frequency unit1=Hz label2=Wavenumber color=j')
Grey('shotsfka',' label1=Frequency unit1=Hz label2=Wavenumber color=j')
Grey('shotsfkr',' label1=Frequency unit1=Hz label2=Wavenumber color=j')
Grey('shotsfki',' label1=Frequency unit1=Hz label2=Wavenumber color=j')
Grey('shotsfkr_filt',' label1=Frequency unit1=Hz label2=Wavenumber color=j')
Grey('shotsfki_filt',' label1=Frequency unit1=Hz label2=Wavenumber color=j')

#deblend using fft
deblend2.deblendfft('orig',
          'uorig',
          'blended754',
          'blended763',
          'shotsfft',
          'ushotsfft',
          'shotdeblendedfft1',
          'shotdeblendedfft2',
          'shottime754',
          'shottime763',
           mute1,
           mute2,
           n1,
           n2,
           niter,
           mode,
           thr1,
           clip,
           fraction)

#deblend using seislet
deblend2.deblendslet('orig',
          'uorig',
          'blended754',
          'blended763',
          'shotsslet',
          'ushotsslet',
          'shotdeblendedslet1',
          'shotdeblendedslet2',
          'shottime754',
          'shottime763',
           mute1,
           mute2,
           n1,
           n2,
           r1,
           r2,
           padno,
           niter,
           ddip,
           mode,
           thr2,
           clip,
           fhi,
           fraction)

#deblend using fxdecon
deblend2.deblendfxdecon('orig',
          'uorig',
          'blended754',
          'blended763',
          'shotsfxdecon',
          'ushotsfxdecon',
          'shotdeblendedfxdecon1',
          'shotdeblendedfxdecon2',
          'shottime754',
          'shottime763',
           mute1,
           mute2,
           n1,
           n2,
           niter,
           clip,
           fraction)


## Ploting difference, error, deblended sections
Flow('shotdifffft1','blended754 shotdeblendedfft1','add scale=1,-1 ${SOURCES[1]}')
Flow('shotdifffft2','blended763 shotdeblendedfft2','add scale=1,-1 ${SOURCES[1]}')
Flow('shotdiffslet1','blended754 shotdeblendedslet1','add scale=1,-1 ${SOURCES[1]}')
Flow('shotdiffslet2','blended763 shotdeblendedslet2','add scale=1,-1 ${SOURCES[1]}')
Flow('shotdifffxdecon1','blended754 shotdeblendedfxdecon1','add scale=1,-1 ${SOURCES[1]}')
Flow('shotdifffxdecon2','blended763 shotdeblendedfxdecon2','add scale=1,-1 ${SOURCES[1]}')

Flow('shoterrorfft1','unblended754 shotdeblendedfft1','add scale=1,-1 ${SOURCES[1]}')
Flow('shoterrorfft2','unblended763 shotdeblendedfft2','add scale=1,-1 ${SOURCES[1]}')
Flow('shoterrorslet1','unblended754 shotdeblendedslet1','add scale=1,-1 ${SOURCES[1]}')
Flow('shoterrorslet2','unblended763 shotdeblendedslet2','add scale=1,-1 ${SOURCES[1]}')
Flow('shoterrorfxdecon1','unblended754 shotdeblendedslet1','add scale=1,-1 ${SOURCES[1]}')
Flow('shoterrorfxdecon2','unblended763 shotdeblendedslet2','add scale=1,-1 ${SOURCES[1]}')

Grey('shotdifffft1','title="" clip=%g'%clip)
Grey('shotdifffft2','title="" clip=%g'%clip)
Grey('shotdiffslet1','title="" clip=%g'%clip)
Grey('shotdiffslet2','title="" clip=%g'%clip)
Grey('shotdifffxdecon1','title="" clip=%g'%clip)
Grey('shotdifffxdecon2','title="" clip=%g'%clip)

Grey('shoterrorfft1','title="" clip=%g'%clip)
Grey('shoterrorfft2','title="" clip=%g'%clip)
Grey('shoterrorslet1','title="" clip=%g'%clip)
Grey('shoterrorslet2','title="" clip=%g'%clip)
Grey('shoterrorfxdecon1','title="" clip=%g'%clip)
Grey('shoterrorfxdecon2','title="" clip=%g'%clip)
#Grey('shotdeblendedfft1','title="Deblended 1 (fft)"clip=%g'%clip)
#Grey('shotdeblendedfft2','title="Deblended 1 (fft)"clip=%g'%clip)
#Grey('shotdeblendedslet1','title="Deblended 1 (Seislet)" clip=%g'%clip)
#Grey('shotdeblendedslet2','title="Deblended 1 (Seislet)" clip=%g'%clip)
Grey('shotdeblendedfft1',' clip=%g'%clip)
Grey('shotdeblendedfft2',' clip=%g'%clip)
Grey('shotdeblendedslet1',' clip=%g'%clip)
Grey('shotdeblendedslet2',' clip=%g'%clip)
Grey('shotdeblendedfxdecon1',' clip=%g'%clip)
Grey('shotdeblendedfxdecon2',' clip=%g'%clip)

## Ploting
Flow('shotsnrsa','shotsfft-snrsa shotsslet-snrsa shotsfxdecon-snrsa','cat axis=2 ${SOURCES[1:3]}')
Flow('shotsnrsb','ushotsfft-snrsb ushotsslet-snrsb ushotsfxdecon-snrsb','cat axis=2 ${SOURCES[1:3]}')

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