# -*- coding: utf-8 -*- 
#Nicolas Enfon - LSIS DYNI - NortekMED
from pylab import *
from time import *
from scipy.io.wavfile import *
import pywt

def scalogram(data, wavelet='db2', level=6, order='freq', normalized=False):
    ''' Computes and return a scalogram with PyWavelets

	@data: input data. Can be the path to the raw .wav file, or the .wav file itself
	@wavelet: PyWavelets wavelet type
	@level: level of decomposition. Level=6 gives 2 ^ 6 = 64 rows
	@order: pecifies nodes order - natural (natural) or frequency (freq)
	@normalized: normalization of the scalogram: scalo = scalo / norm(scalo)

	returns: a scalogram'''

    if type(data) == str:
	sample_rate, data = read(data)

    wp = pywt.WaveletPacket(data, wavelet, 'sym', maxlevel=level)
    nodes = wp.get_level(level=level, order=order)
    values = array([n.data for n in nodes], 'd')
    values = abs(values)
    #flip it so that the low frequencies are in the last rows of the matrix
    values = values[::-1]
    if normalized:
	values = values / norm(values)
	#normalization could make it easier for threshold setting,
	#but would alter the distance information

    return values