#!/usr/bin/python

from rsf.proj import *

##### make model
#curve1
Flow('line1', None,
     '''
     spike d1=0.004 d2=0.005 n1=501 n2=201 k1=150 p2=-0.5  mag=1.5 |
     ricker1 frequency=11
     ''')
Flow('mo1', None, 'math n1=201 output="x1*x1*0.00004" ')
Flow('curve1', 'line1 mo1', 'stretch datum=${SOURCES[1]} rule=d inv=y')
#curve2
Flow('line2', None,
     '''
     spike d1=0.004 d2=0.005 n1=501 n2=201 k1=350 p2=0.5  mag=1. |
     ricker1 frequency=15
     ''')
Flow('mo2', None, 'math n1=201 output="x1*x1*0.00005" ')
Flow('curve2', 'line2 mo2', 'stretch datum=${SOURCES[1]} rule=d ')
#line3
Flow('line3', None,
     '''
     spike d1=0.004 d2=0.005 n1=501 n2=201 k1=40 p2=3  mag=1 |
     ricker1 frequency=13
     ''')

Flow('mod', 'curve1 curve2 line3',
     'add ${SOURCES[1:3]} | put d2=1 label2=Trace unit2=')
Result('mc','mod',
       '''
       grey
       label2=Trace label1=Time unit1=s unit2=
       title="" wheretitle=b wherexlabel=t 
       ''')

##### random mask
Flow('mask', 'mod',
     '''
     window n1=1 | noise type=n rep=y seed=150 |
     mask min=-0.1 | cut n1=3 f1=118 | dd type=int
     ''')
Flow('mask2', 'mask', 'reverse which=1')
##### irregular missing data
Flow('gap', 'mod mask', 'headercut mask=${SOURCES[1]}')
Result('gc','gap',
       '''
       grey
       label2=Trace label1=Time unit1=s unit2=
       title="" wheretitle=b wherexlabel=t 
       ''')

Flow('pmask','mask','spray axis=1 n=501 d=0.004 | dd type=float')
Flow('pmask-int','mask','spray axis=1 n=501 d=0.004')

##### interpolation using t-x SPF with varying smoothness 
#####                     and noncausal structure
# forward interpolation 1 ->
lambda1 = 0.05
lambda2 = 0.8
a1 = 25
a2 = 23
verb = 1
smooth = 1
epst = 0.5
epsx = 0.5

Flow('gap1','gap','reverse which=1')

Flow('inter','gap1 pmask-int',
     '''
     txspfvsint2 lambda1=%f lambda2=%f a=%d,%d  
     known=${SOURCES[1]} verb=%d smooth=%d epst=%f epsx=%f |
     reverse which=1 
     ''' %(lambda1, lambda2, a1, a2, verb, smooth, epst, epsx))

# backward interpolation 1 <-
Flow('prev','pmask-int','reverse which=2 | reverse which=1 ')
Flow('rev','inter','reverse which=2 | reverse which=1')

Flow('inter2','rev prev',
     '''
     txspfvsint2 lambda1=%f lambda2=%f a=%d,%d  
     known=${SOURCES[1]} verb=%d smooth=%d epst=%f epsx=%f |
     reverse which=2 | reverse which=1 
     ''' %(lambda1, lambda2, a1, a2, verb, smooth, epst, epsx))

# forward interpolation 2 ->

Flow('inter3','inter2 pmask-int',
     '''
     txspfvsint2 lambda1=%f lambda2=%f a=%d,%d 
     known=${SOURCES[1]} verb=%d smooth=%d epst=%f epsx=%f
     ''' %(lambda1, lambda2, a1, a2, verb, smooth, epst,epsx))

# add backward 1 and forward 2
Flow('add','inter3 inter2','add scale=0.5,0.5 ${SOURCES[1]}')
Flow('dif','add mod','add scale=-1,1 ${SOURCES[1]}')
Result('ac','add',
       '''
       grey label2=Trace label1=Time unit1=s unit2=
       title="" wheretitle=b wherexlabel=t clip=0.132428
       ''')
Result('dcs','dif',
       '''
       grey label2=Trace label1=Time unit1=s unit2=
       title="" wheretitle=b wherexlabel=t clip=0.132428
       ''')

####### interpolation using Fourier POCS with patching
Flow('patch','gap','patch w=150,150')

cat3 = []
for n3 in range(0,6):
    cat4 = []
    datap = 'datap%d' %n3
    pocscat = 'pocscat%d'   %n3
    Flow(datap,'patch','window n3=1 f3=%d' % n3)

    for n4 in range(0,3):
        data  = 'data%d%d'  %(n3, n4)
        pocs = 'pocs%d%d' %(n3, n4)
        mask  = 'mask%d%d'  %(n3, n4)
        Flow(data,datap,'window n3=1 f3=%d' %n4)
        Flow(mask,data,
             '''
             math output="abs(input)*10^10" | 
             stack axis=1 | mask min=0.00001 | spray axis=1 n=150
             ''')
        Flow(pocs, '%s %s ' %(data, mask),
             '''
             fourmis2 mask=${SOURCES[1]} niter=150 oper=p ordert=1.
             perc=99 verb=n error=n
             ''')
        cat4.append(pocs)

    Flow(pocscat,cat4,'cat ${SOURCES[1:%d]} axis=4 ' % len(cat4))

    cat3.append(pocscat)

Flow('pocsp',cat3,
     'cat ${SOURCES[1:%d]} axis=3 | patch inv=y weight=y' % len(cat3))

Result('pocsp',
       '''
       grey label2=Trace label1=Time unit1=s unit2=
       title='' wheretitle=b wherexlabel=t 
       ''')

Flow('dcpp','pocsp mod','add scale=-1,1 ${SOURCES[1]}')
Result('dcpp',
       '''
       grey label2=Trace label1=Time unit1=s unit2=
       title="" wheretitle=b wherexlabel=t clip=0.132428
       ''')

##### interpolation using t-x SPF with causal structure
lambdas1 = 0.05
lambdas2 = 0.8
as1 = 20
as2 = 12

# forward interpolation 1 ->
Flow('inters',' gap1 pmask-int',
     '''
     txspfint2 lambda1=%f lambda2=%f a=%d,%d known=${SOURCES[1]} |
     reverse which=1
     ''' %(lambdas1, lambdas2, as1, as2))

# backward interpolation 1 <-
Flow('revs','gap','reverse which=2 | reverse which=1')
Flow('inters2',' revs prev',
     '''
     txspfint2 lambda1=%f lambda2=%f a=%d,%d known=${SOURCES[1]} |
     reverse which=2 | reverse which=1 
     ''' %(lambdas1, lambdas2, as1, as2))

# add backward 1 and forward 1
Flow('adds','inters inters2','add scale=0.5,0.5 ${SOURCES[1]}')
Flow('difs','adds mod','add scale=-1,1 ${SOURCES[1]}')
Result('adds',
       '''
       grey label2=Trace label1=Time unit1=s unit2=
       title="" wheretitle=b wherexlabel=t clip=0.132428
       ''')
Result('difs',
       '''
       grey label2=Trace label1=Time unit1=s unit2=
       title="" wheretitle=b wherexlabel=t clip=0.132428
       ''')

####### interpolation using SPEF on a helix
Flow('interh','gap1 pmask-int',
     'streamissh a=20,12 eps=3 known=${SOURCES[1]} | reverse which=1 ')
Result('interh',
       '''
       grey label2=Trace label1=Time unit1=s unit2=
       title="" wheretitle=b wherexlabel=t clip=0.132428
       ''')
Flow('difh','interh mod','add scale=-1,1 ${SOURCES[1]}')
Result('difh',
       '''
       grey label2=Trace label1=Time unit1=s unit2=
       title="" wheretitle=b wherexlabel=t clip=0.132428
       ''')

End()