from rsf.proj import*
from rsf.prog import RSFROOT

def Grey(data,other):
        Result(data,'window f2=560 n2=180 | put o2=0| real | grey clip=0.7 label2=Trace color=d unit2="" label1=Time unit1="s" title="" wherexlabel=t scalebar=n   %s'%other)

def Grey1(data,other):
        Result(data,'window f2=560 n2=180 | put o2=0 | math output="abs(input)" | real | grey label2=Trace unit2="" label1=Time unit1="s" wherexlabel=t allpos=y color=j title= minval=-0.3 maxval=0.3 scalebar=y unit2= clip=0.3 %s'%other)

def Grey2(data,other):
        Result(data,'window f2=560 n2=180 |put o2=0 | grey clip=0.7 label2=Trace color=d unit2="" label1=Time unit1="s" title="" wherexlabel=t scalebar=n   %s'%other)
#       Result(data,'window f2=560 n2=380 |put o2=0 | grey clip=0.7 label2=Trace color=d unit2="" label1=Time unit1="s" title="" wherexlabel=t scalebar=n   %s'%other)

def Grey3(data,other):
#       Flow([data+'2',data+'1'],data,'window n2=180 >${TARGETS[1]} && <${SOURCES[0]} %s n2=180 f2=200 | sfcat axis=2 ${TARGETS[1]}'%(WhereIs('sfwindow')))
#       Result(data,data+'2','grey clip=0.7 label2=Trace color=d unit2="" label1=Time unit1="s" title="" wherexlabel=t scalebar=n  %s'%other)
        Result(data,'grey clip=0.7 label2=Trace color=d unit2="" label1=Time unit1="s" title="" wherexlabel=t scalebar=n  %s'%other)

def Greyzoom(data,other):
        Result(data,'put d1=0.002 d2=1 o1=1.5 o2=0 | grey minval=-0.3 maxval=0.3 clip=0.3 label2=Trace color=g unit2="" label1=Time unit1="s" title="" wherexlabel=t scalebar=n  %s'%other)

Flow('data','stack.bin',' echo in=$SOURCE n1=751 n2=1285 data_format=native_float | put n1=751 n2=1285 d1=0.004 d2=1 o1=0 o2=0 | scale axis=2' )

#Plot('data','grey clip=0.9 title="Input Data"')
Flow('data-f','data','window f2=560 n2=180 |put o2=0 ')
Grey3('data-f','title="Noisy data"')


# find adaptive PEF
ns = 4

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.9 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.004

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]}')

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)
    cwht = 'cwht%d' % n
    Flow(cwht,'cwht','window n3=1 f3=%d' % n)

Grey1('cwht0','title="Component 1"')
Grey1('cwht1','title="Component 2"')
Grey1('cwht2','title="Component 3"')
Grey1('cwht3','title="Component 4"')
Grey('sign0','title="Component 1"')
Grey('sign1','title="Component 2"')
Grey('sign2','title="Component 3"')
Grey('sign3','title="Component 4"')

Flow('cfit-f','cfit','window f2=560 n2=180 |put o2=0 | real')
Flow('cdif-f','cdif','window f2=560 n2=180 |put o2=0 |real')
Grey3('cfit-f','title="Denoised data (SDRNAR)"')
Grey3('cdif-f','title="Noise (SDRNAR)"')

Flow('data-fx','data','window f2=560 n2=180 | put o2=0 | fxdecon n2w=180 lenf=10')
Grey3('data-fx','title="Denoised data (FX)"')

Flow('data-fx-dif','data data-fx','window f2=560 n2=180 | put o2=0 | add scale=1,-1 ${SOURCES[1]}')
Grey3('data-fx-dif','title="Noise (FX)"')

Flow('data-mf0','data','window f2=560 n2=180 | put o2=0 | transp | mean nfw=30 | transp')
Grey3('data-mf0','title="Denoised data (mean filter)"')

Flow('data-mf-dif0','data data-mf0','window f2=560 n2=180 | put o2=0 | add scale=1,-1 ${SOURCES[1]}')
Grey3('data-mf-dif0','title="Noise (mean filter)"')

## Create label A
Plot('labela',None,
        '''
        box x0=5.6 y0=4.25 label="A" xt=0.5 yt=0.5 length=0.75 
        ''')

Plot('labelb',None,
        '''
        box x0=10.2 y0=4.25 label="B" xt=0.5 yt=0.5 length=0.75 
        ''')

Plot('labelc',None,
        '''
        box x0=8.3 y0=4.2 label="C" xt=0.5 yt=0.5 length=0.75 
        ''')

Result('data-mf','Fig/data-mf0.vpl labela labelb','Overlay')
Result('data-mf-dif','Fig/data-mf-dif0.vpl labelc','Overlay')


End()