#!/usr/bin/env python

import math
import numpy
from gmisclib import ortho_poly
from gmisclib import die
from gmisclib import gpkmisc

from gpk_voicing import gammatone as GT


class cont_kernel_c(ortho_poly.ortho):
	"""An instance of this class produces a window with a specified number of nodes.
	So, you construct the class (telling it the window size), and that
	gives you a class instance.  Then, you use the class instance (by
	calling the self.P method) it creates a window with the specified
	number of nodes.
	"""
	def __init__(self, n=None, x=None):
		"""@type n: int or None
		@param n: how many points in the window
		"""
		ortho_poly.ortho.__init__(self, n, x)

	def wt_(self):
		return numpy.square(1.0-numpy.square(self.x))

	def compute(self, nnodes):
		"""@param nnodes: Now many zeros does the window have?
		"""
		if nnodes%2 == 0:
			wiggle = numpy.cos((math.pi*(nnodes//2))*self.x)
		else:
			wiggle = numpy.sin((math.pi*((nnodes+1)//2))*self.x)
		return (1.0 + numpy.cos(math.pi*self.x)) * wiggle


_wincache = {}
_WIN_CACHE_SIZE = 20

def window(n, k, norm=2):
	"""Make a window of length n with k nodes.
	The window is normalized so that numpy.sum(w**norm)==1."""

	assert n > 0
	assert k >= 0
	cachekey = (n,k,norm)
	if cachekey in _wincache:
		return _wincache[ cachekey ]

	opn = cont_kernel_c(n)
	# opn.x ranges over (-1, 1)
	w = opn.P(k)
	if norm > 0:
		numpy.divide(w, (numpy.absolute(w)**norm).sum()**(1.0/norm), w)
	while len(_wincache) > _WIN_CACHE_SIZE:
		_wincache.popitem()
	_wincache[ cachekey ] = w
	return w



def test_cont_kernel():
	import pylab
	c = cont_kernel_c(20)
	pylab.plot(c.P(0))
	pylab.plot(c.P(1))
	pylab.plot(c.P(2))
	pylab.plot(c.P(3))
	pylab.show()


def cont_kernel(n, k, norm=2, kernel=window):
	"""Make a window with width n,
	except that it changes continuously when n
	is a floating point number.   The resulting window is
	always an odd length.
	"""
	assert n >= k
	n2 = n/2.0
	nu = 2*int(math.ceil(n2)) + 1
	assert nu > 1
	nl = 2*int(math.floor(n2)) + 1
	assert nu==nl+2 or nu==nl
	if nu==nl or nu<=k:
		w = kernel(nu, k, norm=norm)
	else:
		f = n2 - math.floor(n2)
		w = kernel(nu, k, norm=norm)
		numpy.multiply(w, f, w)
		wl = window(nl, k, norm=norm)
		wul = w.shape[0]
		dl = (wul-wl.shape[0])//2
		assert dl>=0 and wul==wl.shape[0]+2*dl, "dl=%d wul=%d wll=%d" % (dl, wul, wl.shape[0])
		numpy.multiply(wl, 1-f, wl)
		numpy.add(w[dl:wul-dl], wl, w[dl:wul-dl])
	if norm > 0:
		numpy.divide(w, (numpy.absolute(w)**norm).sum()**(1.0/norm), w)
	return w


def xexp(a):
	"""exp(a), except protected from underflows."""
	MIN = -100
	ok = numpy.greater(a, MIN)
	return ok * numpy.exp(numpy.maximum(a, MIN))




def start_end(t, dt, window, len_d):
	"""Find start and end points for a window.
	@param t: desired window center (seconds).
	@type t: float
	@param window: window width (in samples).
	@type window: int
	@param dt: sampling interval (seconds).
	@type dt: float
	@param len_d: length of data array (samples).
	@type len_d: int
	@return (starting sample, end) where end=sample beyond the last
	@rtype (int, int)
	"""
	assert type(window) == type(1)
	s = int(round(t/dt - 0.5*window))
	if s < 0:
		s = 0
	elif s+window > len_d:
		s = len_d - window
	e = s + window
	return (s, e)


def _filter_win_size(n, cutoff_freq):
	np = n + max(10.0/cutoff_freq, n*0.01)
	nx = near_win_size(min=np+1, tol=1+n//11, real=True)
	return nx - nx%2



_fftworkcache = {}
def _fftwork(n):
	"""This estimates the work required to compute a FFT of size N,
	using the Cooley-Tukey algorithm.
	See http://en.wikipedia.org/wiki/Cooley-Tukey_FFT_algorithm
	"""
	try:
		return _fftworkcache[n]
	except KeyError:
		pass
	f1 = gpkmisc.a_factor(n)
	if f1 == n:
		tmp = n**2
	else:
		nof1 = n/f1
		tmp = f1*_fftwork(nof1) + n + nof1*_fftwork(f1)
	_fftworkcache[n] = tmp
	return tmp


def _fftwork_real(n):
	"""This is the work for a forward and reverse transform of
	real data back to real data, using Num.FFT.rfft and Num.FFT.irfft
	"""
	return _fftwork(n) + _fftwork(n-n%2)


def near_win_size(near=None, tol=None, min=None, max=None, real=False):
	"""Find a window size near the specified size where
	the FFT algorithm is fastest."""

	if near is not None and tol is not None:
		nmax = int(round(near + tol))
		nmin = int(round(near - tol))
	elif min is not None and max is not None:
		nmax = int(round(max))
		nmin = int(round(min))
	elif min is not None and tol is not None:
		nmin = int(round(min))
		nmax = int(round(min + tol))
	elif max is not None and tol is not None:
		nmax = int(round(max))
		nmin = int(round(max - tol))
	else:
		raise ValueError, "Needs near&tol, min&max, min&tol, or max&tol."

	if real:
		fftwork = _fftwork_real
	else:
		fftwork = _fftwork

	besti = None
	bestwork = None
	Since_lim = 100
	since = 0
	for i in frommid_generator(nmin, nmax):
		t = fftwork(i)
		if bestwork is None or t < bestwork:
			bestwork = t
			besti = i
			since = 0
		else:
			# This branch is just to keep near_win_size
			# from trying too hard.   Otherwise, it's quite
			# possible to spend more time finding the optimal
			# size than one spends taking the actual FFT.
			since += 1
			if since > Since_lim:
				break
	return besti


def frommid_generator(s, e):
	"""Generate numbers in [s,e] starting at the midpoint."""
	mid = (s+e)//2
	yield mid
	for i in range(1, min(mid-s, e-mid)+1):
		if mid+i <= e:
			yield mid + i
		if mid-i >= s:
			yield mid - i

def near_win_size_real(near=None, tol=None, min=None, max=None):
	return near_win_size(near, tol, min, max, real=True)


def test_win_size():
	assert near_win_size(min=4, max=7) == 4
	assert near_win_size(min=5, max=7) != 7
	assert near_win_size(near=15, tol=2) != 17
	assert near_win_size(near=14, tol=1) != 13
	assert near_win_size(near=97, tol=1) != 97
	assert near_win_size(near=197, tol=1) != 197
	assert near_win_size(near=23, tol=1) != 23


def test_print_win():
	N = 100
	w0 = window(N, 0)
	w1 = window(N, 1)
	w2 = window(N, 2)
	w3 = window(N, 3)
	w4 = window(N, 4)
	w5 = window(N, 5)
	for i in range(N):
		print i, w0[i], w1[i], w2[i], w3[i], w4[i], w5[i]



# def hipass_first_order(d, cutoff_freq):
	# """A first-order high-pass filter.
	# Cutoff freq is measured in cycles per point."""
# 
	# assert cutoff_freq >= 0.0
	# assert len(d.shape) == 1
	# n = d.shape[0]
	# nx = _filter_win_size(n, cutoff_freq)
	# dx = Num.FFT.rfft(d, nx)
	# f = real_fft_freq(dx.shape[0])
	# fr = 1j* (f / cutoff_freq)
	# ff = fr / (1 + fr)
	# ff[0] = ff[0].real
	# ff[-1] = ff[-1].real
	# return Num.FFT.irfft(dx*ff, nx)[:n]
# 
# 
# def lopass_first_order(d, cutoff_freq):
	# """A first-order low-pass filter.
	# Cutoff freq is measured in cycles per point."""
# 
	# assert cutoff_freq >= 0.0
	# assert len(d.shape) == 1
	# n = d.shape[0]
	# nx = _filter_win_size(n, cutoff_freq)
	# dx = Num.FFT.rfft(d, nx)
	# f = real_fft_freq(dx.shape[0])
	# fr = 1j* (f / cutoff_freq)
	# ff = 1 / (1 + fr)
	# ff[0] = ff[0].real
	# ff[-1] = ff[-1].real
	# return Num.FFT.irfft(dx*ff, nx)[:n]

hipass_first_order = GT.hipass_first_order
lopass_first_order = GT.lopass_first_order


def bandpass_first_order(d, cutoff_hp, cutoff_lp):
	"""A combinatio of a first-order high-pass filter
	and a low-pass filter.
	Cutoff freq is measured in cycles per point."""

	assert cutoff_lp >= 0.0
	assert cutoff_hp >= 0.0
	assert len(d.shape) == 1
	n = d.shape[0]
	nx = _filter_win_size(n, 1.0/math.hypot(1.0/cutoff_lp, 1.0/cutoff_hp))

	dx = numpy.fft.rfft(d, nx)
	f = real_fft_freq(dx.shape[0])

	frlp = 1j* (f / cutoff_lp)
	frhp = 1j* (f / cutoff_hp)
	ff = frhp / ((1 + frhp)*(1+frlp))
	ff[0] = ff[0].real
	ff[-1] = ff[-1].real
	return numpy.fft.irfft(dx*ff, nx)[:n]


def test_hpf():
	q = numpy.cos(numpy.arange(10000)*0.003*2*math.pi)
	sq = numpy.sum(q*q)
	q1 = hipass_first_order(q, 0.0003)
	sq1 = numpy.sum(q1*q1)
	assert 0.98*sq < sq1 < 0.995*sq
	q4 = hipass_first_order(q, 0.003)
	sq4 = numpy.sum(q4*q4)
	assert 0.4*sq < sq4 < 0.6*sq
	q5 = hipass_first_order(q, 0.03)
	sq5 = numpy.sum(q5*q5)
	# print sq5/sq
	assert  0.003 < sq5 < 0.03*sq




def fft_freq(n, d=1.0):
	"""Returns an array of indices whose frequencies are between f0 and f1
	after you've called fft().
	This will produce negative values where appropriate.
	"""
	if hasattr(numpy.fft, 'fftfreq'):
		return numpy.fft.fftfreq(n, d)
	fidx = numpy.arange(n)
	fneg = fidx - n
	useneg = numpy.less(-fneg, fidx)
	df = 0.5/(n*d)
	return df * numpy.where(useneg, fneg, fidx)


def fft_indices(f0, f1, n, d=1.0):
	"""Returns an array of indices whose frequencies are between f0 and f1
	after you've called real_fft().
	"""
	raise RuntimeError, "Sorry! Not implemented"


def real_fft_freq(n, d=1.0):
	"""Returns the frequencies associated with an index, after you've
	called real_fft().
	@param d: The time interval between samples.
	@type d: float
	@param n: the size of the transformed array (i.e. frequency domain, complex)
	@type n: int
	@rtype: numpy.ndarray
	@return: the frequency associated with each element in a Fourier Transform
	"""
	fidx = numpy.arange(n)
	df = 0.5/(n*d)
	return df * fidx


def real_fft_indices(f0, f1, n, d=1.0):
	"""Returns the indices whose frequencies are between f0 and f1
	after you've called real_fft().
	"""
	df = 0.5/(n*d)
	f0 = max(0.0, f0)
	imin = int(math.ceil(f0/df))
	imax = int(math.floor(f1/df))
	return (imin, imax+1)



def dataprep_flat_real(data, dt, edgewidth, pad=None):
	return dataprep_flat_generic(data, dt, edgewidth, pad, True)

def dataprep_flat(data, dt, edgewidth, pad=None):
	return dataprep_flat_generic(data, dt, edgewidth, pad, False)


def dataprep_flat_generic(data, dt, edgewidth, pad, isreal):
	"""Pad the data, subtract the average, and round
	the edges of the window.
	"""
	assert len(data.shape) == 1
	m = data.shape[0]
	if pad is None:
		pad = 2*edgewidth + m//10
	min_win_size = m + pad
	n = near_win_size(min=min_win_size, tol=m//11 + 1, real=isreal)
	d = numpy.zeros((n,))
	avg = numpy.average(data)
	d[:m] = data - avg
	nedge = int(round(edgewidth/dt))
	assert nedge <= m
	ew = 0.5*(1.0-numpy.cos((math.pi/nedge)*(numpy.arange(nedge)+0.5)))
	numpy.multiply(d[:nedge], ew, d[:nedge])
	numpy.multiply(d[m-nedge:m], ew[::-1], d[m-nedge:m])
	return d


def dataprep2_flat_generic(data, dt, edgewidth, pad, isreal):
	"""Pad the data, subtract the average, and round
	the edges of the window.
	"""
	assert len(data.shape) == 2
	m = data.shape[0]
	if pad is None:
		pad = 2*edgewidth + m//10
	min_win_size = m + pad
	n = near_win_size(min=min_win_size, tol=m//11 + 1, real=isreal)
	d = numpy.zeros((n,data.shape[1]))
	avg = numpy.average(data, axis=0)
	d[:m,:] = data - avg
	nedge = int(round(edgewidth/dt))
	assert nedge <= m
	ew = 0.5*(1.0-numpy.cos((math.pi/nedge)*(numpy.arange(nedge)+0.5)))
	numpy.multiply(d[:nedge], ew[:,numpy.newaxis], d[:nedge])
	numpy.multiply(d[m-nedge:m], ew[::-1,numpy.newaxis], d[m-nedge:m])
	return d


def lowpass_sym_butterworth(d, cutoff_freq, order=4):
	"""A time-symmetric filter with a magnitude response that
	is the same as the Butterworth filter.
	Cutoff freq is measured in cycles per point."""
	assert len(d.shape) == 1
	if cutoff_freq<=0 or cutoff_freq>10.0:
		die.warn('Silly value of cutoff_freq: %g cycles per point.' % cutoff_freq)
	n = d.shape[0]
	nx = _filter_win_size(n, cutoff_freq)
	dx = numpy.fft.rfft(d, nx)
	f = real_fft_freq(dx.shape[0])
	ff = numpy.sqrt(1.0/(1 + (f/cutoff_freq)**(2*order)))
	return numpy.fft.irfft(dx*ff, nx)[:n]




def lowpass_sym_Gaussian(d, cutoff_freq):
	"""A time-symmetric filter with a Gaussian impulse response.
	Cutoff freq is measured in cycles per point.
	"""
	assert len(d.shape) == 1
	if cutoff_freq<=0 or cutoff_freq>10.0:
		die.warn('Silly value of cutoff_freq: %g cycles per point.' % cutoff_freq)
	n = d.shape[0]
	nx = _filter_win_size(n, cutoff_freq)
	dx = numpy.fft.rfft(d, nx)
	f = real_fft_freq(dx.shape[0])
	q = 1.0/(2*cutoff_freq**2)
	ff = numpy.exp(-numpy.square(f) * q)
	return numpy.fft.irfft(dx*ff, nx)[:n]

lowpass_sym_gaussian = lowpass_sym_Gaussian



def hipass_sym_butterworth(d, cutoff_freq, order=4):
	"""A time-symmetric filter with a magnitude response that
	is the same as the Butterworth filter.
	Cutoff freq is measured in cycles per point."""
	assert len(d.shape) == 1
	if cutoff_freq<=0 or cutoff_freq>10.0:
		die.warn('Silly value of cutoff_freq: %g cycles per point.' % cutoff_freq)
	n = d.shape[0]
	nx = _filter_win_size(n, cutoff_freq)
	dx = numpy.fft.rfft(d, nx)
	f = real_fft_freq(dx.shape[0])
	ff1 = (f/cutoff_freq)**(2*order)
	ff = numpy.sqrt( ff1 / (1 + ff1) )
	return numpy.fft.irfft(dx*ff, nx)[:n]


def test_lpb():
	q = numpy.cos(numpy.arange(10000)*0.003*2*math.pi)
	sq = numpy.sum(q*q)
	q1 = lowpass_sym_butterworth(q, 0.1)
	sq1 = numpy.sum(q1*q1)
	# print sq1/sq
	assert 0.99*sq < sq1 < 1.00001*sq
	q2 = lowpass_sym_butterworth(q, 0.05)
	sq2 = numpy.sum(q2*q2)
	# print sq2/sq
	assert 0.99*sq < sq2 < 1.00001*sq
	q4 = lowpass_sym_butterworth(q, 0.003)
	sq4 = numpy.sum(q4*q4)
	# print sq4/sq
	assert 0.4*sq < sq4 < 0.6*sq
	q5 = lowpass_sym_butterworth(q, 0.001)
	sq5 = numpy.sum(q5*q5)
	# print sq5/sq
	assert  0 < sq5 < 0.2*sq
	

def test_hpb():
	q = numpy.cos(numpy.arange(10000)*0.003*2*math.pi)
	sq = numpy.sum(q*q)
	q1 = hipass_sym_butterworth(q, 0.0003)
	sq1 = numpy.sum(q1*q1)
	# print sq1/sq
	assert 0.99*sq < sq1 < 1.00001*sq
	q4 = hipass_sym_butterworth(q, 0.003)
	sq4 = numpy.sum(q4*q4)
	# print sq4/sq
	assert 0.4*sq < sq4 < 0.6*sq
	q5 = hipass_sym_butterworth(q, 0.03)
	sq5 = numpy.sum(q5*q5)
	# print sq5/sq
	assert  0 < sq5 < 0.2*sq


def pink_noise(n):
	"""Create pink noise where the power spectral density
	goes as 1/f (except right at f=0).
	"""
	m = near_win_size_real(near=2*n+2, tol=n//3+1)

	rtmp = numpy.random.standard_normal(m)
	itmp = numpy.random.standard_normal(m)
	tmp = rtmp + 1j*itmp
	f = real_fft_freq(m)
	assert f[0] == 0.0
	tmp[0] = 0.0
	f[0] = 1.0
	numpy.multiply(tmp, 1.0/numpy.sqrt( f ), tmp)
	return numpy.fft.inverse_real_fft(tmp)[:n]

if __name__ == '__main__':
	test_win_size()
	test_print_win()
	test_hpf()
	test_lpb()
	test_hpb()
	# test_cont_kernel()
