from rsf.proj import *

sgy = 'boonsvillel249.sgy'

Fetch(sgy,'boonsville')
Flow('boon tboon boon.asc boon.bin',sgy,
     '''
     segyread tfile=${TARGETS[1]} hfile=${TARGETS[2]} bfile=${TARGETS[3]} |
     put label1=Time unit1=s label2=Trace o2=1 d2=1
     ''')

Result('boon','grey wanttitle=n')

# Window a trace
Flow('trace','boon','window n2=1 min2=100')

# Apply phase rotations
shifts = []
for a in range(-90,271,3):
    shift = 'shift%d' % a
    shifts.append(shift)
    Flow(shift,'trace','envelope order=100 hilb=y phase=%d' % a)
Flow('shifts',shifts,
     '''
     cat ${SOURCES[1:%d]} axis=2 | put o2=-90 d2=3 
     ''' % len(shifts))

Result('shifts',
       '''
       wiggle poly=y transp=y yreverse=y wanttitle=n
       label2="Phase Shift" unit2="\^o\_" 
       ''')

# Compute similarity with the square
Flow('sqr','shifts','add mode=p $SOURCE')

# Compute local kurtosis
Flow('kur','shifts','focus rect1=500')

Flow('foc','shifts sqr kur',
     'similarity other=${SOURCES[1]} rect1=500 | math foc=${SOURCES[2]} output=foc/input')

Plot('foc','grey allpos=y wanttitle=n color=j')

# Pick a trend (nonstationary phase rotation)
Flow('pik','foc','scale axis=2 | pick rect1=50 vel0=200 gate=100')
Plot('pik',
     '''
     graph transp=y yreverse=y min2=-90 max2=270 plotcol=7 plotfat=5
     pad=n wanttitle=n wantaxis=n
     ''')

Result('foc','foc pik','Overlay')

# Take a slice through phase shifts
Flow('zero','shifts pik','slice pick=${SOURCES[1]}')

# Rotate by 90 degrees
Flow('ninety','zero','envelope order=100 hilb=y phase=90')

Result('zero','trace ninety',
       '''
       cat axis=2 ${SOURCES[1]} | 
       grey color=G wantaxis2=n
       title="Input versus 0\^o\_ Phase"
       ''')

# Apply to the whole data
#########################

# Apply phase rotations
shifts = []
for a in range(-90,271,3):
    shift = 'bshift%d' % a
    shifts.append(shift)
    Flow(shift,'boon','envelope order=100 hilb=y phase=%d' % a)

Flow('bshifts',shifts,
     '''
     cat ${SOURCES[1:%d]} axis=3 | put o3=-90 d3=3 
     ''' % len(shifts))

# Square the signal
Flow('bsqr','bshifts','mul $SOURCE')

# Compute local kurtosis
Flow('bkur','bshifts','focus rect1=500 rect2=50',split=[3,121])

# 1/input instead of foc/input

Flow('bfoc','bshifts bsqr bkur',
     '''
     similarity other=${SOURCES[1]} rect1=500 rect2=50 |
     math foc=${SOURCES[2]} output=1/input  |
     transp plane=23 |
     put  label2="Phase Shift" unit2="\^o\_" 
     ''',split=[3,121],reduce='cat axis=2')

Plot('bfoc','grey color=j gainpanel=all allpos=y',view=1)

Result('bfoc',
       '''
       scale axis=3 |
       byte gainpanel=all allpos=y bar=bar.rsf bias=0.7 |
       grey3 wanttitle=n color=j frame1=1000 frame2=60 frame3=99 flat=n
       label2="Phase Rotation" scalebar=y
       ''')

# Pick a trend (nonstationary phase rotation)
Flow('bpik','bfoc','scale axis=3 | pick rect1=50 rect2=100 vel0=90 gate=100')

Result('bpik','window | grey color=j scalebar=y title="Phase Rotation" unit2=')
Plot('bpik-win','bpik',
     '''
     window n3=25 min1=0.5 max1=1.5 | math output="90-input" |
     grey color=j scalebar=y title="Estimated Phase" 
     ''')

# Take a slice through phase shifts
Flow('b90','bshifts bpik','transp plane=23 | slice pick=${SOURCES[1]} | window')
Result('b90','grey wanttitle=n')

# Rotate by 90 degrees
Flow('b0','b90','envelope order=100 hilb=y phase=+90')
Result('b0','grey wanttitle=n')

Plot('boon-win','boon',
     'window n2=25 min1=0.25 max1=1 | grey title=Input color=G')
Plot('b0-win','b0',
     '''
     window n2=25 min1=0.25 max1=1 | 
     grey title="Phase Corrected" color=G
     ''')
Plot('b90-win','b90',
     '''
     window n2=25 min1=0.25 max1=1 | 
     grey title="Corrected to 90\^o\_" color=G
     ''')


Result('boon-win','boon-win b0-win','SideBySideAniso')

End()