from rsf.proj import *
import math

pi2 = 2*math.pi

Flow('sig1',None,
     '''
     math n1=2001 o1=-0.5 d1=0.001 output="cos(%g*x1+15*sin(%g*x1))/(1.5+cos(%g*x1))"
     label1=Time unit1=s
     ''' % (30*pi2,pi2,pi2))
Flow('sig2','sig1',
     '''
     math output="cos(%g*x1+sin(%g*x1))/(1.5+sin(%g*x1))"
     ''' % (80*pi2,8*pi2,pi2))
Flow('sig3','sig1',
     '''
     math output="(2+cos(%g*x1))*cos(%g*(x1+1)*(x1+1))"
     ''' % (4*pi2,70*pi2))

Flow('sig','sig1 sig2 sig3','add ${SOURCES[1:3]}')

Result('hsig','sig',
       '''
       window min1=0 max1=1 |
       graph title=Input wanttitle=n screenratio=0.8 crowd1=0.75 crowd2=0.3
       labelsz=6 pad=n label2=Amplitude
       ''')

sigs = []
for s in range(3):
    sig = 'sig%d' % (s+1)
    Plot(sig,'graph wanttitle=n min1=0 max1=1 min2=-4 max2=4 label2=Amplitude labelsz=10')
    sigs.append(sig)
Result('sigs','sig3 sig2 sig1','OverUnderAniso')

# Time-frequency analysis

Flow('timefreq','sig','timefreq rect=25 dw=0.5 nw=601')

Result('htf','timefreq',
       '''
       transp | window min2=0 max2=1 |
       grey wanttitle=n color=j scalebar=n bartype=v allpos=y
       label2=Time unit2=s label1=Frequency unit1=Hz grid=y
       screenratio=0.8 crowd1=0.75 crowd2=0.3 wherexlabel=b labelsz=6
       ''')

Flow('emd','sig','emd')

# Non-stationary auto-regression

Flow('shift','sig','shift1 ns=8')

# 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=20 pred=${TARGETS[1]}')
Flow('cerr','cpre ctrace','add scale=-1,1 ${SOURCES[1]}')

Result('cerr',
     '''     
     real | window min1=0 max1=1 |
     graph title="Nonstationary Deconvolution" min2=-1.5 max2=1.5
     ''')

# 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=r | transp')

# Frequency components

dt = 0.001

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

Result('hgroup','group',
       '''
       graph title=Frequencies yreverse=y pad=n
       wanttitle=n scalebar=n bartype=v
       plotfat=3 grid=y label2=Frequency unit2=Hz 
       min1=0 max1=1 min2=0 max2=300 plotcol=6,5,4,6,5,4,6,5,4
       screenratio=0.8 crowd1=0.75 crowd2=0.3 labelsz=6
       ''')

# From group to phase

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

Result('freqs',
       '''
       graph title=Frequencies yreverse=y pad=n wanttitle=n scalebar=n bartype=v
       plotfat=3 grid=y label2=Frequency unit2=Hz 
       min1=0 max1=1 min2=0 max2=300
       ''')

# Fitting components by non-stationary regression

Flow('comps','freqs',
     'rtoc | math output="exp(I*input*(x1+0.5)*%g)" ' % pi2 )
s
Flow('cwht cfit','comps ctrace',
     'clpf match=${SOURCES[1]} rect1=40 pred=${TARGETS[1]}')

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

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

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

imfs = []
for emd in (5,6,7):
    imf = 'nar%d' % emd
    Flow(imf,'csign','window n2=1 f2=%d | real' % emd)
    Plot(imf,'graph wanttitle=n min1=0 max1=1 min2=-4 max2=4 label2=Amplitude unit2= label1=Time unit1=s labelsz=10')
    imfs.append(imf)

Result('nar',imfs,'OverUnderAniso')

# Extend to avoid boundary conditions


End()