
"""An extension of mcmc that includes new stepping algorithms.
"""

from __future__ import with_statement

import math
import random
import numpy

import die
import mcmc
from mcmc import Debug
import gpk_lsq

SIGFAC = 3.0

NotGoodPosition = mcmc.NotGoodPosition

def N_maximum(a):
	return a[numpy.argmax(a)].item()


def _parab_interp_guts(xylist):
	ND = 3
	assert len(xylist) >= 3

	yy = numpy.zeros((len(xylist), 1))
	a = numpy.zeros((len(xylist), ND))
	for (i, (x, y)) in enumerate(xylist):
		# Order of parameter is a Constant, linear slope, then quadratic.
		yy.itemset(i, 0, y)
		a.itemset(i, 0, 1.0)
		a.itemset(i, 1, x)
		a.itemset(i, 2, x*x)

	lls = gpk_lsq.linear_least_squares(a, yy, copy=False)

	try:
		sv = lls.sv()
		x = lls.x()
	except (ValueError, numpy.linalg.linalg.LinAlgError), ex:
		raise mcmc.NoBoot, "stepParab: %s in _compute_move" % str(ex)
	if not (sv[-1] > 1e-6*sv[0]):
		raise mcmc.NoBoot, "Sv ratio too extreme: %g/%g" % (sv[-1], sv[0])
	c = x.item(0,0)
	b = x.item(1,0)
	a = x.item(2,0)
	# print 'c=', c
	# print 'cc=', cc
	return (a, b, c)


def _parab_interp(points):
	xlist = [x for (x, pos) in points]
	tmp = [pos.logp_nocompute() for (x, pos) in points]
	tmax = max(tmp)
	ylist = [q-tmax for q in tmp]
	# print 'XY=', (x0, y0), (x1, y1), (x2, y2)

	# print 'xy=', (x0, y0), (x1, y1), (x2, y2)
	a, b, c = _parab_interp_guts(zip(xlist, ylist))
	# print 'abc=', a, b, c
	mn = min(xlist)
	mx = max(xlist)
	w = mx - mn
	if a >= 0.0:
		raise mcmc.NoBoot, "parab: positive curvature"
	else:
		# This looks like a nice parabola with a maximum.   We step
		# into the vicinity of the maximum.
		xmin = -b/(2*a)
		sigma = math.sqrt(-SIGFAC/(2*a))

	TOOCLOSE = 0.02
	TOOFAR = 2.0
	if xmin < -TOOFAR*w:
		xmin = -TOOFAR*w
	elif xmin > TOOFAR*w:
		xmin = TOOFAR*w
	if sigma > TOOFAR*w:
		sigma = TOOFAR*w
	elif sigma < 2*TOOCLOSE*w:
		sigma = 2*TOOCLOSE*w

	xtmp = random.normalvariate(xmin, sigma)
	while min([abs(xtmp-x) for x in xlist]) < TOOCLOSE*w:
		xtmp = random.normalvariate(xtmp, sigma)
		# The minimum is too nearly a replication of an existing point,
		# so we try again to make sure we get some new information.
	return xtmp


def test():
	a, b, c = _parab_interp_guts([(0.0, 1.0), (1.0, 0.0), (2.0, 1.0)])
	assert abs(a-1.0) < 0.001
	assert abs(-b/(2*a) - 1.0) < 0.001
	assert abs(c-1.0) < 0.001



def find_closest_p(v, vp):
	"""Searches a list of (v,p) pairs and finds the one whose v
	is closest to the first argument.   Returns
	(v,p) of the closest pair.
	"""
	assert len(vp) > 0
	vc, pc = vp[0]
	for (vi, pi) in vp[1:]:
		if abs(vi-v) < abs(vc-v):
			vc = vi
			pc = pi
	return (vc, pc)



def _fast_adjust(moveB, moveV, vfac, aV, aB, accepted):
	"""@note: this breaks the Markov assumption."""
	SLOP = 5.0
	# If the moveB completely dominates moveV, then we can apply our knowledge of
	# whether the step was accepted or not to adjusting the size of the Bootstrap step.
	if numpy.greater(numpy.absolute(moveB), SLOP*numpy.absolute(moveV)).all():
		aB.inctry(accepted)
	# Likewise if moveV dominates:
	if numpy.greater(numpy.absolute(moveV), SLOP*numpy.absolute(moveB)).all():
		aV.inctry(accepted)
	# Otherwise, we assume that the bootstrap steps are fairly successful,
	# so the bootstrap size is approximately correct, and we
	# adjust the stepV step size to approximately match the stepBoot step size.
	if numpy.greater(numpy.absolute(moveB), SLOP*numpy.absolute(moveV)/vfac).all():
		aV.inctry(1)	# moveV is too small: pretend stepV succeeded so it gets bigger.
	if numpy.greater(numpy.absolute(moveV)/vfac, SLOP*numpy.absolute(moveB)).all():
		aV.inctry(0)	# moveV is too big: pretend stepV failed so it gets smaller.



def _has_enough(points, dP):
	if len(points) < 3:
		return False
	if len(points) > 10:	# Something is wrong.   Stop wasting time on a parabolic approximation.
		return True
	lp = [pos.logp_nocompute() for (x, pos) in points]
	spread = max(lp) - min(lp)
	return spread > 3*dP/math.sqrt(len(points)-2)


class BootStepper(mcmc.BootStepper):
	def __init__(self, lop, v, strategy=mcmc.BootStepper.SSAUTO,
				maxArchSize=None, parallelSizeDiv=1):
		mcmc.BootStepper.__init__(self, lop, v, strategy=strategy,
						maxArchSize=maxArchSize,
						parallelSizeDiv=parallelSizeDiv)


	def step(self):
		mcmc.stepper.step(self)
		Wboot = max(self.archive.distinct_count()-2, 0)
		WV = self.np
		Wmixed = math.sqrt(Wboot * WV)
		# Parabolic steps are a bootstrap method, so they needs a large archive.
		# But, Step_parab is not really Markov, so we only want to do it 
		if self.archive.sorted:
			Wparab = max(self.archive.distinct_count()-2, 0)
		else:
			Wparab = 0
		Wprobe = (Wboot+WV+Wmixed+1.5*Wparab) * self.F
		W = float(Wboot + Wmixed + Wparab + WV + Wprobe)
		Pboot = Wboot/W
		Pmixed = Pboot + Wmixed/W
		Pparab = Pmixed + Wparab/W
		PV = Pparab + WV/W

		again = True
		while again:
			P = random.random()
			try:
				if P < Pboot:
					accepted = self.step_boot()
				elif P < Pmixed:
					accepted = self.step_mixed()
				elif P < Pparab:
					# Note! This violates MCMC requirements!
					# The resulting probability distribution will not be correct
					# when step_parab() is used!
					accepted = self.step_parab()
				elif P < PV:
					accepted = self.stepV()
				else:
					accepted = self.step_probe()
			except mcmc.NoBoot, x:
				if Debug > 0:
					die.info('NoBoot: %s' % str(x))
				again = True
			else:
				again = False
		return accepted


	def step_mixed(self):
		vfac = math.exp(random.random()*math.log(1e-5))
		self.steptype = 'step_mixed'
		if len(self.archive) <= 2:
			# print 'NoBoot', len(self.archive)
			raise mcmc.NoBoot, "mixed short archive"
		p1 = self.archive.choose()
		p2 = self.archive.choose()
		if p1.uid() == p2.uid():
			# print 'NoBoot dup'
			raise mcmc.NoBoot, "mixed duplicate"

		vsV = self.aV.vs() * vfac
		if self.archive.distinct_count() >= min(5, self.np_eff):
			vsV *= (self.archive.variance()/self.v.diagonal())**(1./4.)

		moveB  = (p1.vec() - p2.vec()) * self.aB.vs()
		moveV  = vsV * self.V.sample()
		move = moveB + moveV
		try:
			tmp = self.current().new(move)
			delta = tmp.logp() - self.current().logp()
		except mcmc.NotGoodPosition:
			die.warn('NotGoodPosition')
			# Should I call self.aM.inctry(accepted) ?
			return 0
		if self.acceptable(delta):
			accepted = 1
			self._set_current(tmp)
			if Debug>2:
				die.info('StepMixed: Accepted logp=%g' % tmp.logp_nocompute())
		else:
			self._set_failed( tmp )
			accepted = 0
			self._set_current(self.current())
			if Debug>2:
				die.info('StepMixed: Rejected logp=%g vs. %g, T=%g'
						% (tmp.logp_nocompute(), self.current().logp_nocompute(), self.acceptable.T()))

		# Only do this after a reset or when we don't care about breaking Markov assumptions:
		if self.archive.sorted:
			_fast_adjust(moveB, moveV, vfac, self.aV, self.aB, accepted)
		return accepted



	def step_parab(self):
		self.steptype = 'step_parab'
		# The parabolic fit will be degenerate if we pick vs0 too close
		# to either of the existing points.
		if len(self.archive) <= 2:
			raise mcmc.NoBoot, "parab short archive"

		p1 = self.current()
		p2 = self.archive.choose()
		if p1.uid() == p2.uid():
			raise mcmc.NoBoot, "parab duplicate"
		points = [(0.0, p1), (1.0, p2)]
		pp = []
		vsB = self.aB.vs()
		while True:
			vs0 = random.normalvariate(0.0, vsB)
			# Don't pick a point too close to the existing points.
			tooclose = False
			for (x, pos) in points:
				if abs(vs0-x) < 0.5/len(points):
					tooclose = True
					break
			if tooclose:
				continue

			vbase, pbase = find_closest_p(vs0, points)
			move  = (p2.vec() - p1.vec()) * (vs0-vbase)
			# This is an approximation to a MCMC step, at least when the
			# archive is large and well-annealed.  So, we might as well add
			# it to the archive.
			try:
				tmp = pbase.new(move)
				pp.append( (tmp, tmp.logp()) )
			except mcmc.NotGoodPosition:
				die.warn('NotGoodPosition2a at %.3f' % vs0)
				return 0

			points.append( (vs0, tmp) )
			if _has_enough(points, self.acceptable.jitter()):
				break
			vsB *= 1.2
		if Debug > 0:
			die.info("step_parab: %d points at %s" % (len(points), [q for (q,pos) in points]))

		# Now, we take one of the (potentially many) new points and treat that as
		# a MCMC step and add it to the archive.   We don't want to add them all,
		# because we might build up strong linear structures in the archive and
		# make it more likely that the optimization gets stuck in a subspace.
		tmp, tmp_logp = random.choice(pp)
		try:
			delta = tmp_logp - self.current().logp()
		except mcmc.NotGoodPosition:
			die.warn('NotGoodPosition2a at %.3f' % vs0)
			return 0
		if self.acceptable(delta):
			if Debug > 2:
				die.info('StepParab Accepted1 at %.3f, logp=%g' % (vs0, tmp_logp))
			self._set_current(tmp)
		else:
			self._set_failed( tmp )
			self._set_current(self.current())
			# Return from here?  No.  It still can be used to
			# define the interpolation parabola, even if logP
			# isn't particularly good.

		try:
			vsnew = _parab_interp(points)
		except mcmc.NotGoodPosition:
			die.warn('NotGoodPosition preparing for _parab_interp().')
			return 0
		vbase, pbase = find_closest_p(vsnew, points)
		move = (p2.vec() - p1.vec()) * (vsnew-vbase)

		try:
			tmp = pbase.new(move)
			delta = tmp.logp() - self.current().logp()
		except mcmc.NotGoodPosition:
			die.warn('NotGoodPosition')
			accepted = 0
		if self.acceptable(delta):
			if Debug > 2:
				die.info('StepParab Accepted2 at %.3f, logp=%g' % (vsnew, tmp.logp_nocompute()))
			accepted = 1
			self._set_current(tmp)
		else:
			self._set_failed( tmp )
			accepted = 0
			self._set_current(self.current())
			if Debug>2:
				die.info('StepParab: Rejected logp=%g vs. %g, T=%g'
						% (tmp.logp_nocompute(), self.current().logp_nocompute(), self.acceptable.T()))
		return accepted




def bootstepper(logp, x, v, c=None, strategy=BootStepper.SSAUTO, fixer=None, repeatable=True):
	"""This is (essentially) another interface to the class constructor.
	It's really there for backwards compatibility.
	"""
	pd = mcmc.problem_definition_F(logp_fcn=logp, c=c, fixer=fixer)
	position_constructor = [mcmc.position_nonrepeatable, mcmc.position_repeatable][repeatable]
	return BootStepper(mcmc.make_list_of_positions(x, position_constructor, pd), v, strategy=strategy)

diag_variance = mcmc.diag_variance
stepper = mcmc.stepper
problem_definition = mcmc.problem_definition
problem_definition_F = mcmc.problem_definition_F
problem_definition = mcmc.problem_definition
position_repeatable = mcmc.position_repeatable
position_nonrepeatable = mcmc.position_nonrepeatable

def test2d(stepper):
	import dictops
	start = numpy.array([9, 1])
	c = numpy.array([1.0, 0.1])
	V = numpy.identity(len(c), numpy.float)
	# V = numpy.array([[ 20.69626808,  20.6904984 ], [ 20.6904984,   20.69477235]])

	x = stepper(mcmc._logp1, start, V, c=c)
	v = numpy.zeros((x.np, x.np))
	n_per_steptype = dictops.dict_of_averages()
	for i in range(1000):
		accepted = x.step()
		n_per_steptype.add(x.steptype, accepted)
	lpsum = 0.0
	lps2 = 0.0
	na = 0
	psum = numpy.zeros((x.np,))
	N = 30000
	nreset = 0
	nsorted = 0
	rid = x.reset_id()
	for i in range(N):
		accepted = x.step()
		na += accepted
		n_per_steptype.add(x.steptype, accepted)
		nsorted += x.archive.sorted
		if x.reset_id() != rid:
			rid = x.reset_id()
			nreset += 1
		lp = x.current().logp()
		lpsum += lp
		lps2 += lp**2
		p = x.current().vec()
		numpy.add(psum, p, psum)
		v += numpy.outer(p, p)
	for steptype in n_per_steptype.keys():
		print 'Step %s has %.0f successes out of %.0f trials: %.2f' % (
			steptype, n_per_steptype.get_both(steptype)[0],
			n_per_steptype.get_both(steptype)[1],
			n_per_steptype.get_avg(steptype)
			)
	# assert N*x.PBootLim**1.5 < nboot < N*x.PBootLim**0.7
	assert nsorted < N//30
	assert N//8 < na < N//2
	assert nreset < 30
	lpsum /= N
	assert abs(lpsum+0.5*x.np) < 0.15
	lps2 /= N-1
	lpvar = lps2 - lpsum**2
	assert abs(lpvar-0.5*x.np) < 1.0
	numpy.divide(psum, N, psum)
	assert numpy.alltrue(numpy.less(numpy.absolute(psum), 6.0/numpy.sqrt(c)))
	numpy.divide(v, N-1, v)
	for i in range(x.np):
		for j in range(x.np):
			if i == j:
				# print '2*v[ii]*c=', 2*v[i,j]*c[i]**2
				assert abs(math.log(2*v[i,j]*c[i]*c[j])) < 0.1
			else:
				assert abs(v[i,j]) < 20/(c[i]*c[j])


if __name__ == '__main__':
	# mcmc.test(stepper=bootstepper)
	Debug = 3
	test()
	mcmc.test_(stepper=bootstepper)
