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

nt = 1124
dt = 0.01

Flow('s1', None, 'math n1=%d d1=%g o1=-1 output="0.3*cos(%g*x1)" | cut min1=6' % (nt, dt, 10*pi))
Flow('s21', None, 'math n1=%d d1=%g o1=-1 output="0.8*cos(%g*x1)" | cut min1=6' % (nt, dt, 30*pi))
Flow('s22', None, 'math n1=%d d1=%g o1=-1 output="0.7*cos(%g*x1+sin(%g*x1))" | cut max1=6' % (nt, dt, 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 crowd2=0.25 pad=n labelsz=5
     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 | scale axis=1 | 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):
    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[:6], '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]}')

Result('cdecon', 'cerr cpre ctrace',
       '''
       cat axis=2 ${SOURCES[1:3]} | real |
       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')

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

Result('group',
       '''
       graph title=Frequencies yreverse=y pad=n wanttitle=n scalebar=y
       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=y
       bartype=v plotfat=5 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"')
Flow('nar_emd', 'csign', 'window n2=5 | real')

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

Flow('nar_emd_par', [imfs[1], imfs[3], imfs[4]], 'cat axis=2 ${SOURCES[1:3]}')


nimf3 = 3
imfs3 = []
for emd in range(nimf3):
    imf = 'nar_emd_par%d' % emd
    Flow(imf, 'nar_emd_par', 'window n2=1 f2=%d' % emd)
    Plot(imf,
         '''
         window min1=0 max1=10 |
         graph wanttitle=n min2=-1 max2=1 labels2=Amplitude unit2=
         label1= unit1= wantxlabel=n labelsz=15
         ''')
    imfs3.append(imf)
Result('nar_emd_par', imfs3, 'OverUnderAniso')

# mask low amplitude

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=y 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')

# initiation for matlab

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

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

# test for systhetic 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 o2=-1 d2=%g o1=0 d1=%g | window min2=0 max2=10 ' %(11.23/1123,40./400))
Result('mtfemd', 'tfemd',
       '''
       scale axis=2 | 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
       ''')


# test for systhetic trace based on emd

Flow('tfnar1', [os.path.join(matROOT, matfun+'.m'), 'nar_emd_par'],
     '''
     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 o2=-1 d2=%g o1=0 d1=%g | window min2=0 max2=10 ' %(11.23/1123,40./400))

Result('mtfnar', 'tfnar',
       '''
       scale axis=2 | 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
       ''')




End()