from rsf.proj import *

Fetch('hector.HH','ground')

def grey(title):
    return '''
    grey crowd1=0.8 label1="Time (s)" label2="Offset (km)"
    title="%s"
    ''' % title

Flow('orig','hector.HH', 'dd form=native')
Flow('pad','orig','put d3=0.017 | pad beg1=250')
Flow('data','pad','lmostretch v0=1 half=n delay=0')

Flow('aorig','orig','window j2=2')
Flow('adata','data','window j2=2')

def twodip2(nj1,nj2,extra):
    return '''
    twodip2 order=3 nj1=%d nj2=%d eps=8 verb=y gauss=n
    ''' % (nj1,nj2) + extra

for case in ('','a'):
    orig=case+'orig'
    data=case+'data'
    noiz=case+'noiz'
    sign=case+'sign'
    dat2=case+'dat2'
    mask=case+'mask'

    Plot(orig,grey('Original Data'))
    Plot(data,grey('Data after LMO'))

    Flow(noiz,data,'bandpass fhi=10 | mutter v0=1 half=n')
    Flow(sign,data,'bandpass flo=10')

    Plot(noiz,grey('Bandpassed'))

    Flow(mask,noiz,
         'math output="abs(input)" | mask min=0.02 | dd type=float')
    Flow(dat2,[data,mask],'add mode=p ${SOURCES[1]}')

    Result(case+'hector-dat',[orig,data,noiz],'SideBySideAniso')

    if case=='a':
        nj=(5,4)
    else:
        nj=(3,2)

    ndip=case+'ndip'
    sdip=case+'sdip'

    Flow(ndip,noiz,twodip2(nj[0],nj[0],'p0=2 q0=-2'))
    Flow(sdip,sign,twodip2(nj[1],nj[1],'p0=-2 q0=2'))

    pq = case+'pq'

    Flow(pq,[dat2,sdip,ndip,mask],
         twodip2(nj[1],nj[0],
                 'dip1=${SOURCES[1]} dip2=${SOURCES[2]} mask=${SOURCES[3]}'))

    sdip2 = sdip+'2'
    ndip2 = ndip+'2'

    Flow(sdip2,pq,'window n3=1')
    Flow(ndip2,pq,'window f3=1')

    signoi = case+'signoi'
    noiz2 = noiz+'2'
    sign2 = sign+'2'
    signoi2 = signoi+'2'
    orig2 = orig+'2'

    Flow(signoi,[dat2,sdip2,ndip2],
         '''
         planesignoi sdip=${SOURCES[2]} ndip=${SOURCES[1]} 
         order=3 nj1=%d nj2=%d eps=2 verb=y niter=100
         ''' % nj)
    Flow(noiz2,signoi,
         'window n3=1 | lmostretch inv=y v0=1 half=n delay=0 | window f1=250')
    Flow(sign2,[signoi,mask,data,dat2],
         '''
         window f3=1 |
         add mode=p ${SOURCES[1]} | add scale=1,1,-1 ${SOURCES[2:4]} |
         lmostretch inv=y v0=1 half=n delay=0 | window f1=250
         ''')

    lowf = case+'lowf'
    noiz3 = noiz+'3'
    sign3 = sign+'3'

    Flow(lowf,sign2,'bandpass fhi=8')
    Flow(sign3,[sign2,lowf],'add scale=1,-1 ${SOURCES[1]}')
    Flow(noiz3,[noiz2,lowf],'add ${SOURCES[1]}')

    Flow(signoi2,[orig,sign3,noiz3],
         'cat ${SOURCES[1:3]} | byte pclip=95 gainpanel=all')

    Plot(orig2,signoi2,'window n3=1 f3=0 | ' + grey('Data'))
    Plot(sign2,signoi2,'window n3=1 f3=1 | ' + grey('Estimated Signal'))
    Plot(noiz2,signoi2,'window n3=1 f3=2 | ' + grey('Estimated Noise'))
    Result(case+'hector-sn',[orig2,sign2,noiz2],'SideBySideAniso')

End()