from rsf.proj import *
import rsf.recipes.beg as beg

Fetch('chev.HH','chevron',beg.server)

Flow('chev','chev.HH',
     '''
     dd form=native |
     put o2=7923.3979 d2=12.5 o3=10298.643 d3=12.5
     label1=Time unit1=s label2=Inline unit2=m label3=Crossline unit3=m |
     scale dscale=0.01
     ''')

Flow('spec','chev','pad n1=256 | spectra all=y')

def cubeplot(title,clip='',extra=''):
    return '''
    window n1=34 min1=0.402 |
    byte gainpanel=all %s |
    transp plane=13 |
    grey3 frame3=18 frame2=350 frame1=250 flat=y point1=0.8 point2=0.8
    label3=Time unit3=s label2=Inline unit2=ft label1=Crossline unit1=ft
    title="%s" %s screenratio=1 labelsz=5
    ''' % (clip,title,extra)

def Cubeplot(name,frame=18):
    t=0.402+0.006*frame
    Plot(name+'0',name,
         '''
         window n1=1 f1=%g | transp | grey label2=Inline unit2=ft label1=Crossline unit1=ft wanttitle=n
         grid=y g1num0=12298.4 g1num=10000 g2num0=13423.6 g2num=10000 gridcol=6 gridfat=5 labelsz=5
         screenratio=1 scalebar=y barlabel=Amplitude parallel2=n
         ''' % frame)

    Plot(name+'1',name,
         '''
         window n3=1 f3=250 | grey label2= unit2= label1=Time unit1=s labelsz=5
         grid=y g1num0=12298.4 g1num=10000 g2num0=%g g2num=10000 gridcol=6 gridfat=5 screenratio=1
         scalebar=y barlabel=Amplitude parallel2=n format2="%%3.1f" o2num=0.4 d2num=0.1 formatbar="%%g"
         title="Inline (ft)" titlesz=5 wheretitle=t dbarnum=1 nbartic=10
         ''' % t)
    Plot(name+'2',name,
         '''
         window n2=1 f2=350 | grey label2= unit2= label1=Time unit1=s labelsz=5
         grid=y g1num0=13423.6 g1num=10000 g2num0=%g g2num=10000 gridcol=6 gridfat=5 screenratio=1
         scalebar=y barlabel=Amplitude parallel2=n format2="%%3.1f" o2num=0.4 d2num=0.1 formatbar="%%g"
         title="Crossline (ft)" titlesz=5 wheretitle=t dbarnum=1 nbartic=10
         ''' % t)
    Plot(name+'12',[name+'1',name+'2'],'OverUnderAniso',vppen='txscale=2')
    Result(name,[name+'0',name+'12'],'OverUnderIso')

#Result('chev',cubeplot(' '))

Cubeplot('chev')

Plot('chev',cubeplot(' '))

patch = 'patch w=84,250,250 p=1,2,2 | put n4=1 n5=4 n6=1'

Flow('chev-pad','chev','pad beg1=25 end1=25')
Flow('patch','chev-pad',patch)
Flow('mask','chev','math output=1 | pad beg1=25 end1=25 |' + patch)

Flow('dip','patch mask','dip rect1=10 rect2=10 rect3=10 mask=${SOURCES[1]}',
     split=[5,4,[0,1]])
Flow('dipx','dip','window n4=1 | put n4=1 n5=2 n6=2 | patch inv=y weight=y')
Flow('dipy','dip','window f4=1 | put n4=1 n5=2 n6=2 | patch inv=y weight=y')
Flow('dips','dipx dipy','cat axis=4 ${SOURCES[1]}')

Result('dipx',
       cubeplot('Inline Dip','bar=bar1.rsf',
                '''
                color=j scalebar=y bar=bar1.rsf wanttitle=n
                barlabel="Inline Dip" bartype=h
                '''))
Result('dipy',
       cubeplot('Crossline Dip','bar=bar2.rsf',
                '''
                color=j scalebar=y bar=bar2.rsf wanttitle=n
                barlabel="Crossline Dip" bartype=h
                '''))

Flow('shift1','chev','window f2=1')
Flow('shift2','chev','window f3=1')

Flow('last1','chev','window f2=449 squeeze=n')
Flow('last2','chev','window f3=449  squeeze=n')

Flow('ref1','shift1 last1','cat axis=2 ${SOURCES[1]}')
Flow('ref2','shift2 last2','cat axis=3 ${SOURCES[1]}')

Flow('ref1s','ref1','add mode=p $SOURCE | stack axis=1 norm=n')
Flow('ref2s','ref2','add mode=p $SOURCE | stack axis=1 norm=n')

Flow('corr1','ref1 chev','add mode=p ${SOURCES[1]} | stack axis=1 norm=n')
Flow('corr2','ref2 chev','add mode=p ${SOURCES[1]} | stack axis=1 norm=n')

Flow('chev2','chev','add mode=p $SOURCE | stack axis=1 norm=n')

Flow('cos1','corr1 chev2 ref1s',
     '''
     math s1=${SOURCES[1]} s2=${SOURCES[2]} output="(s1*s2)/(input*input)"
     ''')
Flow('cos2','corr2 chev2 ref2s',
     '''
     math s1=${SOURCES[1]} s2=${SOURCES[2]} output="(s1*s2)/(input*input)"
     ''')
Flow('cos','cos1 cos2',
     'cat axis=3 ${SOURCES[1]} ')

Plot('cos1',
     '''
     grey color=j scalebar=y bartype=v allpos=y bias=1 
     title="West-East Cost" transp=n yreverse=n
     ''')
Plot('cos2',
     '''
     grey color=j scalebar=y bartype=v allpos=y bias=1
     title="South-North Cost" transp=n yreverse=n
     ''')
Result('chos','cos1 cos2','SideBySideIso')

Flow('seed','dips','window n2=1 n3=1 n4=1 | math output=x1')
Flow('pick','dips seed cos',
     '''
     pwpaint3 seed=${SOURCES[1]} cost=${SOURCES[2]} ref2=250 ref3=350 | 
     bandpass fhi=40
     ''')

Result('pick',cubeplot('','allpos=y','color=j wanttitle=n'))

Flow('flat','chev-pad pick','iwarp warp=${SOURCES[1]} eps=1 n1=34 o1=0.402')
Result('flat',cubeplot(''))


# multiple references

picks = []
for ref in ((150,150),(300,150),(300,300)):
        pick = 'pick%d-%d' % ref
        picks.append(pick)

        Flow(pick,'dips seed cos',
             '''
             pwpaint3 seed=${SOURCES[1]} cost=${SOURCES[2]} ref2=%d ref3=%d | 
             bandpass fhi=40
             ''' % ref)

np = len(picks)
Flow('picks',picks,
     'add ${SOURCES[1:%d]} | scale dscale=%g' % (np,1.0/np))

Flow('flats','chev-pad picks','iwarp warp=${SOURCES[1]} eps=1 n1=34 o1=0.402')
#Result('flats',cubeplot('') + ' frame3=20')

Cubeplot('flats',frame=20)

Flow('spec-flat','flats','pad n1=256 | spectra all=y')

Result('spec','spec spec-flat','cat axis=2 ${SOURCES[1]} | scale axis=1 | graph dash=1,0 title=Spectrum label2= unit2=')

Flow('timefreq','flats','timefreq rect=5 nw=12 dw=5')

Result('slice','flats','window n1=1 f1=20 | grey transp=n title=Amplitude unit1=ft unit2=ft screenratio=1 scalebar=y parallel2=n labelsz=6')

for freq in (10,15,20,25,30,35,40,45,50):
    slice = 'slice%d' % freq
    Result(slice,'timefreq',
           'window n1=1 f1=20 n2=1 min2=%d | grey title="%g Hz" transp=n unit1=ft unit2=ft allpos=y scalebar=y parallel2=n labelsz=6 screenratio=1' % (freq,freq))

Flow('den','timefreq','stack')
Flow('num','timefreq','math output="x2*input" | stack')

Flow('freq','num den','div ${SOURCES[1]}')

Result('freq','window n1=1 f1=20 | grey color=j scalebar=y title="Dominant Frequency" bias=15 allpos=y barlabel=Frequency barunit=Hz')

End()