from rsf.proj import *
from rsf.recipes.beg import server as private

import random, math

random.seed(2005)

nr = 0

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

Fetch('elf0.H','elf',private)

Flow('elf','elf0.H',
     '''
     dd form=native | cut n3=1 n2=1 n1=300 f3=663 f2=67 |
     bandpass flo=5 fhi=60
     ''')

Flow('gath','elf','window n2=128 n3=1 f3=500 | put d2=0.0125 o2=0.05')

data='gath'
n1=800
n2=128          # data dimensions
o2=50
d2=12.5         # lateral scale
rect1=10
rect2=10        # smoothing for dip estimation
p0=0
pmin=0          # initial and minimum dips
clip=5          # clip percentile
eps=0.1         # regularization
nsp=200         # number of spikes

Result(data,'grey unit1=s unit2=km  title=Input')

dip = data+'dip'
Flow(dip,data,
     'dip rect1=%d rect2=%d p0=%g pmin=%g' % (rect1,rect2,p0,pmin))
Result(dip,'grey unit1=s unit2=km  color=j title=Slope allpos=y scalebar=y')

seis = data+'seis'
Flow(seis,[data,dip],
     'seislet dip=${SOURCES[1]} eps=%g adj=y inv=y unit=y' % eps)
Result(seis,
       '''
       put o2=0 d2=1 | 
       grey unit1=s title="Seislet Transform" label2=Scale unit2=
       ''')

#    sseis = data+'sseis'
#    Flow(sseis,[data,dip],
#         'seislet dip=${SOURCES[1]} eps=%g adj=y niter=100' % eps)
#    Result(sseis,'grey unit1=s unit2=m title="Sparse Seislet Transform" ')

sinv = data+'sinv'
#    ssinv = data+'ssinv'

Flow(sinv,[seis,dip],'seislet dip=${SOURCES[1]} eps=%g inv=y unit=y' % eps)
Result(sinv,'grey unit1=s unit2=km  title="Inverse Seislet Transform" ')

#    Flow(ssinv,[sseis,dip],'seislet dip=${SOURCES[1]} eps=%g' % eps)
#    Result(ssinv,'grey unit1=s unit2=m  title="Inverse Seislet Transform" ')

wvlt = data+'wvlt'

Flow(wvlt,data,'transp | dwt inv=y unit=y | transp')
Result(wvlt,'grey unit1=s title="Wavelet Transform" label2=Scale unit2=')

recs = []
for c in (1,clip,25):
    rec = '%ssrec%d' % (data,c)
    recs.append(rec)

    Flow(rec,[seis,dip],
         '''
         threshold pclip=%d |
         seislet dip=${SOURCES[1]} eps=%g inv=y unit=y
         ''' % (c,eps))
    Result(rec,
           '''
           mutter v0=1.3 |
           grey unit1=s unit2=km  title="Inverse Seislet Transform (%d%%)"
           ''' % c)

    wrec = '%swrec%d' % (data,c)
    Flow(wrec,wvlt,
         'threshold pclip=%d | transp | dwt adj=y inv=y unit=y | transp' % c)
    Result(wrec,
           '''
           mutter v0=1.3 |
           grey unit1=s unit2=km  title="Inverse Wavelet Transform (%d%%)"
           ''' % c)

max=int(math.log(n2)/math.log(2))
for m in xrange(max):
    scale = int(math.pow(2,m))
    slet = '%sslet%d' % (data,scale)
    Flow(slet,[seis,dip],
         '''
         cut f2=%d | 
         seislet dip=${SOURCES[1]} eps=%g inv=y unit=y
         ''' % (scale,eps))
    Result(slet,
           'mutter v0=1.3 | grey unit1=s unit2=km  title="Denoising result" ')
#    diff = '%sdiff%d' % (data,scale)
#    Flow(diff,[seis,dip],
#         'cut n2=%d | seislet dip=${SOURCES[1]} eps=%g' % (scale,eps))
#    Result(diff,
#           'mutter v0=1300 | grey unit1=s unit2=m  title="Scale=%d" ' % scale)

slet = '%sslet%d' % (data,16)

Result('sign',[slet,dip],
     '''
     seislet dip=${SOURCES[1]} eps=%g adj=y inv=y unit=y |
     put o2=0 d2=1 | 
     grey unit1=s title="Seislet Transform" label2=Scale unit2=
     ''' % eps)

Flow('rand',seis,'noise rep=y var=200000 seed=2006 | mutter v0=1.3')

Flow('seis2',[seis,'rand'],'cat axis=2 ${SOURCES[1]} | put d2=0.00625')
Plot('seis2','grey unit1=s title="Seislet Transform" label2=Scale unit2=')

Flow('dip2',dip,
     'transp | remap1 d1=0.00625 o1=0.050 n1=256 | transp | scale dscale=0.5')
Flow('gath2','seis2 dip2',
     'seislet dip=${SOURCES[1]} eps=%g inv=y unit=y' % eps)
Plot('gath2',
     '''
     window n2=255 | mutter v0=1.2 |
     grey unit1=s unit2=km title="Resampled by 2"
     ''')

Flow('rand2','seis2','noise rep=y var=200000 seed=2006 | mutter v0=1.3')

Flow('seis4','seis2 rand2','cat axis=2 ${SOURCES[1]} | put d2=0.003125')
Result('seis4',
       '''
       put o2=0 d2=1 | 
       grey unit1=s title="Seislet Transform" label2=Scale unit2=
       ''')

Flow('dip4','dip2',
     'transp | remap1 d1=0.003125 o1=0.05 n1=512 | transp | scale dscale=0.5')
Flow('gath4','seis4 dip4',
     'seislet dip=${SOURCES[1]} eps=%g inv=y unit=y' % eps)
Result('gath4',
       '''
       window n2=509 | mutter v0=1.2 |
       grey unit1=s unit2=km title="Resampled by 4"
       ''')

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

imps = data+'imps'
Flow(imps,dip,
     '''
     spike nsp=%d k1=%s k2=%s n1=%d n2=%d o2=%g d2=%g label2="Half offset" |
     seislet dip=$SOURCE eps=%g inv=y unit=y
     ''' % (nsp,k1,k2,n1,n2,o2,d2,eps),stdin=0)
Result(imps,'mutter v0=1.3 | grey unit1=s unit2=m  title=Seislets')


impw = data+'impw'
Flow(impw,dip,
     '''
     spike nsp=%d k1=%s k2=%s n1=%d n2=%d o2=%g d2=%g label2="Half offset" |
     transp | dwt eps=%g adj=y inv=y unit=y | transp
     ''' % (nsp,k1,k2,n1,n2,o2,d2,eps),stdin=0)
Result(impw,'mutter v0=1.3 | grey unit1=s unit2=m  title=Wavelets')

nmos = []
for i2 in range(n2):
    trace = 'trace%d' % i2
    if i2 == 0:
        Flow(trace,'gath','cut f2=1')
    elif i2 == n2-1:
        Flow(trace,'gath','cut n2=%d' % i2)
    else:
        Flow(trace,'gath','cut n2=%d | cut f2=%d' % (i2,i2+1))

    nmo = 'nmo%d' % i2
    nmos.append(nmo)
    Flow(nmo,[trace,dip],
         '''
         seislet dip=${SOURCES[1]} eps=%g adj=y inv=y unit=y type=haar |
         window n2=1         
         ''' % eps)

Flow('snmo',nmos,
     '''
     cat axis=2 ${SOURCES[1:%d]} 
     ''' % len(nmos))

Result('snmo','grey unit1=s unit2=km title="Seislet Moveout" ')

Flow('scn','gath',
     'mutter v0=%g | vscan semblance=y v0=%g nv=%d dv=%g' % (1.4,1.4,120,0.025))
Flow('vel','scn','pick rect1=50 | window')
Flow('nmo','gath vel','mutter v0=%g | nmo velocity=${SOURCES[1]} str=0' % 1.4)

Result('nmo','grey unit1=s unit2=km title="Normal Moveout" ')

Flow('gaths','elf','window n2=128')

Flow('pat','gaths','patch w=800,128,150')
Flow('dips','pat','dip n4=0 rect1=10 rect2=10 rect3=10 p0=0 pmin=0',split=[6,11])
Flow('dip','dips','patch inv=y weight=y')

Flow('seis','gaths dip',
     'seislet dip=${SOURCES[1]} eps=%g adj=y inv=y unit=y' % eps)

Flow('sstack','seis','window n2=1')

Plot('sstack',
     '''
     put d2=0.0133333 |
     grey label2=Midpoint unit2=km label1=Time unit1=s
     title="Seislet Stack"
     ''')

Fetch('elf-stk.rsf','masha')

Plot('stack','elf-stk.rsf',
     '''
     dd form=native | put d2=0.0133333 |
     grey label2=Midpoint unit2=km label1=Time unit1=s
     title="NMO Stack"
     ''')

Result('sstack','stack sstack','SideBySideAniso')

End()