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

# module defining


def Grey(hyper, other):
    Result(hyper,
           '''        
          grey  font=2 labelsz=6 titlefat=4 title= wheretitle=t
       labelfat=4 screenratio=1.4 clip=0.13167 wherexlabel=b
       label2="Offset" unit2=m 
       %s ''' % other)


def Greydip(hyper1, hyper2, other):
    Result(hyper1, hyper2,
           '''
       grey font=2 titlefat=4 labelfat=4 barwidth=0.1
       label2=Offset unit2=m barlabel=Slope barunit=samples
       allpos=y scalebar=y clip=2.0 maxval=2 color=j title= wheretitle=t       
       wherexlabel=b screenratio=1.4 %s
       ''' % other)


def Greyfk(data, other):
    Result(data, 'grey label2=Trace  unit2="" clip=0.8 label1=Time unit1="s" title="" wherexlabel=b wanttitle=y wheretitle=t screenratio=1.4 %s ' % other)


def Greyplot(data, other):
    Plot(data,
         '''
                grey font=2 titlefat=4 labelfat=4 barwidth=0.1
       label2=Offset unit2=m barlabel=Slope barunit=samples
       allpos=y scalebar=y clip=2.0 maxval=2 color=j title= wheretitle=b
       screenratio=1.4 %s
       ''' % other)


def Graph(data, other):
    Result(data, 'graph label1="Iter #no" label2="SNR" unit2=dB unit1="" title="" wherexlabel=b wheretitle=t %s' % other)


# setup model
Flow('vrms', None, 'math d1=0.004 n1=1001 o1=0 output="x1*125+2000" ')
Flow('vint', 'vrms', 'math output="125*sqrt((16+x1)*(16+3*x1))" ')

for vel in ('vrms', 'vint'):
    Plot(vel,
         '''
         graph transp=y yreverse=y min2=1995 max2=3005
         pad=n wantaxis=n wanttitle=n plotcol=5
         ''')
for vel in ('vrms', 'vint'):
    Plot(vel+'1', vel,
         '''
         graph transp=y yreverse=y min2=1995 max2=3005
         pad=n wantaxis=n wanttitle=n plotcol=5
         screenratio=1.4 screenht=8 labelsz=6 plotfat=4 plotcol=7 dash=3
         ''')

Flow('hyper', 'vrms',
     '''
     spike d1=0.004 n1=1001 o1=0 nsp=17 n2=128 d2=20 o2=0
     label2=Offset unit2=m
     k1=%s
     mag=1,1,1,1,-1,1,1,-1,1,1,1,-1,1,1,-1,1,1 |
     bandpass flo=4 fhi=20 |
     inmo velocity=$SOURCE half=n 
     ''' % ','.join(map(str, range(100, 916, 48))), stdin=0)


Flow('hyper0', 'hyper', 'pad beg2=1 | rmtrace factor=2 | noise seed=201617 var=0.0005')
Flow('hyper-zero', 'hyper0', 'addtrace ratio=2')


Flow('idip', 'hyper',
     '''
     math output="(0.00125*(x2+10)/x1)" |
     mutter half=n v0=10000 t0=0.004 tp=0.1
     ''')

# Riesz transform dips
Flow('riesz', 'hyper-zero', 'riesz order=10')

Flow('rx', 'riesz', 'window n3=1')
Flow('ry', 'riesz', 'window f3=1 | scale dscale=-1')
Plot('rx', 'grey title=Rx')
Plot('ry', 'grey title=Ry')
Result('riesz', 'rx ry', 'SideBySideIso')

Flow('rdip', 'ry rx', 'divn den=${SOURCES[1]} rect1=5 rect2=5')
Result('rdip',
       'grey color=j scalebar=y title="Dip from Riesz Transform"')


# PWD dips
Flow('sdip', 'hyper idip',
     '''
     dip order=3 liter=100 pmin=0
     idip=${SOURCES[1]} niter=10 rect1=40 rect2=5
     ''')
Plot('sdip',
     '''
     grey color=j scalebar=y title="Ideal dip"
     font=2 titlefat=4 labelfat=4
     allpos=y scalebar=y clip=2 maxval=2 color=j
     ''')
Result('sdip',
       '''
       grey title="PWD slope" font=2 titlefat=4 labelfat=4 barwidth=0.1
       allpos=y scalebar=y clip=2 maxval=2 color=I wanttitle=n
       barlabel=Slope barunit=samples
       wanttitle=n screenratio=1.7 screenht=8 labelsz=6
       ''')

Flow('pdip', 'hyper-zero',
     '''
     bandpass fhi=15 | dip order=3 liter=100
      niter=10 rect1=10 rect2=10
     ''')
Plot('pdip',
     '''
     grey title="PWD slope" font=2 titlefat=4 labelfat=4
     allpos=y scalebar=y clip=2 maxval=2 color=j wanttitle=n
     ''')
Result('pdip',
       '''
       grey title="PWD slope" font=2 titlefat=4 labelfat=4 barwidth=0.1
       allpos=y scalebar=y clip=2 maxval=2 color=I wanttitle=n
       barlabel=Slope barunit=samples
       wanttitle=n screenratio=1.7 screenht=8 labelsz=6       
       ''')


# velocity dependent (VD) dips
Flow('svsc', 'hyper-zero', 'vscan semblance=y v0=2000 dv=10 nv=101 half=n')

Plot('svsc',
     '''
     envelope |
     grey allpos=y label2=Velocity unit2=m/s labelfat=4
     title="Velocity Scan" font=2 titlefat=4 wanttitle=n
     ''')

Flow('svsc-t', 'hyper', 'vscan semblance=y v0=2000 dv=10 nv=101 half=n')

Plot('svsc-t',
     '''
     envelope |
     grey allpos=y label2=Velocity unit2=m/s labelfat=4
     title="Velocity Scan" font=2 titlefat=4 wanttitle=n
     ''')

Plot('svsc1', 'svsc',
     '''
     envelope |
     grey allpos=y label2=Velocity unit2=m/s labelfat=4
     title="Velocity Scan" font=2 titlefat=4 wanttitle=n
     screenratio=1.55 screenht=8 labelsz=6 color=I
     ''')

Flow('vpick', 'svsc', 'scale axis=2 | pick rect1=40 | window')

Plot('vpick',
     '''
     graph transp=y yreverse=y min2=1995 max2=3005
     pad=n wantaxis=n wanttitle=n plotcol=3
     ''')
Plot('vpick1', 'vpick',
     '''
     graph transp=y yreverse=y min2=1995 max2=3005
     pad=n wantaxis=n wanttitle=n plotcol=3
     screenratio=1.55 screenht=8 labelsz=6 plotfat=4 plotcol=7 dash=0
     ''')

Plot('svsc2', 'svsc vrms vpick', 'Overlay')
Result('svsc2', 'svsc1 vrms1 vpick1', 'Overlay')

### Test Wrong convert equation (should be v(t0) but not v(t))###
Flow('vdip1', 'vpick',
     '''
     spray axis=2 n=128 d=20 o=0 |
     math output="x2*20./(x1*input*input*0.004+0.00001)" |
     mutter half=n v0=10000 t0=0.004 tp=0.1
     ''')
#################################################################

Flow('vdip', 'vpick',
     '''
     v2d n=128 d=20 o=0 mute=y half=n v0=10000 t0=0.004 tp=0.1
     ''')


Grey('hyper', 'title="Original"')
Grey('hyper0', 'title="Under-sampled"')
Grey('hyper-zero', 'title="Zero-padded"')


Greydip('hyper-sdip', 'sdip', 'title="True"')
Greydip('hyper-pdip', 'pdip', 'title="PWD"')
Greydip('hyper-rdip', 'rdip', 'title="Riesz"')
Greydip('hyper-vdip', 'vdip', 'title="Velocity->Dip"')


# Create mask
Flow('mask', 'hyper0', 'math output="1" | addtrace ratio=2')
Flow('mask1', 'mask', 'math output="1-input"')
Grey('mask', 'color=j')


Flow('hyper-n', 'hyper mask1 hyper-zero-n',
     'noise seed=201617 var=0.0005 |math b=${SOURCES[1]} c=${SOURCES[2]} output="input*b+c"')
Flow('hyper0-n', 'hyper0', 'cp')
Flow('hyper-zero-n', 'hyper-zero', 'cp')

Grey('hyper-n', 'title="Noisy"')
Grey('hyper0-n', 'title="Under-sampled"')
Grey('hyper-zero-n', 'title="Zero-padded"')


Flow('fk-hyper-zero', 'hyper-zero',
     'put d2=1 o2=0 | rtoc | fft3 axis=1 pad=2 | fft3 axis=2 pad=2| cabs | window f1=1001')
Flow('fk-hyper', 'hyper',
     'put d2=1 o2=0 | rtoc | fft3 axis=1 pad=2 | fft3 axis=2 pad=2| cabs | window f1=1001')
Flow('fk-hyper0', 'hyper0',
     'put d2=1 o2=0 |rtoc | fft3 axis=1 pad=2 | fft3 axis=2 pad=2| cabs | window f1=1001')
Greyfk('fk-hyper-zero', 'allpos=y color=j clip=100 label1=Frequency unit1=Hz label2=Wavenumber unit2= min2=-0.5 max2=0.5 title="Zero-padded"')
Greyfk('fk-hyper', 'allpos=y color=j clip=100 label1=Frequency unit1=Hz label2=Wavenumber unit2= min2=-0.5 max2=0.5 title="Original"')
Greyfk('fk-hyper0', 'allpos=y color=j clip=100 label1=Frequency unit1=Hz label2=Wavenumber unit2= min2=-0.5 max2=0.5 title="Under-sampled"')

# parameters
ddip = 3
fhi = 20
r1 = 10
r2 = 3
padno = 128
thr0 = 3
niter = 100
mode = 's'
# POCS (thresholding in the seislet domain)
## define forward and backward seislet transform strings####
forw = 'seislet dip=${SOURCES[1]} adj=y inv=y eps=0.1 type=b'
back = 'seislet dip=${SOURCES[1]} inv=y eps=0.1 type=b'

sig = 'hyper-zero'
data_pocs = sig
plots_pocs = [sig]
diffsa_pocs = []
diffsb_pocs = []
diffs_pocs = []
snrs_pocs = []
dips_pocs = []
datas_pocs = []
snrs_pocs = []
# Create mask for seislet transform
Flow('dipmask', 'hyper', 'math output=1 | pad n2=%d' % (padno))
Flow('dip', ['hyper', 'dipmask'],
     '''
                bandpass fhi=%d | pad n2=%d | 
                dip mask=${SOURCES[1]} rect1=%d rect2=%d liter=30
                ''' % (fhi, padno, r1, r2))
Grey('dip', 'clip=1 color=j')

for iter in range(niter):
    thr = thr0+((thr0-thr0)*iter*iter/((niter-1)*(niter-1)))
    if iter % ddip == 0:
        dip_pocs = 'dip-pocs%d' % int(iter/ddip)
        Flow(dip_pocs, [data_pocs, 'dipmask'],
             '''
                bandpass fhi=%d | pad n2=%d | 
                dip mask=${SOURCES[1]} rect1=%d rect2=%d liter=30
                ''' % (fhi, padno, r1, r2))
        dips_pocs.append(dip_pocs)
        Greyplot(dip_pocs, 'clip=2 color=j')
    old_pocs = data_pocs
    data_pocs = 'data-pocs%d' % iter
    diffa_pocs = 'diffa-pocs%d' % iter
    diffb_pocs = 'diffb-pocs%d' % iter
    diff_pocs = 'diff-pocs%d' % iter
    snr_pocs = 'snr-pocs%d' % iter
    # 1. Forward seislet
    # 2. Multiply by seislet mask
    # 3. Inverse seislet
    # 4. Multiply by space mask
    # 5. Add data outside of hole
    Flow(data_pocs, [old_pocs, dip_pocs, 'mask1', sig],
         '''
         %s | threshold1 ifperc=1 mode=%s thr=%g | 
         %s | mul ${SOURCES[2]}  | 
         add ${SOURCES[3]} | %s | threshold1 ifperc=1 mode=%s thr=%g | %s 
         ''' % (forw, mode, thr, back, forw, mode, thr, back))
    Flow(diff_pocs, ['hyper', data_pocs], 'add scale=1,-1 ${SOURCES[1]}')
    Flow(snr_pocs, ['hyper', diff_pocs], 'snr2 noise=${SOURCES[1]}')

    Greyplot(data_pocs, 'title="Iteration %d"' % (iter+1))
    datas_pocs.append(data_pocs)
    snrs_pocs.append(snr_pocs)

Flow('snrs-pocs', snrs_pocs, 'cat axis=1 ${SOURCES[1:%d]}' % (len(snrs_pocs)))
Plot('dips-pocs', dips_pocs, 'Movie')
Plot('datas-pocs', datas_pocs, 'Movie')

Flow('hyper-seis', data_pocs, 'cp')
Flow('hyper-seis-dif', 'hyper hyper-seis', 'add scale=1,-1 ${SOURCES[1]}')

Grey('hyper-seis', 'title="Traditional"')
Grey('hyper-seis-dif', 'title="Traditional"')

Flow('fk-hyper-seis', 'hyper-seis',
     'put o2=0 d2=1 | rtoc | fft3 axis=1 pad=2 | fft3 axis=2 pad=2| cabs | window f1=1001')
Greyfk('fk-hyper-seis', 'allpos=y color=j clip=100 label1=Frequency unit1=Hz label2=Wavenumber unit2= min2=-0.5 max2=0.5 title="Traditional"')

Graph('snrs-pocs', 'title="Convergence diagram"')

Greydip('hyper-dipiter', dip_pocs, 'title="Iterative"')

########################################################################
# INITIALIZATION
########################################################################
matlab = WhereIs('matlab')
matROOT = '../matfun'
matfun = 'Spitz'
matlabpath = os.environ.get('MATLABPATH', os.path.join(RSFROOT, 'lib'))

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

n1 = 1001
n2 = 64
n22 = 127
d1 = 0.004
d2 = 20
o1 = 0
o2 = 0
npef = 40
pre1 = 1
pre2 = 1
flow = 1
fhigh = 120
############################################################
# with parameter
############################################################
Flow('hyper-fx-t', [os.path.join(matROOT, matfun+'.m'), 'hyper0'],
     '''MATLABPATH=%(matlabpath)s %(matlab)s 
     -nosplash -nojvm -r "addpath %(matROOT)s;%(matfun)s('${SOURCES[1]}','${TARGETS[0]}',%(n1)d,%(n2)g,%(d1)g,%(npef)d,%(pre1)g,%(pre2)g,%(flow)g,%(fhigh)g);quit"
     ''' % vars(), stdin=0, stdout=-1)
Flow('hyper-fx', 'hyper-fx-t', 'put d1=%g d2=%g o1=%g o2=%g' % (d1, d2, o1, o2))
Flow('hyper-fx-dif', 'hyper hyper-fx',
     'window n2=%d | add scale=1,-1 ${SOURCES[1]}' % n22)

Grey('hyper-fx', 'label1=Time uni1=s title="Spitz"')
Grey('hyper-fx-dif', 'title="Spitz"')


Flow('fk-hyper-fx', 'hyper-fx',
     'put d2=1 o2=0 | rtoc | fft3 axis=1 pad=2 | fft3 axis=2 pad=2| cabs | window f1=1001')
Greyfk('fk-hyper-fx', 'allpos=y color=j clip=100 label1=Frequency unit1=Hz label2=Wavenumber unit2= min2=-0.5 max2=0.5 title="Spitz"')


###################################
# Seislet VD
###################################
thr0 = 3
niter = 15
mode = 's'

sig = 'hyper-zero'
data_pocs = sig
plots_pocs = [sig]
diffsa_pocs = []
diffsb_pocs = []
diffs_pocs = []
snrs_pocs = []
dips_pocs = []
datas_pocs = []
snrs_pocs = []
# Create mask for seislet transform

for iter in range(niter):
    thr = thr0+((thr0-thr0)*iter*iter/((niter-1)*(niter-1)))
    old_pocs = data_pocs
    data_pocs = 'data-pocsvd%d' % iter
    diffa_pocs = 'diffa-pocsvd%d' % iter
    diffb_pocs = 'diffb-pocsvd%d' % iter
    diff_pocs = 'diff-pocsvd%d' % iter
    snr_pocs = 'snr-pocsvd%d' % iter
    # 1. Forward seislet
    # 2. Multiply by seislet mask
    # 3. Inverse seislet
    # 4. Multiply by space mask
    # 5. Add data outside of hole
    Flow(data_pocs, [old_pocs, 'vdip', 'mask1', sig],
         '''
         %s | threshold1 ifperc=1 mode=%s thr=%g | 
         %s | mul ${SOURCES[2]}  | 
         add ${SOURCES[3]} | %s | threshold1 ifperc=1 mode=%s thr=%g | %s
         ''' % (forw, mode, thr, back, forw, mode, thr, back))
    Flow(diff_pocs, ['hyper', data_pocs], 'add scale=1,-1 ${SOURCES[1]}')
    Flow(snr_pocs, ['hyper', diff_pocs], 'snr2 noise=${SOURCES[1]}')

    Greyplot(data_pocs, 'title="Iteration %d"' % (iter+1))
    datas_pocs.append(data_pocs)
    snrs_pocs.append(snr_pocs)

Flow('snrs-pocsvd', snrs_pocs,
     'cat axis=1 ${SOURCES[1:%d]}' % (len(snrs_pocs)))
Plot('datas-pocsvd', datas_pocs, 'Movie')

Flow('hyper-seisvd', data_pocs, 'cp')
Flow('hyper-seisvd-dif', 'hyper hyper-seisvd', 'add scale=1,-1 ${SOURCES[1]}')

Grey('hyper-seisvd', 'title="Proposed"')
Grey('hyper-seisvd-dif', 'title="Proposed"')

Graph('snrs-pocsvd', 'title="Convergence diagram"')


Flow('fk-hyper-seisvd', 'hyper-seisvd',
     'put d2=1 o2=0 | rtoc | fft3 axis=1 pad=2 | fft3 axis=2 pad=2| cabs | window f1=1001')
Greyfk('fk-hyper-seisvd', 'allpos=y color=j clip=100 label1=Frequency unit1=Hz label2=Wavenumber unit2= min2=-0.5 max2=0.5 title="Proposed"')


Flow('hypern-seisvd', 'hyper-seisvd', 'cp')
Flow('hypern-seis', 'hyper-seis', 'cp')
Flow('hypern-fx', 'hyper-fx', 'cp')
Flow('hypern-seisvd-dif', 'hyper-seisvd-dif', 'cp')
Flow('hypern-seis-dif', 'hyper-seis-dif', 'cp')
Flow('hypern-fx-dif', 'hyper-fx-dif', 'cp')

Grey('hypern-seis', 'title="Traditional"')
Grey('hypern-seis-dif', 'title="Traditional"')

Grey('hypern-seisvd', 'title="Proposed"')
Grey('hypern-seisvd-dif', 'title="Proposed"')

Grey('hypern-fx', 'label1=Time uni1=s title="Spitz"')
Grey('hypern-fx-dif', 'title="Spitz"')


Flow('snrsn-pocsvd', 'snrs-pocsvd', 'cp')
Flow('snrsn-pocs', 'snrs-pocs', 'cp')
Graph('snrsn-pocsvd', 'title="Convergence diagram"')
Graph('snrsn-pocs', 'title="Convergence diagram"')

Flow('diff-fx', ['hyper', 'hypern-fx'],
     'window n2=127|add scale=1,-1 ${SOURCES[1]}')
Flow('snr-fx', ['hyper', 'diff-fx'], 'window n2=127|snr2 noise=${SOURCES[1]}')

End()