from rsf.proj import *
import math
from rsf.prog import RSFROOT
from rsf.recipes.beg import server as private

pi2 = 2*math.pi
dt = 0.001

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))"
     ''' % (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',
       '''
       graph title=Input wanttitle=n labelsz=5 min1=0 max1=1 min2=-4 max2=4
       crowd2=0.3 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 | scale axis=2 | window min2=0 max2=1 |
       grey wanttitle=n color=j scalebar=y  minval=0 maxval=1
       allpos=y screenratio=0.35 label2=Time unit2=s label1=Frequency unit1=Hz grid=y
       wherexlabel=b labelsz=11
       ''')

# time-frequency analysis of st transform

Flow('st-tf', 'sig', 'st | math output="abs(input)" | window n2=601 |real')
Result('st-tf', 'st-tf',
       '''
       transp | scale axis=2 | window min2=0 max2=1 min1=0 max300|
       grey wanttitle=n color=j scalebar=y  minval=0 maxval=1
       allpos=y label2=Time unit2=s label1=Frequency unit1=Hz grid=y
       screenratio=0.35 wherexlabel=b labelsz=11
       ''')

#EEMD

nimf = 6
emds = []
for seed in range(25):
    emd = 'emd%d' % seed
    sn = 'sn%d' % seed
    Flow(sn, 'sig', 'noise var=0.01 seed=%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=1 |
        graph wanttitle=n min2=-4 max2=4 labelsz=10  pad=n
        label2=Amplitude label1= unit1=
        ''')
    imfs.append(imf)

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

# nos-stationary auto-regression

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

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

Plot('cerr',
       '''
       real | window min1=0 max1=1 |
       graph title="Nonstationary Deconvolution" min2=-1.5 max2=1.5
       ''')
Result('clpf_fit', 'cerr cpre ctrace',
       '''
       cat axis=2 ${SOURCES[1:3]} | real | window min1=0 max1=1 |
       dots labels=Difference:Fit:Original gaineach=n
       ''')
# 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

Flow('group', 'croots',
     '''
     math output="-arg(input)/%g" | real
     ''' % (pi2*dt))
Result('hgroup', 'group',
       '''
       graph title=Frequencies yreverse=y pad=n wanttitle=n scalebar=y
       bartype=v plotfat=5 grid=y label2=Frequency unit=Hz min1=0 max1=1
       min2=0 max2=300 plotclol=6,5,4,6,5,4,6,5,4
       screenratio=0.8 crowd1=0.75 crowd2=0.3 labelsz=6
       ''')

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=y bartype=v
       plotfat=5 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)
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
       ''')

# components for nar

Flow('csign', 'comps cwht', 'math other=${SOURCES[1]} output="input*other"')
Flow('nar_emd', 'csign','window n2=6 | real')
ns = nimf
imfs = []
for emd in (4,3,5):
    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=
         lable1=Time unit1=s labelsz=10
         ''')
    imfs.append(imf)

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

# initialization for matlab

matlab         = WhereIs('matlab')
matROOT = '../matfun'
matfun = 'hou_test'
matlabpath = os.environ.get('MATLABPATH',os.path.join(RSFROOT,'lib'))

if not matlab:
    sys.stderr.write('\nCannot find Matlab.\n')
    sys.exit(1)

# test for synthetic trace based on nar

Flow('tfnar1', [os.path.join(matROOT, matfun+'.m'), 'nar_emd'],
     '''
     MATLABPATH=%(matlabpath)s %(matlab)s
     -nosplash -nojvm -r "addpath %(matROOT)s;%(matfun)s('${SOURCES[1]}', '${TARGET}'); quit"
   '''%vars(), stdin=0, stdout=-1)
Flow('tfnar', 'tfnar1', 'put o1=0 d1=0.5 o2=-0.5 d2=0.001 | window min2=0 max2=1')
Result('htfnar', 'tfnar',
       '''
       scale axis=2 | grey wanttitle=n color=j scalebar=y allpos=y
       label2=Time unit2=s label1=Frequency unit1=Hz grid=y
       screenratio=0.35 wherexlabel=b labelsz=11
       ''')

# test for synthetic trace based on emd

Flow('tfemd1', [os.path.join(matROOT, matfun+'.m'), 'emd'],
     '''
     MATLABPATH=%(matlabpath)s %(matlab)s
     -nosplash -nojvm -r "addpath %(matROOT)s;%(matfun)s('${SOURCES[1]}', '${TARGET}'); quit"
   '''%vars(), stdin=0, stdout=-1)
Flow('tfemd', 'tfemd1', 'put o1=0 d1=0.5 o2=-0.5 d2=0.001 | window min2=0 max2=1')
Result('htfemd', 'tfemd',
       '''
       scale axis=2 | grey  wanttitle=n color=j scalebar=y allpos=y
       label2=Time unit2=s label1=Frequency unit1=Hz grid=y
       screenratio=0.35 wherexlabel=b labelsz=11
       ''')

End()