from rsf.proj import *
import math

nt=5001
pad=500
dt = 0.0004
wf = 2*math.pi

Flow('sig',None,
    '''
    math n1=%d o1=%g d1=%g n2=1 output="cos(%g*(2.5*x1*x1+8)*x1)+cos(%g*(5*x1*x1+25)*x1)"
    ''' % (nt+2*pad, -pad*dt, dt, wf, wf))

Flow('true1','sig','math output="2.5*x1*x1+8" ')
Flow('true2','sig','math output="5*x1*x1+25"  ')

Flow('true','true1 true2','cat axis=2 ${SOURCES[1]}')


window = 'window f1=%d n1=%d | ' % (pad,nt)
graph = 'graph pad=n crowd1=0.75 label2=Amplitude label1=Time unit1=s crowd2=0.3 wanttitle=n min2=-2 max2=2'

Plot('sig',window + graph)
Result('sig','Overlay')

remap = '| remap1 n1=%d o1=%g d1=%g' % (nt+2*pad, -pad*dt, dt)

Flow('tf','sig','window j1=10 | timefreq rect=17' + remap)

Result('tf',
       window + '''
       transp | scale axis=2 | grey wanttitle=n color=j screenratio=0.8 scalebar=n bartype=v
       barwidth=0.2 crowd1=0.75  crowd2=0.3 wherexlabel=b allpos=y
       label2=Time unit2=s label1=Frequency unit1=Hz n2tic=30 grid=y
       ''')

# apply shifts in time

Flow('shift','sig','window j1=10 | shift1 ns=2')

# find adaptive PEF

# analytical trace
Flow('itrace','sig','window j1=10 | envelope hilb=y')
Flow('ctrace','sig itrace','window j1=10 | cmplx ${SOURCES[1]}')

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

Flow('cpef cpre','cshift ctrace',
     'clpf match=${SOURCES[1]} rect1=5 pred=${TARGETS[1]}')
Flow('cerr','cpre ctrace','add scale=-1,1 ${SOURCES[1]} | real' + remap)

Plot('cerr',
     window + '''     
     graph title="Nonstationary Deconvolution" min2=-1 max2=1
     pad=n crowd1=0.75 label2=Amplitude label1=Time unit1=s crowd2=0.3 wanttitle=n
     ''')

Result('cdecon','sig cerr','OverUnderAniso')

Result('cerr',window + graph)

# Find complex polynomial roots

Flow('cpoly','cpef','window n2=1 | math output=-1 | cat axis=2 $SOURCE')
Flow('croots','cpoly',
     'transp | roots verb=n niter=100 sort=p | transp')

# Frequency components

Flow('group','croots',
     '''
     math output="-arg(input)/%g" | real | clip2 lower=0  
     ''' % (10*wf*dt) + remap)

Result('group',
       window + '''
       graph title=Frequencies yreverse=y min2=0 max2=125 pad=n wanttitle=n screenratio=0.8 scalebar=n bartype=v
       barwidth=0.2 crowd1=0.75  crowd2=0.3 wherexlabel=b plotcol=5,4 plotfat=3
       label1=Time unit1=s label2=Frequency unit2=Hz n2tic=30 grid=y
       ''')

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

Plot('freqs',
       '''
       graph title=Frequencies yreverse=y min2=0 max2=50
       pad=n wanttitle=n screenratio=0.8 scalebar=n bartype=v
       barwidth=0.2 crowd1=0.75  crowd2=0.3 wherexlabel=b plotfat=2 plotcol=5,4
       label1=Time unit1=s label2=Frequency unit2=Hz n2tic=30 grid=y
       ''')

Plot('true',
     window + '''
     graph title=Frequencies yreverse=y min2=0 max2=50 pad=n wanttitle=n screenratio=0.8 scalebar=n bartype=v
     barwidth=0.2 crowd1=0.75  crowd2=0.3 wherexlabel=b 
     label2=Time unit2=s label1=Frequency unit1=Hz n2tic=30 grid=y dash=1 wantaxis=n
     ''')

Result('freqs','true freqs','Overlay')

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

# analytical trace
Flow('itrace1','sig','envelope order=1000 hilb=y')
Flow('ctrace1','sig itrace1','cmplx ${SOURCES[1]} | ' + window + 'window')

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

Flow('cdif','cfit ctrace1','add scale=1,-1 ${SOURCES[1]}')

Result('cfit','cdif cfit ctrace1',
       '''
       cat axis=2 ${SOURCES[1:3]} | real | dots labels=Difference:Fit:Original gaineach=n
       ''')

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

Result('csign','real | dots gaineach=n label1=Time unit1=s')

Result('sign1','csign','window n2=1 | real  | ' + graph)
Result('sign2','csign','window f2=1 | real  | ' + graph)

End()