#!/usr/bin/env python
# -*- coding: utf-8 -*-
#from matplotlib.pyplot import *
from numpy import *
from multiprocessing import Process, Pool, cpu_count
from fftshift import fftshift
#import fftw3
import pyfftw


"""
    Continuous Multi-Wavelet analyzis  0.001
    
    Description:
    The continius wavelet analysis od the multiple signal is calculated. Instead
    of the single one dimensional wavelet a n-dimensional Morlet wavelet with shifted component is created.
    The moving scalar product with such wavelets helps to identify and amplify the  correctly time shifted 
    signals and suppress the other. 



"""




__all__ = ["cwt",]


def angularfreq(N, dt):
    """Compute angular frequencies.

    :Parameters:   
      N : integer
        number of data samples
      dt : float
        time step
    
    :Returns:
      angular frequencies :  1d numpy array
    """

    # See (5) at page 64.
    N2 = N / 2.0
  
    w = fftshift(2*pi*arange(-N2,N2)/ (N * dt))

    
    
    return w


def scales(N, dj, dt, s0):
    """Compute scales.

    :Parameters:
      N : integer
        number of data samples
      dj : float
        scale resolution
      dt : float
        time step

    :Returns:
      scales : 1d numpy array
        scales
    """


    J = floor(dj**-1 * log2((N * dt) / s0))
    s = empty(J + 1)
    
 
    s = s0 * 2**(arange(s.shape[0]) * dj)
   
    return s


def compute_s0(dt, p):
    """Compute s0.
    
    :Parameters:    
      dt :float
        time step
      p : float
        omega0 ('morlet') or order ('paul', 'dog')


    :Returns:    
      s0 : float
    """


    return (dt * (p + sqrt(2 + p**2))) / (2 * pi)
    
   

    
    
    
    
    
  

def cwt((x, dt, dj, p , res, fmin,fmax)):

    """
    x - transformed signal
    dt - time step
    dj - frequency resolution
    p - tradeoff between spatial and time resolution
    res - number of the timepoint of the spectrogram
    fmin,fmax  - frequency range
    
    """
    lenght_ = size(x,0)
    
    #reduce lenght to increase speed of fft
    div = 2**int(log2(lenght_/256))
    lenght = div*(lenght_/div)   
    
   
    x = csingle(x[:lenght])

   
    w = angularfreq(lenght, dt)
    s0 = compute_s0(dt, p)
    scale = scales(lenght, dj, dt, s0)
    freq = (p + sqrt(2.0 + p**2))/(4*pi * scale)	
    
    ind = (freq>fmin)&(freq<fmax)
    scale =scale[ind]
    freq = freq[ind]
    x = fft.fft(x, axis=0)[:lenght/2]
    x[::2]*= -1  #magic :-)

   
    step  = max(int(lenght/res),1)
    hres = len(w[0:lenght:step])
    spec = zeros((len(scale), hres),dtype=csingle)
    
    pyfftw.interfaces.cache.enable()
    pyfftw.interfaces.cache.set_keepalive_time(30)

    
    #wft = fftw3.create_aligned_array(len(w))
    #out = fftw3.create_aligned_array(len(w))
    #fft_backward = fftw3.Plan(wft,out, direction='backward', flags=['estimate',],nthreads=1)
    
    pyfftw.interfaces.cache.enable()
    pyfftw.interfaces.cache.set_keepalive_time(30)
    
    wft = pyfftw.n_byte_align_empty(len(w), 16,dtype=csingle)
    out = pyfftw.n_byte_align_empty(len(w), 16,dtype=csingle)

    w = w[:len(w)/2]


    fft_backward = pyfftw.FFTW(wft,out, direction='FFTW_BACKWARD', flags=['FFTW_ESTIMATE','FFTW_DESTROY_INPUT'])

    for i,s in enumerate(scale):
	
	i2 = min((5+p)/s*(lenght * dt)/(2*pi),len(w))
	i1 = max((p-5)/s*(lenght * dt)/(2*pi),0)
	slc = slice(int(i1),int(i2))

	wft[:] = 0


	wavelet= exp(-(s*w[slc]-p)**2/2)*sqrt(s/dt)/lenght
	wft[slc] = wavelet*x[slc]
	
	fft_backward()
	

	spec[i,:] = fftshift(out[::step])
	

    return spec, scale, freq
    
    

def DFT(x, y,nf):
    
    Ftheta = arange(-nf/2,nf/2)

    E = exp(-1j*outer(Ftheta,x))
    Fx_ = dot(y,conj(linalg.pinv(E)))*len(Ftheta)
    
    return  Fx_
	

    
def main():
    random.seed(1)
    x = random.rand(1e4,16)/1.e5
    tvec = arange(1e4)*1e-6

    x+= sin(tvec**2*2*pi*1e4**2)[:,None]
    #plot(x)
    #show()

    x = DFT(2*pi*arange(16)/16., x,16)
    x = fftshift(x,axes=1)
    #plot(x)
    
    #plot(x[:,0])
    #show()
    dt = 1e-6
    dj = 0.05
    fmin = 0
    w0 = 10
    fmax = 1e6/2
    import time
    t = time.time()
    
    p = Pool(cpu_count())
    out = map(cwt, [ (x[:,i], dt, dj,20 ,500, fmin,fmax) for i in range(16) ])
    p.close()
    p.join()

    spec, s,f  = out[0]
    #spec, s,f = cwt((x[:,0], dt, dj, w0, 500, fmin,fmax))
    print  (time.time()-t)
    exit()

    #imshow(abs(spec),aspect='auto', interpolation='bicubic')
    #colorbar()
    #show()
    #print t
    
    fig = figure()

    ax = fig.add_axes([0.1, 0.1, 0.8, 0.85])
    ax.set_yscale('log', nonposy='clip')
    img = ax.imshow(abs(spec), extent=[tvec[0],tvec[-1] ,f[-1], f[0]],
    aspect='auto',interpolation='bicubic') 
    minorLocator   = MultipleLocator(1)
    ax.xaxis.set_minor_locator(minorLocator)
    ax.axis([tvec[0],tvec[-1] ,amin(f), amax(f)])
    ax.set_xlabel('time [s]')
    ax.set_ylabel('Frequency [Hz]')
    
    show()
    
    
    
    
    
 
if __name__ == "__main__":
    main()