from rsf.proj import*
import math

nt=250
pad=50
dt = 0.004
wf = 2*math.pi

Flow('sig0',None,
    '''
    math n1=%d o1=%g d1=%g n2=1 output="cos(%g*10*x1)+cos(%g*30*x1)"
    ''' % (nt+2*pad, -pad*dt, dt, wf, wf))
Flow('sig','sig0','math output=input')

Flow('true1','sig','math output="10" ')
Flow('true2','sig','math output="30"  ')


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


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

Result('sig','Overlay')


Flow('tf','sig','timefreq rect=17')

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','shift1 ns=2')

# find adaptive PEF

# analytical trace
Flow('itrace','sig','envelope hilb=y')
Flow('ctrace','sig itrace','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]}')

Plot('cerr',
     window + '''     
     real |
     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 + 'real | ' + 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',
     window + '''
     math output="-arg(input)/%g" | real | clip2 lower=0  
     ''' % (wf*dt))

Result('group',
       '''
       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',
     '''
     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) )

Flow('ctrace1','ctrace',window + 'window')

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

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

Result('cfit','cresid 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 | graph gaineach=n label1=Time unit1=s')

Plot('sig',window + graph)
Result('sig1','csign','window n2=1 | real  | ' + graph )
Result('sig2','csign','window n2=1 f2=1 | real  | ' + graph)
Result('cresid','cresid',' real  | ' + graph)

#Result('sigsep','sig sign1 sign2 cresid','OverUnderAniso')



###################################################################################
### FOR EMD
###################################################################################

Flow('sigemd','sig','emd | pad n2=3')

Result('sigemd2','sigemd','window n2=1 | '+window + graph )
Result('sigemd1','sigemd','window n2=1 f2=1 | '+window + graph)
Result('sigemd3','sigemd','window n2=1 f2=2 | '+window + graph)





End()