from rsfproj import *

import os
private = {'login':os.environ.get('BEG_LOGIN'),
           'password':os.environ.get('BEG_PASSWORD'),
           'server':os.environ.get('BEG_SERVER')}

Fetch('cmp.hh','acig',private)
Flow('cmp','cmp.hh','dd form=native | tpow tpow=2 | mutter v0=1.4 half=n')


def grey(title,other=''):
    return '''
    grey title="%s" labelsz=10 titlesz=12
    label1="Time (s)" label2="Offset (km)" %s
    ''' % (title,other)

Plot('cmp',grey('(a) Data','clip=228'))

# Playing with amplitude gains
##############################
Flow('env','cmp','envelope | scale axis=2')
Flow('amp','env',
     '''
     math output="atan(input)" |
     math output="input-sin(input)" |
     add mode=d $SOURCE
     ''')

Flow('acmp','cmp amp','add mode=p ${SOURCES[1]}')
Flow('ncmp','cmp',
     '''
     scale axis=2 |
     math output="atan(input)" |
     math output="input-sin(input)"
     ''')

# Interpolate near offsets
Flow('cmp0','cmp','window max2=1 | pad beg2=5')
Flow('cmp1','cmp0','reverse which=2 opt=n | cat axis=2 $SOURCE')
Flow('mask0','cmp','math output=1 | window max2=1 | pad beg2=5')
Flow('mask1','mask0','reverse which=2 opt=n | cat axis=2 $SOURCE')
Flow('dip1','cmp1','math output="x2*%g/(x1+0.001)" ' % (0.05/(0.004*1.5*1.5)))
Flow('dip2','cmp1 dip1 mask1',
     'twodip2 eps=100 lam=10 dip1=${SOURCES[1]} mask=${SOURCES[2]} q0=0')
Flow('mis','cmp1 dip2 mask1',
     '''
     planemis2 dip=${SOURCES[1]} mask=${SOURCES[2]} verb=y prec=0 niter=10000
     ''')
Flow('mis2','mis cmp','window min2=0 n2=6 | cat axis=2 ${SOURCES[1]}')
Flow('mis3','mis2','window f2=1 | reverse which=2 opt=n | cat axis=2 $SOURCE')

Plot('mis2','grey title="Near Offsets Interpolated" ')

# Predict multiples
###################
Flow('ccmp','mis2','pad n1=2048 | fft1 | fft3')
Flow('mult','ccmp',
     'add mode=p $SOURCE | fft3 inv=y | fft1 inv=y | window n1=1000')
Plot('mult','window f2=6 | ' + grey('(b) SRME-predicted Multiples'))

# Mask the important part
#########################
Flow('mask','mult',
     'math output=1 | mutter hyper=y t0=0.7 v0=2 half=n | smooth rect1=5 rect2=5')
Flow('cmp2','mask mis2','add mode=p ${SOURCES[1]}')
Flow('mult2','mask mult','add mode=p ${SOURCES[1]}')

# Estimate dips
###############
Flow('mdip','mult2','dip rect1=20 rect2=10 liter=40 pmin=0')
Flow('vdip','cmp2',
     'math output="%g*x2/(x1+0.004)" ' % (0.05/(2.5*2.5*0.004)))
Flow('mask2','mdip','mask max=5 | dd type=float | smooth rect1=5 rect2=5')

Flow('mdip2','cmp2 mdip mask2 vdip',
     '''
     twodip2 eps=15 lam=5 dip2=${SOURCES[1]} dip1=${SOURCES[3]} mask=${SOURCES[2]} verb=y 
     ''')

Flow('mdip0','mdip2','window n3=1 f2=6')
Flow('mdip1','mdip2','window f3=1 f2=6')

Plot('mdip0',
     grey('(c) Signal Slope',
          'color=j scalebar=y allpos=y clip=5 minval=0 maxval=6 barlabel="Slope (samples)" '))
Plot('mdip1',
     grey('(d) Noise Slope',
          'color=j scalebar=y allpos=y clip=5 minval=0 maxval=6 barlabel="Slope (samples)" '))

Result('cmp','cmp mult mdip0 mdip1','SideBySideAniso')

Flow('pvel0','cmp mdip0','pveltran half=n v0=1 dv=0.02 nv=100 dip=${SOURCES[1]}')
Flow('pvel1','cmp mdip1','pveltran half=n v0=1 dv=0.02 nv=100 dip=${SOURCES[1]}')
Flow('pvel','pvel0 pvel1','stack | envelope')

# Separate signal/noise
#######################
Flow('comp','cmp2 mdip2','pwdsigk dips=${SOURCES[1]} verb=y niter=10000 eps=0.01')

Flow('nois','comp mask','window f3=1 | add mode=p ${SOURCES[1]}')
Plot('nois','window f2=6 | ' + grey('(b) Estimated Noise','clip=228'))

Flow('sign','mis2 nois','add scale=1,-1 ${SOURCES[1]}')
Plot('sign','window f2=6 | ' + grey('(a) Estimated Signal','clip=228'))

Flow('cmp2s','cmp2','window n2=64')
Flow('mdip2s','mdip2','window n2=64')
Flow('masks','mask','window n2=64')

Flow('comps','cmp2s mdip2s','seisigk dips=${SOURCES[1]} verb=y niter=1000')

Flow('noiss','comps masks','window f3=1 | add mode=p ${SOURCES[1]}')
Plot('noiss','window f2=6 | ' + grey('(b) Estimated Noise','clip=228'))

Flow('signs','mis2 noiss','window n2=64 | add scale=1,-1 ${SOURCES[1]}')
Plot('signs','window f2=6 | ' + grey('(a) Estimated Signal','clip=228'))

# Velocity scans before and after
#################################
for dat in ('cmp','sign'):
    Flow('v'+dat,dat,'vscan semblance=y v0=1 nv=100 dv=0.02 half=n')
    Plot('v'+dat,grey('(%s) Velocity Scan (%s)' % (('c','Data'),('d','Estimated Signal'))[dat=='sign'],
                      'label2="Velocity (km/s)" color=j allpos=y'))

# Final result
##############
Result('super','sign nois vcmp vsign','SideBySideAniso')

Result('supers','signs noiss','SideBySideAniso')

End()