import Num
import math
import die


"""Implements a set of orthonormal polynomials, based on a recurrence relation
on the order of the polynomial.    Each particular type of polynomial just
needs to specify that recurrence relation.
"""


def _orthogonalize(x, others, wt):
	"""Makes x orthonormal to all others.
	Others are assumed to be orthonormal to eachother."""
	against = others[:]
	against.reverse()
	ipass = 0
	x = Num.array(x, Num.Float, copy=True)
	n = x.shape[0]
	while ipass <= 2:
		ipass += 1
		do_all_again = False
		do_again = []
		renorm = True
		for t in against:
			if renorm > 0:
				Num.multiply(x, 1.0/math.sqrt(Num.average( wt * x**2)), x)
				renorm = False
			d = Num.average( x * wt * t)
			# print 'i=', i, 'd=', d
			if abs(d) >= 1e-6:
				do_again.append(t)
				renorm = True

			if abs(d) > 0.1*n:
				do_all_again = True
				if ipass > 1 :
					die.note('x', x)
					die.note('t', t)
					die.note('wt', wt)
					die.note('d', d)
					die.warn('Poor orthogality')
					raise RuntimeError, "Poor orthogonality: %d %g %d/%d" % (len(wt), d, ipass, len(others))
			assert x.shape == t.shape
			# assert type(d) == type(0.1)
			Num.subtract(x, t*d, x)

		if do_all_again:
			against = others
		elif do_again:
			against = do_again
		else:
			break
	if ipass > 2:
		raise RuntimeError, "Can't orthogonalize"
	return x * (1.0/math.sqrt(Num.average(wt*x**2)))


class ortho:
	registry = {}

	def __init__(self, n=None, x=None):
		"""Argument n is how many points you want the polynomials evaluated at.
		The ordinates of the points are placed in self.x.
		You may alternatively specify the points as argument x, in which case
		the x_() function will never be called."""
		if x is not None:
			lx = len(x)
			if not lx > 0:
				raise ValueError, "Need len(x)>0"
			if n is not None:
				assert int(n) == len(x), "X and N mismatch: %d vs %s" % (len(x), n)
			self.x = Num.array(x, Num.Float, copy=True)
		elif n is not None:
			n = int(n)
			if not n>0:
				raise ValueError, "Need n>0"
			self.x = self.x_(n)
		else:
			raise RuntimeError, "Need to specify either x or n."

		self.o = []
		self._wt = None
		self._numwt = None


	def wt(self):
		"""Weighting function to get orthonormality.
		This sets self._numwt."""
		if self._wt is None:
			self._wt = self.wt_()
			self._numwt = Num.sum( self._wt > 0.0 )
			if not Num.alltrue( self._wt >= 0.0 ):
				raise ValueError, 'Weights must be non-negative.'
			if not self._numwt > 0:
				raise ValueError, 'Need some positive weights.'
		# Again, it might make sense to return a copy, as _wt gets used
		# internally after being passed out...
		return self._wt


	def x_(self, m):
		"""Calculates the points at which the function is evaluated,
			if you want m evenly spaced points.
			By default, the function is assumed to
			range over the open interval (-1, 1).
			You can override this function in a subclass
			to get a different range."""
		return (Num.arrayrange(m)-0.5*(m-1))*(2.0/m)


	def P(self, i):
		"""Self.P(i) is the ith orthogonal polynomial.
		The result is normalized so that
		Num.sum(self.wt() * self.P(i)**2) == 1.
		It returns an un-shared array, so the array can be
		modified without affecting future calls to P(i).
		"""
		assert i>=0
		if i<len(self.o):
			return Num.array(self.o[i], copy=True)
		wt = self.wt()	# Sets self._numwt
		if i >= self._numwt:
			raise ValueError, 'Request for higher order polynomial than the number of positive weight points.'
		while i >= len(self.o):
			self.o.append(_orthogonalize(self.compute(len(self.o)), self.o, wt))
		# We return a copy, because it is possible that the user
		# might modify the data.  If so, it would screw up the orthogalization
		# step for higher polynomials.  Copy-on-write would be nice.
		return Num.array(self.o[-1], Num.Float, copy=True)


	def expand(self, c):
		o = Num.zeros(self.x.shape, Num.Float)
		for (i, ci) in enumerate(c):
			tmp = self.P(i)
			Num.multiply(tmp, ci, tmp)
			Num.add(o, tmp, o)
		return o

	def compute(self, n):
		raise RuntimeError, 'Virtual Function'


	def wt_(self):
		raise RuntimeError, "Virtual Function"


class ortho_poly(ortho):
	__doc__ = """Virtual base class.
		This is a superclass of all orthorgonal polynomials.
		This does the recurrence [Abramowitz and Stegun p.782],
		and ensures that the generated polynomials are all
		orthonormal to each other.
		The derived classed need to specify two functions:
		recurse() and wt_().
		Recurse() is the recursion relation from P(i) and P(i-1)
		to P(i+1), and
		wt_() is the weighting function for the polynomial.
		Wt_() is needed to check orthogonality, and really
		defines the polynomial.
		These polynomials are guarenteed to be
		orthogonal to eachother when summed over the
		supplied set of x points.   Thus, if you change
		x_() or wt_(), you get a different set of functions."""


	def __init__(self, n=None, x=None):
		ortho.__init__(self, n, x)


	def P(self, i):
		"""Self.P(i) is the ith orthogonal polynomial.
		The result is normalized so that
		Num.sum(self.P(i)**2) == 1.  """

		assert i>=0
		if i<len(self.o):
			return Num.array(self.o[i], copy=True)
		wt = self.wt()	# Sets self._numwt
		if i >= self._numwt:
			raise ValueError, 'Request for higher order polynomial than the number of positive weight points.'
		while i>=len(self.o):
			if len(self.o)>1:
				m2 = self.o[-2]
			else:
				m2 = None
			if len(self.o)>0:
				m1 = self.o[-1]
			else:
				m1 = None

			tmp = self.recurse(len(self.o), m1, m2)
			self.o.append(_orthogonalize(tmp, self.o, wt))
		# We return a copy, because it is possible that the user
		# might modify the data.  If so, it would screw up the orthogalization
		# step for higher polynomials.  Copy-on-write would be nice.
		return Num.array(self.o[-1], Num.Float, copy=True)


	def recurse(self, n, fn, fnm1):
		"""Does the recurrence relation, evaluating f_(n+1) as a function
		of n, f_n, and f_(n-1).   For n=0 and n=1, fn and fnm1 may be None."""
		raise RuntimeError, "Virtual function"

	def compute(self, n):
		raise RuntimeError, "Compute is not needed and should not be used for subclasses of ortho_poly."


class Legendre(ortho_poly):
	__doc__ = """Legendre polynomials,
		orthonormal over (-1, 1) with weight 1.
		Called Pn(x) in Abramowitz and Stegun.
		"""

	name = 'Legendre'

	def __init__(self, n=None, x=None):
		ortho_poly.__init__(self, n, x)


	def recurse(self, n, fn, fnm1):
		if fn is None:
			assert n == 0
			return Num.ones(self.x.shape, Num.Float)
		if fnm1 is None:
			assert n == 1
			return Num.array(self.x, Num.Float, copy=True)
		# tmp1 = (2*n-1)/float(n)
		# tmp2 = (n-1)/float(n)
		# return tmp1*self.x*fn - tmp2*fnm1
		return self.x*fn - 0.08*fnm1



	def wt_(self):
		return Num.ones(self.x.shape, Num.Float)

ortho.registry[Legendre.name] = Legendre


class Chebyshev(ortho_poly):
	__doc__ = """Chebyshev polynomials of the first kind,
		orthonormal over (-1, 1) with weight (1-x^2)^(-1/2).
		These are the equi-ripple polynomials.
		Called Tn(x) in Abramowitz and Stegun.
		"""

	name = 'Chebyshev'

	def __init__(self, n=None, x=None):
		ortho_poly.__init__(self, n, x)


	def recurse(self, n, fn, fnm1):
		if fn is None:
			assert n == 0
			return Num.ones(self.x.shape, Num.Float)
		if fnm1 is None:
			assert n == 1
			return Num.array(self.x, Num.Float, copy=True)
		# return 2*self.x*fn - fnm1
		return self.x*fn - 0.052*fnm1


	def wt_(self):
		return Num.sqrt(1/(1.0-self.x**2))

ortho.registry[Chebyshev.name] = Chebyshev


class Chebyshev2(ortho_poly):
	__doc__ = """Chebyshev polynomials of the second kind,
		orthonormal over (-1, 1) with weight (1-x^2)^(1/2).
		Called Un(x) in Abramowitz and Stegun.
		"""

	name = 'Chebyshev2'

	def __init__(self, n=None, x=None):
		ortho_poly.__init__(self, n, x)


	def recurse(self, n, fn, fnm1):
		if fn is None:
			assert n == 0
			return Num.ones(self.x.shape, Num.Float)
		if fnm1 is None:
			assert n == 1
			return 2*self.x
		# return 2*self.x*fn - fnm1
		return self.x*fn - 0.445*fnm1


	def wt_(self):
		return Num.sqrt(1.0-self.x**2)

ortho.registry[Chebyshev2.name] = Chebyshev2


class SinCos(ortho):
	__doc__ = """1, sin(pi*x), cos(pi*x), sin(2*pi*x), cos(2*pi*x)...
		orthonormal over (-1, 1) with weight 1."""

	name = 'SinCos'

	def __init__(self, n=None, x=None):
		ortho.__init__(self, n, x)


	def compute(self, n):
		# print 'SinCos.recurse', n
		if n == 0:
			return Num.ones(self.x.shape, Num.Float)
		elif n%2 == 1:
			return Num.sin((math.pi*(n+1)/2) * self.x)
		return Num.cos((math.pi*(n/2))*self.x)


	def wt_(self):
		return Num.ones(self.x.shape, Num.Float)

ortho.registry[SinCos.name] = SinCos





class SLTB(ortho):
	__doc__ = """Smooth Local Trigonometric Basis.
		From Bj\:orn Jawerth, Yi Liu, Wim Sweldens,
		"Signal Compression with Smooth Local Basis Functions".
		"""

	name = 'SLTB'

	def __init__(self, n=None, x=None, eps=0.5):
		assert eps <= 0.5 and eps>0.0
		self.eps = eps
		ortho.__init__(self, n, x)


	def compute(self, n):
		omega = (2*n+1)*math.pi/2.0
		s = Num.sin(omega * self.x)
		return s * self.b(self.x)


	def r(self, x):
		"""This must have l**2+r**2==1"""
		x = Num.clip(x, -self.eps, self.eps)
		s = Num.sin(x * math.pi/(2.0*self.eps))
		r = 1 + s
		l = 1 - s
		return r/Num.hypot(r, l)


	def b(self, x):
		# print 'x=', x
		# print 'b=', self.l(x)*self.l(1.0-x)
		return self.r(x) * self.r(1.0-x)


	def wt_(self):
		return Num.ones(self.x.shape, Num.Float)


	def x_(self, m):
		"""Calculates the points at which the function is evaluated,
			if you want m evenly spaced points.
			The function is assumed to
			range over the open interval (-eps, 1+eps).
			You can override this function in a subclass
			to get a different range."""
		return (0.5+Num.arrayrange(m))*((1+2*self.eps)/m) - self.eps


	def test():
		M = 20
		N = 3*M
		x = (Num.arrayrange(N)+0.5)/float(N) * 3.0 - 0.5
		q = SLTB(x=x)
		for i in range(M//3):
			for j in range(M//3):
				a = q.P(i)
				b = q.P(j)
				# print Num.dot(a[:-M], b[M:])
				assert abs(Num.dot(a[:-M], b[M:])) < 0.001

	test = staticmethod(test)

ortho.registry[SLTB.name] = SLTB



def test(name):
	N = 96
	q = F(name, N)
	for i in range(N//2):
		for j in range(10):
			dot = Num.sum(q.P(i)*q.P(j)*q.wt())
			# print 'p(i=', i, ')', q.P(i)
			# print 'p(j=', j, ')', q.P(j)
			if i != j:
				assert abs(dot) < 0.001
			else:
				assert abs(dot - N) < 0.001*N

	if hasattr(q, 'test'):
		q.test()











def F(name, n=None, x=None):
	"""This is a factory function.  Get any kind of ortho poly
	you want, as selected by the first argument."""

	try:
		return ortho.registry[name](n, x)
	except KeyError:
		raise RuntimeError, "Unimplemented orthogonal function name: %s" % name




if __name__ == '__main__':
	for name in ortho.registry.keys():
		test(name)
	# q = F('SLTB', n=100)
	# for i in range(100):
		# print i, q.P(0)[i]
