from rsf.proj import *
from rsf.prog import RSFROOT

def Grey(data,other):
        Result(data,'window n1=490  | grey label2=Trace unit2="" label1=Time unit1="s" title="" labelsz=10  labelfat=4 font=2 titlesz=10 titlefat=4  wherexlabel=b wheretitle=t color=g bartype=v screenratio=1.3 %s'%other)

def Grey3(data,other):
        Result(data,'window n1=490| byte clip=0.5 | grey3 frame1=250 frame2=128 frame3=24 flat=n point1=0.8 point2=0.8 label3=Scale label2=Trace unit2="" label1=Time unit1="s" title="" labelsz=10  labelfat=4 font=2 titlesz=10 titlefat=4  wherexlabel=b wheretitle=t color=g bartype=v screenratio=1.3 %s'%other)

def Greyfk(data,other):
        Result(data,'grey label2=Trace unit2="" label1=Time unit1="s" title="" labelsz=10  labelfat=4 font=2 titlesz=10 titlefat=4  wherexlabel=b wheretitle=t color=g bartype=v screenratio=1.3 clip=83 %s'%other)


#############################################################################
# Setting parameters
########################################################################
n1=512
n2=256
lambda1=0.5
niter=30
lvl=2
htype='spline'
thr1=5  #slet alone
thr2=2  #dsd            
thr3=4 #ddtf alone
thrtype='g'

###########################################################################
## four complex events (one horizontal, three dipping with intersection)
###########################################################################
Flow('conflict-0',None,
     '''
     spike n1=512 n2=256 d2=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 | scale axis=2
     ''')
# Add some static

Flow('static','conflict-0','window n1=1 | noise seed=2015 var=0.005 | smooth rect1=10 repeat=2')

Flow('conflict','conflict-0 static','datstretch datum=${SOURCES[1]}')

Flow('conflictn','conflict','noise var=0.02 seed=201314')
Grey('conflict','title=Clean')
Grey('conflictn','title=')


Flow('dip1','conflictn','threshold1 thr=10 | dip rect1=5 rect2=5 order=2')
Flow('dip','conflictn','dip rect1=10 rect2=10 order=2')

# seislet transform
Flow('conflictn-slet','conflictn dip','seislet order=2 dip=${SOURCES[1]} adj=y inv=y eps=0.1 type=b')

# threshold
Flow('conflictn-sletthr','conflictn-slet','threshold1 thr=%g'%thr1)

# transform back
Flow('conflictn-re','conflictn-sletthr dip','seislet order=2 dip=${SOURCES[1]} adj=n iv=y eps=0.1 type=b' )

import random
random.seed(2014)

nr = 0
def rnd(x):
    global nr
    r = str(random.randint(1,nr))
    return r

nsp = 200 # number of basis functions
n1 = 512
n2 = 256

nr = n1
k1 = string.join(map(rnd,range(nsp)),',')
nr = n2
k2 = string.join(map(rnd,range(nsp)),',')

nr = 12544
nsp=10000
k3 = string.join(map(rnd,range(nsp)),',')

#Flow('seislets','conflictn-slet dip',
#     '''
#     spike nsp=%d k1=%s k2=%s | mul ${SOURCES[0]} |
#     seislet order=2 dip=${SOURCES[1]} adj=n inv=y eps=0.1 type=b
#     ''' % (nsp,k1,k2))


N=150
thr0=N/512.0/256*100
thr1=N/512.0/256/49*100


Flow('seislets','conflictn-slet dip',
     '''
     threshold1 type=h thr=%g | 
     seislet order=2 dip=${SOURCES[1]} adj=n inv=y eps=0.1 type=b
     '''%thr0)

#Grey('seislets','title="Seislet basis" clip=0.0005') 
Flow('seisletsr','seislets','cp')
Grey('seisletsr','title="Seislet reconstruction" scalebar=y ')

Flow('conflictn-dif','conflictn conflictn-re','add scale=1,-1 ${SOURCES[1]}')

########################################################################
# INITIALIZATION
########################################################################
matlab         = WhereIs('matlab')
matROOT = '../Matfun'
matfun = 'Syn_ddtf'
matfun1 = 'Syn_ddtf_dec'
matfun2 = 'Syn_ddtf_rec'
matlabpath = os.environ.get('MATLABPATH',os.path.join(RSFROOT,'lib'))

if not matlab:
    sys.stderr.write('\nCannot find Matlab.\n')
    sys.exit(1)

############################################################
## generate and process synthetic data
############################################################
Flow('conflictn-sletddtfthr',[os.path.join(matROOT,matfun+'.m'),'conflictn-slet'],
     '''MATLABPATH=%(matlabpath)s %(matlab)s 
     -nosplash -nojvm -r "addpath %(matROOT)s;%(matfun)s('${SOURCES[1]}','${TARGETS[0]}',%(n1)d,%(n2)d,%(lambda1)g,%(niter)d,%(lvl)d,'%(htype)s',%(thr2)g,'%(thrtype)s');quit"
     '''%vars(),stdin=0,stdout=-1)
Flow('conflictn-sletddtfthr0','conflictn-sletddtfthr','put d2=1 o2=0 d1=0.004 o1=0')

Flow('conflictn-sletddtf-t',[os.path.join(matROOT,matfun1+'.m'),'conflictn-slet'],
     '''MATLABPATH=%(matlabpath)s %(matlab)s 
     -nosplash -nojvm -r "addpath %(matROOT)s;%(matfun1)s('${SOURCES[1]}','${TARGETS[0]}',%(n1)d,%(n2)d,%(lambda1)g,%(niter)d,%(lvl)d,'%(htype)s');quit"
     '''%vars(),stdin=0,stdout=-1)
Flow('conflictn-sletddtf','conflictn-sletddtf-t','put d2=1 o2=0 d1=0.004 o1=0 d3=1 n2=256 n3=49')

Flow('conflictn-sletddtf-t-thr','conflictn-sletddtf-t','threshold1 type=s ifperc=1 thr=2')
Flow('conflictn-sletddtf-thr-recon-t',[os.path.join(matROOT,matfun2+'.m'),'conflictn-sletddtf-t-thr', 'conflictn-slet'],
     '''MATLABPATH=%(matlabpath)s %(matlab)s 
     -nosplash -nojvm -r "addpath %(matROOT)s;%(matfun2)s('${SOURCES[1]}','${SOURCES[2]}','${TARGETS[0]}',%(n1)d,%(n2)d,%(lambda1)g,%(niter)d,%(lvl)d,'%(htype)s');quit"
     '''%vars(),stdin=0,stdout=-1)
Flow('conflictn-sletddtf-thr','conflictn-sletddtf-t-thr','put d2=1 o2=0 d1=0.004 o1=0 d3=1 n2=256 n3=49')
Flow('conflictn-sletddtf-thr-recon','conflictn-sletddtf-thr-recon-t','put d2=1 o2=0 d1=0.004 o1=0')

Flow('conflictn-sletddtf-recon-t',[os.path.join(matROOT,matfun2+'.m'),'conflictn-sletddtf-t', 'conflictn-slet'],
     '''MATLABPATH=%(matlabpath)s %(matlab)s 
     -nosplash -nojvm -r "addpath %(matROOT)s;%(matfun2)s('${SOURCES[1]}','${SOURCES[2]}','${TARGETS[0]}',%(n1)d,%(n2)d,%(lambda1)g,%(niter)d,%(lvl)d,'%(htype)s');quit"
     '''%vars(),stdin=0,stdout=-1)
Flow('conflictn-sletddtf-recon','conflictn-sletddtf-recon-t','put d2=1 o2=0 d1=0.004 o1=0')




#Flow('conflictn-sletddtf-ts','conflictn-sletddtf-t',
#     '''
#     spike nsp=%d k1=%s k2=%s | mul ${SOURCES[0]}
#     ''' % (nsp,k1,k3))


### Basis functions of DSD
Flow('conflictn-sletddtf-ts','conflictn-sletddtf-t',
     '''
     threshold1 type=h thr=%g
     '''%thr1)

Flow('conflictn-sletddtf-recon-ts',[os.path.join(matROOT,matfun2+'.m'),'conflictn-sletddtf-ts', 'conflictn-slet'],
     '''MATLABPATH=%(matlabpath)s %(matlab)s 
     -nosplash -nojvm -r "addpath %(matROOT)s;%(matfun2)s('${SOURCES[1]}','${SOURCES[2]}','${TARGETS[0]}',%(n1)d,%(n2)d,%(lambda1)g,%(niter)d,%(lvl)d,'%(htype)s');quit"
     '''%vars(),stdin=0,stdout=-1)
Flow('conflictn-sletddtf-recons','conflictn-sletddtf-recon-ts','put d2=1 o2=0 d1=0.004 o1=0')

Flow('dsds','conflictn-sletddtf-recons dip',
        '''
     seislet order=2 dip=${SOURCES[1]} adj=n inv=y eps=0.1 type=b
     ''')
#Grey('dsds','title="DSD basis" clip=0.0005')    
Flow('dsdsr','dsds','cp')
Grey('dsdsr','title="DSD reconstruction" scalebar=y')

### Basis functions of DDTF
Flow('conflictn-ddtf-t',[os.path.join(matROOT,matfun1+'.m'),'conflictn'],
     '''MATLABPATH=%(matlabpath)s %(matlab)s 
     -nosplash -nojvm -r "addpath %(matROOT)s;%(matfun1)s('${SOURCES[1]}','${TARGETS[0]}',%(n1)d,%(n2)d,%(lambda1)g,%(niter)d,%(lvl)d,'%(htype)s');quit"
     '''%vars(),stdin=0,stdout=-1)

Flow('conflictn-ddtf-ts','conflictn-ddtf-t',
     '''
     threshold1 type=h thr=%g  
     '''%thr1)

Flow('conflictn-ddtf-recon-ts',[os.path.join(matROOT,matfun2+'.m'),'conflictn-ddtf-ts', 'conflictn'],
     '''MATLABPATH=%(matlabpath)s %(matlab)s 
     -nosplash -nojvm -r "addpath %(matROOT)s;%(matfun2)s('${SOURCES[1]}','${SOURCES[2]}','${TARGETS[0]}',%(n1)d,%(n2)d,%(lambda1)g,%(niter)d,%(lvl)d,'%(htype)s');quit"
     '''%vars(),stdin=0,stdout=-1)
Flow('ddtfs','conflictn-ddtf-recon-ts','put d2=1 o2=0 d1=0.004 o1=0')


#Grey('ddtfs','title="DDTF basis" clip=0.0005')  
Flow('ddtfsr','ddtfs','cp')
Grey('ddtfsr','title="DDTF reconstruction" scalebar=y')



Flow('dsd-denoise','conflictn-sletddtf-thr-recon dip',
        '''
     seislet order=2 dip=${SOURCES[1]} adj=n inv=y eps=0.1 type=b
     ''')
Grey('dsd-denoise','')
Flow('dsd-noise','conflictn dsd-denoise','add scale=1,-1 ${SOURCES[1]}')
Grey('dsd-noise','')
## Seislet + DDTF
Flow('conflictn-ddtf-re','conflictn-sletddtfthr dip','put d2=1 o2=0 d1=0.004 o1=0 | seislet order=2 dip=${SOURCES[1]} adj=n inv=y eps=0.1 type=b')
Flow('conflictn-ddtf-dif','conflictn conflictn-ddtf-re','add scale=1,-1 ${SOURCES[1]}')



## DDTF alone
Flow('conflictn-ddtf0',[os.path.join(matROOT,matfun+'.m'),'conflictn'],
     '''MATLABPATH=%(matlabpath)s %(matlab)s 
     -nosplash -nojvm -r "addpath %(matROOT)s;%(matfun)s('${SOURCES[1]}','${TARGETS[0]}',%(n1)d,%(n2)d,%(lambda1)g,%(niter)d,%(lvl)d,'%(htype)s',%(thr3)g,'%(thrtype)s');quit"
     '''%vars(),stdin=0,stdout=-1)
Flow('conflictn-ddtf','conflictn-ddtf0','put d2=1 o2=0 d1=0.004 o1=0')
Flow('conflictn-ddtf-dif0','conflictn conflictn-ddtf','add scale=1,-1 ${SOURCES[1]}')
Grey('conflictn-ddtf','title="DDTF"')
Grey('conflictn-ddtf-dif0','title="DDTF"')

## Ploting
Grey('dip','color=j')
Grey('dip1','color=j')
Grey('conflictn-slet','clip=0.5 label2=Scale')
Grey('conflictn-sletddtf-recon','clip=0.5 label2=Scale')
Grey('conflictn-sletddtf-thr-recon','clip=0.5 label2=Scale')

Grey('conflictn-sletthr','clip=0.5')
Grey('conflictn-sletddtfthr0','clip=0.5')
Grey('conflictn-re','title="Seislet"')
Grey('conflictn-dif','title="Seislet"')

Grey('conflictn-ddtf-re','title="DSD"')
Grey('conflictn-ddtf-dif','title="DSD"')


## Error section

Flow('conflictn-err','conflict conflictn-re','add scale=1,-1 ${SOURCES[1]}')
Flow('conflictn-ddtf-err','conflict conflictn-ddtf','add scale=1,-1 ${SOURCES[1]}')
Flow('conflictn-ddtf-re-err','conflict conflictn-ddtf-re','add scale=1,-1 ${SOURCES[1]}')

Grey('conflictn-ddtf-re-err','title="DSD" clip=1.2')
Grey('conflictn-ddtf-err','title="DDTF" clip=1.2')
Grey('conflictn-err','title="Seislet" clip=1.2')
## FX decon
Flow('conflictn-fx','conflictn','fxdecon n2w=%d '%n2)
Flow('conflictn-fx-dif','conflictn conflictn-fx','add scale=1,-1 ${SOURCES[1]}')
Grey('conflictn-fx','title="FX"')
Grey('conflictn-fx-dif','title="FX"')


## FK
Flow('conflict-fk','conflict','fft1 | fft3 axis=2 | cabs')
Flow('conflictn-fk','conflictn','fft1 | fft3 axis=2 | cabs')
Flow('conflictn-re-fk','conflictn-re','fft1 | fft3 axis=2 | cabs')
Flow('conflictn-ddtf-fk','conflictn-ddtf','fft1 | fft3 axis=2 | cabs')
Flow('conflictn-ddtf-re-fk','conflictn-ddtf-re','fft1 | fft3 axis=2| cabs')

Greyfk('conflict-fk','title="Clean" color=j allpos=y')
Greyfk('conflictn-fk','title="Noisy" color=j allpos=y')
Greyfk('conflictn-re-fk','title="Seislet" color=j allpos=y')
Greyfk('conflictn-ddtf-fk','title="DDTF" color=j allpos=y')
Greyfk('conflictn-ddtf-re-fk','title="DSD" color=j allpos=y')


Grey3('conflictn-sletddtf','')
Grey3('conflictn-sletddtf-thr','')



End()