import numpy as np

def movingWindowAnalysis(data,winfun,nwin,shift,exp):
    """
    Performs moving window analysis of a time series.
    data:   data array
    winfun: name of the window function to be called
    nwin:   number of window samples (power of 2)
    shift:  displacement of moving window in samples
    exp:    exponent for taking power of spectrum
    """
    fwin = winfun(nwin)                    # compute window values
    npts = len(data)                       # number of total samples
    nseg = int((npts-nwin)/shift)+1        # total number of expected data segment
    mwa = np.zeros((nwin//2+1,nseg))       # array for result (rfft returns N/2+1 samples)  
    wa = 0                                 # start index of data segment
    we = nwin                              # end index of data segment
    jseg = 0                               # initialize data segment counter
    while we < npts:                       # loop over segments
        seg = data[wa:we]*fwin                     # multiply data segment with window
        seg = seg-seg.mean()                       # subtract mean value of segment
        ftseg = np.abs(np.fft.rfft(seg))           # abs value of Fourier transform
        maxft = np.amax(ftseg)                     # max value of Fourier transform
        ftseg = ftseg/maxft+1.e-10                 # normalize spectrum to its maximum, remove zeros
        mwa[:,jseg] = np.power(ftseg,exp)          # assign values to the matrix
        wa = wa+shift                              # move window start by shift
        we = we+shift                              # move window end by shift
        jseg = jseg+1                              # increase segment counter
    return nseg,mwa                         # return number of segments and moving window matrix
#------------------------------------------------------------------------------------
#
def hann(nw):
    arg = 2.*np.pi*np.arange(0,nw)/nw     # argument of cosine
    fwin = 0.5*(1.-np.cos(arg))           # Hann window
    return fwin