from rsf.proj import *
from math import pi

# Make synthetic signal

Flow('s1', None,'math n1=1124 d1=0.01 o1=-1 output="0.3*cos(%g*x1)" | cut min1=6' % (10*pi))
Flow('s21',None,'math n1=1124 d1=0.01 o1=-1 output="0.8*cos(%g*x1)" | cut min1=6' % (30*pi))
Flow('s22',None,'math n1=1124 d1=0.01 o1=-1 output="0.7*cos(%g*x1+sin(%g*x1))" | cut max1=6' % (20*pi,pi))
Flow('s2','s21 s22','add ${SOURCES[1]}')
Flow('s3','s1','math output="0.4*cos(%g*x1+sin(%g*x1))" | cut max1=4 | cut min1=7.8' % (66*pi,2*pi))

Flow('s','s1 s2 s3','add ${SOURCES[1:3]} | put label1=Time unit1=s')

sigs = []
for s in range(3):
    sig = 's%d' % (s+1)
    Plot(sig,'graph title="Signal %d" min2=-1 max2=1 label2=Amplitude' % (s+1))
    sigs.append(sig)
Result('sig','s3 s2 s1','OverUnderAniso')

Plot('s',
     '''
     window min1=0 max1=10 | 
     graph wanttitle=n min2=-1.5 max2=1.5 screenratio=0.8 crowd1=0.75 crowd2=0.3 pad=n
     label2=Amplitude
     ''')
Result('msig','s','Overlay')

# Time-frequency analysis

Flow('mtf','s','timefreq rect=25 dw=0.1 nw=401')

Result('mtf',
       '''
       window min1=0 max1=10 | transp | 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
       ''')

# EEMD

nimf = 6

emds = []
for seed in range(25):
    sn = 'sn%d' % seed
    Flow(sn,'s','noise var=0.01 seed=%d' % seed)

    emd = 'emd%d' % seed
    Flow(emd,sn,'emd | window n2=%d' % nimf)
    emds.append(emd)

# ensemble average
Flow('emd',emds,'cat axis=3 ${SOURCES[1:%d]} | stack axis=3' % len(emds))

imfs = []
for emd in range(nimf):
    imf = 'imf%d' % emd
    Flow(imf,'emd','window n2=1 f2=%d' % emd)
    Plot(imf,
         '''
         window min1=0 max1=10 | 
         graph wanttitle=n min2=-1 max2=1 labelsz=15 screenratio=0.8 pad=n label2=Amplitude
         label1= unit1=
         ''')
    imfs.append(imf)

Result('emd',imfs[:5],'OverUnderAniso')

# Non-stationary auto-regression

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

# find adaptive PEF

# analytical trace
Flow('itrace','s','envelope hilb=y')
Flow('ctrace','s 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=25 pred=${TARGETS[1]}')
Flow('cerr','cpre ctrace','add scale=-1,1 ${SOURCES[1]}')

Plot('cerr',
     '''     
     real |
     graph title="Nonstationary Deconvolution" min2=-1.5 max2=1.5
     ''')

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

# 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.01

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

Result('group',
       '''
       graph title=Frequencies yreverse=y pad=n wanttitle=n scalebar=n bartype=v
       plotfat=3 grid=y label2=Frequency unit2=Hz min2=0 max2=40
       ''')

# Filter out bad frequencies

Flow('gmask','group','stack axis=1 | mask min=0')
Flow('wgroup','group gmask','headerwindow mask=${SOURCES[1]}')

# From group to phase

Flow('freqs','wgroup',
     '''
     causint | math output="input*%g/(x1+1+%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 min2=0 max2=40
       ''')

# Fitting components by non-stationary regression

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

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

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

Result('cfit','cdif cfit ctrace',
       '''
       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"')

nimf2 = 5

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

Result('narall',imfs,'OverUnderAniso')
Result('nar',[imfs[1],imfs[3],imfs[4]],'OverUnderAniso')

# mask low amplitudes

Flow('mask','cwht','math output="abs(input)" | real | mask min=0.2')

group = []
for emd in (1,3,4):
    mask = 'mask%d' % emd
    Flow(mask,'mask','window n2=1 f2=%d' % emd)
    grup = 'group%d' % emd
    Flow(grup,['wgroup',mask],
         '''
         window n2=1 f2=%d | 
         rtoc | math output="x1+I*input" | 
         transp plane=12 | headerwindow mask=${SOURCES[1]} | window
         ''' % emd)
    Plot(grup,
         '''
         graph title=Frequencies yreverse=y pad=n wanttitle=n scalebar=n bartype=v
         plotfat=5 plotcol=%d grid=%d label2=Frequency unit2=Hz min2=0 max2=40 min1=0 max1=10
         screenratio=0.8 crowd1=0.75 crowd2=0.3
         ''' % (7-emd,emd==1))
    group.append(grup)

Result('mgroup',group,'Overlay')

# extend boundaries

End()