from rsf.proj import *
import math, string

#Create regular sinusoid signal with two frequency 

f0 = 1/(2*math.pi)
f1 = 0.2/(2*math.pi)

#Flow('ttrace',None,'math n1=128 d1=1 o1=0 output="0.5*cos(%g*x1+%g)+cos(%g*x1+%g)" ' %(2*math.pi*f0,math.pi*0.1,2*math.pi*f1,math.pi*0.2))
#Result('ttrace','graph title=Data ')

# Make analytical trace

#Flow('thilb','ttrace','envelope hilb=y order=100')
#Flow('tatrace','ttrace thilb','cmplx ${SOURCES[:2]}',stdin=0)
Flow('tatrace',None,
     '''
     math n1=128 d1=1 o1=1 type=complex 
     output="2.5*exp(I*(%g*x1+%g))+5*exp(I*(%g*x1+%g))" |
     dd type=float | noise var=0.05 seed=1234 | dd type=complex
     ''' % (2*math.pi*f0,2*math.pi*f0,
            2*math.pi*f1,0))
Result('tatrace',
       'imag | dots label1=Sample unit1= title="Synthetic Mixed Sinusoidal Signal"')

# Choose a range of frequencies

#Flow('afreq',None,'math n1=100 o1=0 d1=1 output=x1')
#Flow('afreq',None,'spike n1=2 nsp=2 k1=1,2 mag=%g,%g' % (f0,f1))
#Flow('afreq0',None,'spike n1=1 mag=%g' % (f0))
#Flow('afreq1',None,'spike n1=1 mag=%g' % (f1))

# Find frequencies from the data

nf = 2
Flow('freq2',None,'spike n1=2 nsp=2 k1=1,2 mag=%g,%g' % (f0,f1))

Flow('troots','tatrace','cpef nf=%d | roots niter=20' % (nf+1))
Result('troots','graph symbol=x title=Roots')

# Apply freqlet transform

Flow('taft','tatrace troots',
     '''
     freqlet freq=${SOURCES[1]} type=b ncycle=50 perc=95
     ''')

Result('taft',
       '''
       math output="abs(input)" | real |
       dots label1=Scale unit1=  title="1-D Seislet Frame" 
       labels=First:Second yreverse=y 
       ''')
#       put o2=0 d2=1 label2=Frequency unit2=Hz |
#       grey title="Seislet Transform" color=j allpos=y

# Selecting zero frequency is the conventional wavelet transform
#Flow('dwt','aft','window n2=1')

Flow('zero','troots','window n1=1 | math output=1')

Flow('tdwt','tatrace zero',
     '''
     freqlet freq=${SOURCES[1]} type=b  
     ''')

#Flow('tdwt','tatrace','real | dwt type=h unit=y adj=n')
Result('tdwt',
        '''
       math output="abs(input)" | real |
       dots label1=Scale unit1=  title="1-D Wavelet Transform" 
       ''')

Flow('fourier','tatrace','fft3 axis=1')
Result('fourier',
        '''
       math output="abs(input)" | real |
       dots label1=Scale unit1=  title="1-D Fourier Transform" 
       n1=64 d1=0.5 o1=1
       ''')

# Selecting peak frequency

Flow('taft0','taft','window n2=1 f2=0')
Flow('taft1','taft','window n2=1 f2=1')
#Flow('taft2','taft0 taft1','cat axis=2 ${SOURCES[1]}')
#Result('taft2','math output="abs(input)" | real | dots label1=Scale unit1= labels="Seislet-f0":"Seislet-f1"')

# Sort by absolute value and take logarithm
#for case in ('taft0','taft1'):
#    Flow('log'+case,case,'math output="abs(input)" | real | scale axis=1 | sort | math output="log(input)"')
n1 = 128
n2 = 128*nf
p  = 25.0

Flow('logtaft','taft0 taft1',
     '''
     cat axis=1 ${SOURCES[1]} | 
     math output="abs(input)" | 
     real | 
     scale axis=1 | 
     sort | 
     math output="%g*log(input)" |
     window n1=%d j1=%d | 
     put o1=0 d1=%g
     ''' % (10/math.log(10),n2*p/(nf*100),(nf-1),p/(n2*p/100-1)))

for case in ('tdwt','fourier'):
        Flow('log'+case,case,
        '''     
     math output="abs(input)" | 
     real |
     scale axis=1 | 
     sort | 
     math output="%g*log(input)" |
     window n1=%d | 
     put o1=0 d1=%g     
     ''' % (10/math.log(10),n1*p/100,p/(n1*p/100-1)))

Plot('dwt-fourier-aft','logtdwt logfourier logtaft',
     '''
     cat axis=2 ${SOURCES[1:3]} | 
     graph dash=1,2,0 title="Compression Ratio" 
     label1="Percentage (%)" unit1= 
     label2="a\_\s75 n \s100 (dB)" plotcol=7,7,7
     ''' )#10 Log\_\s75 10\^\s100 (a\_\s75 n\^\s100 /a\_\s75 0\^\s100 )
box = '''
      box x0=%g y0=%g label="%s" xt=%g yt=%g
      '''
Plot('ldwt',None,box % (6.55,4.5,"Seislet Frame",-1,-0.5))
Plot('lfourier',None,box % (6,7.85,"Fourier Transform",-1,-0.5))
Plot('laft',None,box % (7.04,8.1,"Wavelet Transform",1,0.5))
Result('tlog','dwt-fourier-aft ldwt laft lfourier','Overlay')

#Result('tlog','logtdwt logtaft0 logtaft1',
#       '''
#       cat axis=2 ${SOURCES[1:3]} | 
#       graph dash=0,1,2 wanttitle=n label1=n unit1= label2="Log(a\_n\^)" unit2=
#       ''')

# Perfect reconstruction

Flow('trecon','taft troots tatrace',
     'freqlet freq=${SOURCES[1]} inv=y type=b | cat axis=2 ${SOURCES[2]}')
Result('trecon','real | graph title=Reconstruction')

for freq in range(nf):
    root = 'troot%d' % freq
    inv = 'tinv%d' % freq

    Flow(root,'troots','window n1=1 f1=%d' % freq)

    Flow(inv,['taft',root],
         '''
         window n2=1 f2=%d |
         freqlet freq=${SOURCES[1]} inv=y type=b
         ''' % freq)
    Result(inv,'real | graph title="Frequency %d" ' % (freq+1))

# Compression

for case in range(3):
    recon = 'trecon%d' % case
    pclip = (5,10,50)[case]

    Flow(recon,'taft troots tatrace',
         '''
         threshold pclip=%g | 
         freqlet freq=${SOURCES[1]} inv=y type=b |
         cat axis=2 ${SOURCES[2]}
         ''' % pclip)
    Result(recon,'real | graph title="2 Sinusoid Reconstruction (%d%%)"' % pclip)

End()