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

## def grey(title,allpos=1):
##     return '''
##     grey transp=y title="%s"
##     ''' % (title)

## Fetch('model0.dat','bgp',private)

## Flow('model0','model0.dat',
##      '''
##      echo in=$SOURCE n1=1500 n2=160 data_format=native_float |
##      put d1=0.002 d2=0.05 o1=0 o2=-1
##      ''',stdin=0)
## Flow('model','model0',
##      '''
##      window n2=128 f2=20
##      ''')
## Result('model',
##        '''
##        grey transp=y title="Synthetic model"
##        label1=Time label2=Lateral unit1=s unit2=km
##        ''')

def grey(title):
    return '''
    grey title="%s" label2=Trace label1="Time" unit1=s screenratio=1.5
    parallel2=n format2="%%3.1f" labelsz=6 titlesz=7.5 clip=675.48
    font=2 labelfat=2 titlefat=2 d2num=0.5 o2num=2 n2tic=4
    ''' % title

Fetch('sean.HH','bp')
Flow('sean','sean.HH',
     'dd form=native | window n3=1 f3=3 n1=500 | bandpass fhi=50')
Plot('sean',grey('Original'))

# Mask for random traces
Flow('mcut','sean',
     '''
     window n1=1 |
     noise rep=y type=n seed=2008 |
     math output="input^1" |
     mask min=-0.1 | cut n1=5 | cut n1=5 f1=115
     ''')

# Zero out random traces
Flow('zero','sean mcut','headercut mask=${SOURCES[1]}')
Result('zero',grey('Missing 40% traces'))

# Yang's recovery
Flow('mask','mcut','spray axis=1 n=500 d=0.004 | dd type=float')
Plot('mask',
     '''
     grey title=Mask label2=Trace
     label1="Time" unit1=s screenratio=1.5
     labelfat=4 font=2 titlefat=4 color=j scalebar=y
     ''')

# Adaptive PEFs
Flow('rzero','zero','reverse which=2 | cat axis=2 ${SOURCES[0]}')
Flow('rmask','mask','reverse which=2 | cat axis=2 ${SOURCES[0]}')

Flow('apef mask2','rzero rmask',
     '''
     apef maskin=${SOURCES[1]} maskout=${TARGETS[1]} jump=1
     a=4,3 niter=500 rect1=50 rect2=10 verb=y
     ''')
Plot('mask2',
     '''
     grey title=Mask2 label2=Trace
     label1="Time" unit1=s screenratio=1.5
     labelfat=4 font=2 titlefat=4 color=j scalebar=y
     ''')

Flow('amiss','rzero apef rmask',
     'miss4 filt=${SOURCES[1]} mask=${SOURCES[2]} niter=1000 verb=y')
Flow('ramiss','amiss','window f2=180 | put d2=1 o2=0')
Result('ramiss',grey('Interpolated result' ))

Flow('fcoe','apef','window n1=1 f1=2 n2=1')
Plot('fcoe',
     '''
     grey title="First" screenratio=1.5 
     wheretitle=b wherexlabel=t color=j gainpanel=e scalebar=y
     labelfat=4 barwidth=0.2 font=2 titlefat=4 
     label1=Time unit1=s label2=Trace 
     ''')
Flow('mcoe','apef',
     'cut n1=8 n2=1 | stack axis=2 norm=n | stack axis=1 norm=n')
Plot('mcoe',
     '''
     grey title="Mean" screenratio=1.5 bias=-0.6
     wheretitle=b wherexlabel=t color=j gainpanel=e scalebar=y
     labelfat=4 barwidth=0.2 font=2 titlefat=4
     label1=Time unit1=s label2=Trace 
     ''')
Flow('diff','sean ramiss','add scale=1,-1 ${SOURCES[1]}')
Result('diff',grey('Difference'))

################
# Compare with POCS
Flow('fourpocs perr','zero mask sean',
     '''
     fourmis2 mask=${SOURCES[1]} niter=200 oper=p ordert=1.
     perc=99 verb=y error=y ref=${SOURCES[2]} res=${TARGETS[1]}
     ''')
Plot('fourpocs',grey('Fourier POCS (500 iterations)' ))


#################
# 3D test
#################
def grey3(title):
    return '''
    put o3=0 |
    byte allpos=n clip=622.105 |
    grey3 title="%s" label2=Trace label1="Time" unit1=s label3=Shot
    point1=0.8 point2=0.7 frame1=250 frame2=50 frame3=10
    parallel2=n o2num=2 d2num=1 n2tic=2 screenratio=1.3
    font=2 labelfat=2 titlefat=2
    ''' % title
Flow('sean3','sean.HH',
     'dd form=native | bandpass fhi=50')
Result('sean3',grey3('Original'))

# Mask for random traces
Flow('mcut3','sean3',
     '''
     window n1=1 |
     noise rep=y type=n seed=2008 |
     math output="input^1" |
     mask min=-0.0 | cut n1=5
     ''')
Plot('mcut3',
     '''
     dd type=float |
     grey transp=n yreverse=n scalebar=y
     label1="In-line midpoint" label2="Cross-line midpoint"
     wheretitle=t wherexlabel=b title="Mask Map"
     ''')
# Zero out random traces
Flow('zero3','sean3 mcut3','headercut mask=${SOURCES[1]}')
Result('zero3',grey3('Missing 50% traces'))
Flow('maskin3','mcut3','dd type=float | spray axis=1 n=500 d=0.004 o=1.9')

Flow('rzero3','zero3',
     '''
     reverse which=2 |
     window f2=130 |
     cat axis=2 ${SOURCES[0]}
     ''')
Flow('rmask3','maskin3',
     '''
     reverse which=2 |
     window f2=130 |
     cat axis=2 ${SOURCES[0]}
     ''')
Plot('rzero3',grey3('Extended missing data'))

# 3D APEF
Flow('apef3 maskout3','rzero3 rmask3',
     '''
     apef maskin=${SOURCES[1]} maskout=${TARGETS[1]} jump=1
     a=4,2,2 niter=500 rect1=50 rect2=10 rect3=10 verb=y
     ''')
Plot('maskout3',
     '''
     window n1=1 f1=100 |
     grey transp=n yreverse=n scalebar=y
     label1="In-line midpoint" label2="Cross-line midpoint"
     wheretitle=t wherexlabel=b
     title="Output Mask Map"
     ''')

Flow('amiss3','rzero3 apef3 rmask3',
     '''
     miss43 filt=${SOURCES[1]} mask=${SOURCES[2]} niter=1000 verb=y |
     window f2=50 | put d2=1 o2=0
     ''')

Result('amiss3',grey3('RNA Result'))

Flow('adiff3','sean3 amiss3','add scale=1,-1 ${SOURCES[1]}')

Result('adiff3',grey3('RNA Error'))

# 3D PWD
Flow('dip3','rzero3 rmask3',
     '''
     dip verb=y rect1=10 rect2=5 rect3=5
     order=2 niter=10 mask=${SOURCES[1]}
     ''')
Flow('miss3','rzero3 rmask3 dip3',
     '''
     planemis3 mask=${SOURCES[1]} dip=${SOURCES[2]}
     order=2 niter=200 verb=y |
     window f2=50 | put d2=1 o2=0
     ''')
Result('miss3',grey3('PWD Result'))

Flow('diff3','sean3 miss3','add scale=1,-1 ${SOURCES[1]}')
Result('diff3',grey3('PWD Error'))


End()