from scipy.signal import get_window
from numpy import *
import _extend

__all__ = ["stft"]



def stft(x, dt, dTime, p,extmethod='none', extlength='powerof2',res = 1000, fmin = 0, fmax = inf):
    """shor time fourier Tranform.

    :Parameters:   
      x : 1d numpy array
        data
      dt : float
         time step
      dTime : float
         time resolution (smaller values of dTime give finer resolution)
      p : float
        wavelet function parameter
      extmethod : string ('none', 'reflection', 'periodic', 'zeros')
                indicates which extension method to use
      extlength : string ('powerof2', 'double')
                indicates how to determinate the length of the extended data
      res:   int
           computed  radial number of the pixels
      fmin, fmax: float   range where fft will be computed
            
    :Returns:
      (X, scales) : (2d numpy array complex, 1d numpy array float)
                  transformed data, scales

    Example:

    >>> import numpy as np
    >>> import mlpy
    >>> x = np.array([1,2,3,4,3,2,1,0])
    >>> stft(x=x, dt=1, dTime=2, wf='dog', p=2)
    (array([[ -4.66713159e-02 -6.66133815e-16j,
             -3.05311332e-16 +2.77555756e-16j,
              4.66713159e-02 +1.38777878e-16j,
              6.94959463e-01 -8.60422844e-16j,
              4.66713159e-02 +6.66133815e-16j,
              3.05311332e-16 -2.77555756e-16j,
             -4.66713159e-02 -1.38777878e-16j,
             -6.94959463e-01 +8.60422844e-16j],
           [ -2.66685280e+00 +2.44249065e-15j,
             -1.77635684e-15 -4.44089210e-16j,
              2.66685280e+00 -3.10862447e-15j,
              3.77202823e+00 -8.88178420e-16j,
              2.66685280e+00 -2.44249065e-15j,
              1.77635684e-15 +4.44089210e-16j,
             -2.66685280e+00 +3.10862447e-15j,
             -3.77202823e+00 +8.88178420e-16j]]), array([ 0.50329212,  2.01316848]))
    """

    x -=  mean(x)
    length = x.shape[0]
    
    if extmethod != 'none':
        x = _extend.extend(x, method=extmethod, length=extlength)
   
    n = x.shape[0]
    freq = fft.fftfreq(n, d=dt)
    #print fmin,fmax
    ind = (freq[:n/2] >fmin )&(freq[:n/2] <fmax)

    step = floor(1/(dt/dTime))

    spec = zeros((res, length/step),dtype=complex)

    for i in range(size(spec,1)):
        #print i 
        s = zeros(size(x))
        #speed improvement
        minimum = max(0,floor(i*step-3*p))
        maximum = min(n,floor(i*step+3*p))
        
        win = get_window('hamming', maximum-minimum, fftbins=False)
        s[minimum:maximum] = x[minimum:maximum]*exp(-(arange(minimum,maximum)-i*step)**2/double(p)**2)


        s = fft.rfft(s)
        s = s[:n/2]
        s = s[ind]
        #print shape(s), shape(s[:(len(s)/res)*res].reshape(res,len(s)/res).mean(1)), res,len(s)/res, size(spec,1)
        spec[:,i] = s[:(len(s)/res)*res].reshape(res,len(s)/res).mean(1)


    scale = fft.fftfreq(n, d=dt)[:(n+1)/2]
    scale = scale[ind]
    
    
    
    return spec, scale[::(n/2)/res]