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

data = 'bend_l1_pmig_enhanc.sgy'
Fetch(data,'vecta',private)
Flow('data',data,'segyread read=data | window n2=471 max1=1.5 | scale axis=2')

Plot('data','grey clip=0.44 title="Input Data" ')
Result('data','Overlay')

# find adaptive PEF
ns = 2

Flow('shift','data','shift1 ns=%d' % ns)

Flow('itrace','data','envelope hilb=y')
Flow('ctrace','data itrace','cmplx ${SOURCES[1]}')

Flow('ishift','shift','envelope hilb=y')
Flow('cshift','shift ishift','cmplx ${SOURCES[1]} | transp plane=23')

Flow('cpef cpre','cshift ctrace',
     'clpf match=${SOURCES[1]} rect1=20 rect2=40 pred=${TARGETS[1]}')
Flow('cerr','cpre ctrace','add scale=-1,1 ${SOURCES[1]}')

Result('cerr','real | grey clip=0.44 title="Residual after RNAR" ')

Flow('cpoly','cpef','window n3=1 | math output=-1 | cat axis=3 $SOURCE')
Flow('croots','cpoly',
     '''
     transp plane=23 | transp plane=12 |
     roots verb=n niter=100 sort=p |
     transp plane=12 | transp plane=23
     ''')

# Frequency components
import math
wf = 2*math.pi
dt = 0.002

Flow('group','croots',
     '''
     math output="-arg(input)/%g" | real 
     ''' % (wf*dt))

Flow('freqs','group',
     '''
     causint | math output="input*%g/(x1+%g)" 
     ''' % (dt,dt))

for n in range(ns):
    group = 'group%d' % n
    Flow(group,'group','window n3=1 f3=%d' % n)
    Plot(group,
    '''
    grey color=j bias=50 clip=25 scalebar=y title="Instantaneous Frequency %d"
    barlabel=Frequency barunit=Hz unit2=
    ''' % (n+1))
    Result(group,'Overlay')

    freq = 'freq%d' % n
    Flow(freq,'freqs','window n3=1 f3=%d' % n)
    Plot(freq,
    '''
    grey color=j bias=50 clip=25 scalebar=y title="Local Frequency %d"
    barlabel=Frequency barunit=Hz 
    ''' % (n+1))
    Result(freq,'Overlay')

Result('freqs','freq0 freq1','OverUnderIso')

Result('vgroup','group0 group1','OverUnderIso')

# Decomposition

Flow('comps','freqs','rtoc | math output="exp(I*input*x1*%g)" ' % wf)

Flow('cwht cfit','comps ctrace',
     'clpf match=${SOURCES[1]} rect1=5 rect2=5 pred=${TARGETS[1]}')

Flow('cdif','cfit ctrace','add scale=1,-1 ${SOURCES[1]}')
Plot('cdif','real | grey clip=0.44 title="(b) Residual" ')
Plot('cfit','real | grey clip=0.44 title="Fit by RNR" unit2=')

Flow('csign','comps cwht','math other=${SOURCES[1]} output="input*other" ')

Plot('label',None,'box x0=5.5 y0=7.2 label="Gas?" size=0.3 xt=0.75 yt=0.75')

for n in range(ns):
    sign = 'sign%d' % n
    Flow(sign,'csign','window n3=1 f3=%d' % n)
    Plot(sign,'real | grey clip=0.44 title="Component %d" unit2=' % (n+1))
    cwht = 'cwht%d' % n
    Flow(cwht,'cwht','window n3=1 f3=%d' % n)
    Plot(cwht,
           '''
           math output="abs(input)" | real | grey allpos=y color=j
           title="Component %d" scalebar=y unit2=
           ''' % (n+1))
    Plot('l'+cwht,[cwht,'label'],'Overlay')
    Result(cwht,'Overlay')

Result('vsign','sign0 sign1','OverUnderIso')
Result('vcwht','lcwht0 lcwht1','OverUnderIso')
Result('vdata','data cfit','OverUnderIso')

Result('vdif','cdif','real | grey clip=0.44 title="Residual" ')

End()