import os
from rsf.proj import *
from math import *

###################
# Shot gather test
###################
############
# Setup model parameters
############
nt = 501
dt = 0.004
ot = 0
nx = 501
dx = 0.01
ox = 0
clip=0.2

############
# Setup shot gather
############
for eve in (1,2,3,4):
    spike='spike%d' % eve
    tpara='tpara%d' % eve
    para='para%d'   % eve
    Flow(spike,None,
        '''
        spike n1=%d d1=%g o1=%g n2=%d d2=%g o2=%g nsp=1 k1=%d mag=1  p2=0| 
        ricker1 frequency=15 | put unit2=km label2=Distance
        ''' % (nt,dt,ot,nx,dx,ox,eve*80-30))

    Flow(tpara,spike,
        '''
        window n1=1 | math output="-sqrt(%g*%g+(x1-2.5)*(x1-2.5)/%g/%g)+%g"
        ''' % (0.004*(eve*80-30),0.004*(eve*80-30),2,2,0.004*(eve*80-30)))

    Flow(para,[spike, tpara],
        'datstretch datum=${SOURCES[1]} ')

Flow('para','para1 para2 para3 para4',
     'add ${SOURCES[1]} ${SOURCES[2]} ${SOURCES[3]}')
Result('gpara','para',
       '''
       grey label2=Distance
       unit2=km title="" clip=%f titlesz=7 wherexlabel=t
       labelfat=2 font=2 titlefat=2
       ''' % clip)

############
# Add random noise
############
Flow('npara','para','noise var=1e-2')
Result('gnpara','npara',
       '''
       grey label2=Distance
       unit2=km title="" clip=%f titlesz=7 wherexlabel=t
       labelfat=2 font=2 titlefat=2
       ''' % clip)

############
# tx SOPF
############
# Step 1: t-x Streaming PF
Flow('txr','npara','txspf a=11,3 lambda1=2 lambda2=0.8')
Flow('tx','npara txr','math ref=${SOURCES[1]} output="input-ref"')

# Step 2: streaming orthogonalization
Flow('r2 h2','tx txr',
     'txsorth gamma1=0.4 gamma2=1.0 noise=${SOURCES[1]} sig2=${TARGETS[1]}')
Result('gr2','r2',
       '''
       grey label2=Distance
       unit2=km title="" clip=%f titlesz=7 wherexlabel=t
       labelfat=2 font=2 titlefat=2
       ''' % clip)
Result('gh2','h2',
       '''
       grey label2=Distance
       unit2=km title="" clip=%f titlesz=7 wherexlabel=t
       labelfat=2 font=2 titlefat=2
       ''' % clip)
Result('diff2','h2 para',
       '''
       add scale=1,-1 ${SOURCES[1]} |
       grey label2=Distance
       unit2=km title="" clip=%f titlesz=7 wherexlabel=t
       labelfat=2 font=2 titlefat=2
       ''' % clip)

############
# FXdecon
############
Flow('fx','npara','fxdecon lenf=4 n2w=20')
Result('gfx','fx',
       '''
       grey label2=Distance
       unit2=km title="" clip=%f titlesz=7 wherexlabel=t
       labelfat=2 font=2 titlefat=2
       ''' % clip)
Flow('fxn2','npara fx','math ref=${SOURCES[1]} output="input-ref"')
Result('gfxn2','fxn2',
       '''
       grey label2=Distance
       unit2=km title="" clip=%f titlesz=7 wherexlabel=t
       labelfat=2 font=2 titlefat=2
       ''' % clip)

############
# fx RNA
############
Flow('nft','npara','fft1 ')
nshifts = []
for s in range(1,5):
    nshift = 'nshift-%d' % s
    Flow(nshift,'nft','window f2=%d | pad end2=%d' % (s,s))
    nshifts.append(nshift)

    nshift = 'nshift+%d' % s
    Flow(nshift,'nft','window n2=%d | pad beg2=%d ' % (nx-s,s))
    nshifts.append(nshift)

Flow('nshifts',nshifts,
     'cat ${SOURCES[1:%d]} axis=3 | put o2=0 ' % len(nshifts))
Flow('nflt npref','nshifts nft',
     '''
     clpf match=${SOURCES[1]} pred=${TARGETS[1]}
     rect2=20 rect1=3 niter=10 verb=n
     ''')

Flow('rna','npref','fft1 inv=y ')
Result('grna','rna',
       '''
       grey label2=Distance
       unit2=km title="" clip=%f titlesz=7 wherexlabel=t
       labelfat=2 font=2 titlefat=2
       ''' % clip)

Flow('rnan2','npara rna','math x=${SOURCES[1]} output="x-input"')
Result('grnan2','rnan2',
       '''
       grey label2=Distance
       unit2=km title="" clip=%f titlesz=7 wherexlabel=t
       labelfat=2 font=2 titlefat=2
       ''' % clip)
Result('diff1','rna para',
       '''
       add scale=1,-1 ${SOURCES[1]} |
       grey label2=Distance
       unit2=km title="" clip=%f titlesz=7 wherexlabel=t
       labelfat=2 font=2 titlefat=2
       ''' % clip)

End()