from rsf.proj import *
import math

Flow('cup2',None,
     'math n2=1 n3=1 n1=256 o1=0 d1=0.008 output="%g*(3+cos(%g*(%g*x1-1)))" ' %
     ((256*0.004/9.0),math.pi,2.0/((1+256/2)*0.008)))
Flow('cup','cup2',
     '''
     unif2 n1=256 o1=0.004 d1=0.004 |
     deriv |
     bandpass flo=10 fhi=50 | transp | bandpass fhi=50 | transp
     ''')
Result('cup',
       'grey label1=Time unit1=s label2=Midpoint unit2=km pclip=99.5 title=Model')

Flow('shots2','cup2',
     '''
     scale dscale=0.666666666 |
     kirmod freq=30 nt=256 dt=0.004 t0=0.004
     type=c vel=1.5 
     ns=316 ds=0.008 s0=-0.48
     nh=61  dh=0.016 h0=0
     ''')
Flow('data2','shots2',
     '''
     put d2=0.008 |
     shot2cmp |
     window min3=0 n3=256 |
     transp plane=23
     ''')

Flow('data','cup',
     '''halfint inv=y |
     preconstkirch zero=y inv=y h0=0 dh=0.008 nh=61 vel=1.5 |
     window''')

Flow('nmo','data',
     'transp plane=23 | nmostretch v0=1.5 | transp plane=23')
Flow('mask','nmo','shotholes perc=0.9')
Flow('hole','nmo mask','headercut mask=${SOURCES[1]}')

nmos = []
holes = []
for slice in [0,25,50]:
    offset = slice*0.008
    grey = '''
         window n1=100 f1=50 n3=1 f3=%d | 
         grey pclip=99.5 title="h=%g"
         label1=Time unit1=s label2=Midpoint unit2=km
         ''' % (slice,offset)
    nmo = 'nmo%d' % slice
    hole = 'hole%d' % slice
    Plot(nmo,'nmo',grey)
    Plot(hole,'hole',grey)
    nmos.append(nmo)
    holes.append(hole)
Plot('nmo',nmos,'OverUnderAniso')
Plot('hole',holes,'OverUnderAniso')
Result('data',['nmo','hole'],'SideBySideAniso',vppen='txscale=3')

nmos = []
holes = []
for slice in [74,99,124]:
    time = 0.004 + slice*0.004
    grey = '''
         window n1=1 f1=%d | 
         grey pclip=99.5 transp=n title="t=%g"
         label1=Midpoint unit1=km label2=Offset unit2=km
         ''' % (slice,time)
    nmo = 'tnmo%d' % slice
    hole = 'thole%d' % slice
    Plot(nmo,'nmo',grey)
    Plot(hole,'hole',grey)
    nmos.append(nmo)
    holes.append(hole)
Plot('tnmo',nmos,'OverUnderAniso')
Plot('thole',holes,'OverUnderAniso')
Result('tslice',['tnmo','thole'],'SideBySideAniso',vppen='txscale=3')

Flow('fftd','hole',
     'window f1=50 | logstretch nout=256 | fft1 | transp | transp plane=23')
Flow('fill',['fftd','mask'],
     'window n3=55 | ofilp known=${SOURCES[1]} niter=500')
Flow('hfill',['fftd','mask'],
     'window n3=55 | ofilp known=${SOURCES[1]} niter=500 simple=y')

ffts = []
fills = []
for slice in [8,13,19]:
    frec = int(slice*0.585393)
    grey = '''
         window n3=1 f3=%d | real | 
         grey title="frequency=%g" transp=n
         label1=Midpoint unit1=km label2=Offset unit2=km
         ''' % (slice,frec)
    fft = 'fft%d' % slice
    fill = 'fill%d' % slice
    Plot(fft,'fftd',grey)
    Plot(fill,'fill',grey)
    ffts.append(fft)
    fills.append(fill)
Plot('fft',ffts,'OverUnderAniso')
Plot('fill',fills,'OverUnderAniso')
Result('fslice',['fft','fill'],'SideBySideAniso',vppen='txscale=3')

back = '''pad n3=136 | transp plane=23 | transp | fft1 inv=y |
     window n1=256 | logstretch inv=y'''

Flow('fild','fill',back)
Flow('hild','hfill',back)

goods = []
bads = []
for slice in [0,25,50]:
    offset = slice*0.008
    grey = '''
         window n1=100 n3=1 f3=%d | 
         grey pclip=99.5 title="h=%g"
         label1=Time unit1=s label2=Midpoint unit2=km
         ''' % (slice,offset)
    good = 'good%d' % slice
    bad = 'bad%d' % slice
    Plot(good,'fild',grey)
    Plot(bad,'hild',grey)
    bads.append(bad)
    goods.append(good)
Plot('bad',bads,'OverUnderAniso')
Plot('good',goods,'OverUnderAniso')
Result('all',['bad','good'],'SideBySideAniso',vppen='txscale=3')

Flow('shot','nmo','transp plane=23 | cmp2shot')
Flow('shot1','shot','window n3=1 f3=227')
Flow('shot2','shot','window n3=1 f3=228')
Flow('shot3','shot','window n3=1 f3=229')

clip=2145.77

grey = 'grey clip=%g' % clip + ' label1=time label2=offset crowd1=0.85 title="%s" '

Plot('shot1',grey % 'shot1')
Plot('shot3',grey % 'shot3')

Flow('simp',['shot1','shot3'],'add ${SOURCES[1]} scale=0.5,0.5')
Plot('simp',grey % '(shot1+shot3)/2')

Flow('simperr',['simp','shot2'],'add scale=1,-1 ${SOURCES[1]}')
Plot('simperr',grey % 'error: (shot1+shot3)/2-shot2')
Result('shot3','shot1 shot3 simp simperr','TwoRows')

Flow('shotftd','shot',
     '''
     window f1=50 | logstretch nout=256 |
     fft1 | transp | transp plane=23
     ''')
Flow('shots','shotftd','window n2=2 f2=227 j2=2')
Flow('shotin','shots','infill | window n2=1 f2=1')
Flow('shotnd','shotin',
     '''
     transp | fft1 inv=y |
     logstretch inv=y | pad beg1=50 
     ''')
Flow('shoterr','shotnd shot2','add scale=1,-1 ${SOURCES[1]}')

Flow('shot0','shotftd','window n2=1 f2=100 squeeze=n')
Flow('shotp','shot0','shotprop ns=100 ds=0.008')

Plot('shotnd',grey % 'Interpolated Shot')
Plot('shoterr',grey % 'Error')
Result('shotin','shotnd shoterr','SideBySideIso')

End()