import ortho_poly
import Num
import die
import math
import gpk_lsq
from oxigen_exceptions import NotEnoughData

SING_VAL_RATIO_LIM = 1e-6


class BadDataSupport(RuntimeError):
	def __init__(self, s=""):
		RuntimeError.__init__(self, "Support for fit is poor: %s" % s)


class solution:
	def __init__(self, polys, coefs, err=None, rank=None, fit=None):
		self.p = polys
		self.c = coefs
		self.e = err
		self.rank = rank
		if fit is not None:
			self._fit = Num.array(fit)
		else:
			self._fit = None

	def __str__(self):
		return '<ortho_fit_solution %s %s>' % (
				getattr(self.p, 'name', '???'),
				['%g' % q for q in self.c]
				)

	__repr__ = __str__

	def fit(self):
		if self._fit is None:
			self._fit = self.p.expand(self.c)
		return self._fit


	def expand(self, time):
		sum = Num.zeros(time.shape, Num.Float)
		for i in range(len(self.c)):
			# print "OPPS=", op.P(i).shape, "SS=", soln.shape
			# print "I=", i, "op.P=", op.P(i), "SOLN=", soln[i]
			sum += self.p.P(i) * self.c[i]
		return sum


def perplexity(x):
	p = x / Num.sum(x)
	return math.exp(-Num.sum(p * Num.log(p)))





def poly(x=None, n=None, name='Chebyshev2'):
	return ortho_poly.F(name, x=x, n=n)


def ortho_fit_svd(x, y, wt, N, svr_lim=None, svrls=None, name='Chebyshev2'):
	"""This uses a SVD fit to handle ill-constrained fits."""
	assert N >= 0
	assert N < 100
	EPS = 1e-8
	die.note('y.shape', y.shape)
	die.note('len(y)', len(y))
	die.note('N', N)
	assert Num.alltrue(Num.absolute(x)<=1+EPS), "Time must be normalized into [-1,1]"
	# die.note('x', x)
	# die.note('y', y)
	if len(y) < N:
		raise NotEnoughData, "y"
	op = ortho_poly.F(name, x=x)
	wwt = op.wt() * wt
	rwwt = Num.sqrt(wwt)
	if N == 0:
		return ( solution(op, Num.zeros((0,), Num.Float)),
				Num.sum(wwt), None
			)

	assert len(wwt.shape) == 1
	yw = y * rwwt
	assert len(yw.shape)==1
	die.note('yw.shape', yw.shape)
	die.note('len(yw)', len(yw))
	# print "CSC=", yw
	if Num.sum( wwt > 1e-10 ) <= N:
		raise NotEnoughData, "wt"
	basisvec = []
	for i in range(N):
		# sys.stderr.writelines('v')
		basisvec.append( op.P(i) * rwwt )
	basisw = Num.transpose(basisvec)
	assert len(basisw.shape) == 2
	# print "b.shape=", b.shape, "y.shape=", y.shape
	# print "bc=", bc
	# print "yw=", yw
	if svrls is None:
		# If you are fitting a linear equation with support at
		# (0,0)*N, (1,0)*N, (1,1), (i.e. there are N point to
		# determine the first coefficient, and only one point that
		# determines the second coefficient, then the singular values
		# are sqrt(N) and 1, so the limit to the ratio of singular
		# values is set to be a little smaller than that. 
		svrls = 0.5/math.sqrt(len(y))
		# So, if the rank is limited because if this, the equations
		# must be both somewhat singular and poorly supported.
	# print 'yw=', yw
	# print 'basisw=', basisw
	raise Exception, "Next line is obsolete!"
	soln, fitw, rank, q = gpk_lsq.linear_least_squares(basisw, yw, svrls)
	# Gotta be careful here: fitw is the fit to the data multiplied
	# by the weight.  It is *not* the fit to the input data.

	assert len(soln) == N
	# print 'soln=', soln
	# print 'rank=', rank
	# print 'q=', q
	# print 'svrls=', svrls
	try:
		svr = min(q) / max(q)
	except ZeroDivisionError:
		die.warn('Largest singular value is apparently zero!')
		raise

	if svr_lim is None:
		svr_lim = SING_VAL_RATIO_LIM
	if svr < svr_lim:
		# print 'q=', q
		# print 'wt=', wt
		# print 'wwt=', wwt
		# print 'basisw=', basisw
		raise BadDataSupport, 'Ortho_fit is nearly degenerate (svr=%g::%g)' % (svr, svr_lim)

	assert len(soln) == N
	swt = Num.sum(wwt)
	soln = solution(op, soln, err=Num.sum((fitw-yw)**2)/swt, rank=rank)
	return (soln, swt, svr)


ortho_fit = ortho_fit_svd


def ortho_fit_reg(x, y, wt, N, regstr, regtgt,
			rscale=None, name='Chebyshev2'):
	"""This uses linear regularization to handle ill-constrained fits."""
	assert N >= 0
	assert N < 100
	EPS = 1e-8
	die.note('y.shape', y.shape)
	die.note('len(y)', len(y))
	die.note('N', N)
	assert Num.alltrue(Num.absolute(x)<=1+EPS), "Time must be normalized into [-1,1]"
	# die.note('x', x)
	# die.note('y', y)
	if len(y) < N:
		raise NotEnoughData, "y"
	op = ortho_poly.F(name, x=x)
	wwt = op.wt() * wt
	rwwt = Num.sqrt(wwt)
	if N == 0:
		raise Exception, "something wrong here -- inconsistent return values"
		return ( solution(op, Num.zeros((0,), Num.Float)),
				Num.sum(wwt), None
			)

	assert len(wwt.shape) == 1
	yw = y * rwwt
	assert len(yw.shape)==1
	die.note('yw.shape', yw.shape)
	die.note('len(yw)', len(yw))
	# print "CSC=", yw
	if Num.sum( wwt > 1e-10 ) <= N:
		raise NotEnoughData, "wt"
	basisvec = []
	for i in range(N):
		# sys.stderr.writelines('v')
		basisvec.append( op.P(i) * rwwt )
	basisw = Num.transpose(basisvec)
	assert len(basisw.shape) == 2
	# print "b.shape=", b.shape, "y.shape=", y.shape
	# print "bc=", bc
	# print "yw=", yw
	# print 'basisw=', basisw
	rlss = gpk_lsq.reg_linear_least_squares(basisw, yw, regstr,
						regtgt, rscale=rscale)
	# Gotta be careful here: this is the fit to the data multiplied
	# by the weight.  It is *not* the fit to the input data.

	assert len(rlss.x()) == N
	swt = Num.sum(wwt)
	rank = rlss.eff_rank()
	soln = solution(op, rlss.x(), err=Num.sum((rlss.fit()-yw)**2)/swt, rank=rank)
	# raise Exception, "something wrong here -- inconsistent return values"
	return (soln, swt)



def ortho_err_rms(soln, x, data, wt):
	wwt = soln.p.wt() * wt
	r = data - soln.expand(x)
	return Num.sum(wwt * r**2)/Num.sum(wwt)


def ortho_err_abs(soln, x, data, wt):
	wwt = soln.p.wt() * wt
	r = data - soln.expand(x)
	return Num.sum(wwt*Num.absolute(r))/Num.sum(wwt)







def eff_perplexity(wt, correlation_length, threshold = 0.2, name='Chebyshev2'):
	assert len(wt.shape) == 1
	n = wt.shape[0]
	assert wt.shape == (n,)
	op = ortho_poly.F(name, n)
	x = []
	sqwt = Num.sqrt(wt * op.wt())
	# print 'EP:', sqwt
	for i in range(n-1):
		x.append( op.P(i) * sqwt )
	x = Num.array(x, Num.Float)
	xx = Num.matrixmultiply(x, Num.transpose(x))
	assert xx.shape == (n-1,n-1), "xx.shape=%s, n=%d" % (str(xx.shape), n)
	ev = Num.LA.Heigenvalues(xx)
	# print 'Eff_perp ev:', ev
	mxev = ev[Num.argmax(ev)]
	# print 'Eff_perp mxev:', mxev, correlation_length, threshold**2 * mxev, Num.sum(ev > threshold**2 * mxev)
	return Num.sum(ev > threshold**2 * mxev) / math.hypot(correlation_length, 1.0)



def test_r():
	import RandomArray
	N = 2
	M = 10
	wt1 = RandomArray.exponential(1.0, (M,))
	wt = Num.concatenate((wt1, wt1[::-1]))
	# wt = Num.ones((2*M,), Num.Float)*2
	y1 = RandomArray.exponential(1.0, (M,))
	y = Num.concatenate((y1, y1[::-1]))
	ya = Num.sum(y*wt)/Num.sum(wt)
	y /= ya
	# Y is constructed to have a mean of 1 and a slope of zero.
	x = (Num.arrayrange(2*M)-M+0.5)/float(2*M)
	soln, swt, svr = ortho_fit(x, y, wt, N, svr_lim=None, svrls=None, name='Legendre')
	assert svr>0.5 and svr<=1.0
	assert abs(swt - 2*M) < 5*math.sqrt(2*M)
	assert abs(soln.c[0]-1) < 1e-6
	assert abs(soln.c[1]) < 1e-6


def test_expand():
	import RandomArray
	N = 5
	M = 60
	x = (Num.arrayrange(2*M)-M+0.5)/float(2*M)
	op = ortho_poly.Legendre(x=x)
	y = op.P(0) + op.P(1) - op.P(2) - op.P(3)
	wt = RandomArray.exponential(1.0, (2*M,))
	soln, swt, svr = ortho_fit(x, y, wt, N, svr_lim=None, svrls=None, name='Legendre')
	assert Num.alltrue( Num.absolute(soln.c - [1, 1, -1, -1, 0]) < 1e-6)
	assert svr > 0.5 and svr<=1.0
	assert abs(swt - Num.sum(wt)) < 1e-6
	fitdiff = soln.fit() - y
	assert Num.sum(fitdiff**2) < 0.001
	diff = soln.expand(x) - y
	assert Num.sum(diff**2) < 0.001


if __name__ == '__main__':
	test_r()
	test_expand()
