import cPickle

import die
import mcmc

import mcmc_indexclass as IC
import mcmc_newlogger as LGI


HUGE = 1e38

def setup(pd_factory, arglist, uid='UID', psn_factory=mcmc.position_repeatable, adjust_prms=None):
	"""This and prepares to restart a mcmc.py optimization
	from a log file.
	@param arglist: the list of arguments to the optimization routine (i.e. the linux command line).
	@param adjust_prms: a function to make arbitrary changes to the parameters before you restart.
	@type adjust_prms: None or C{function(}L{mcmc_indexclass.index_base<IC.index_base>}, L{mcmc.problem_definition}, L{mcmc_indexclass.index<IC.index>}C{) ->} L{mcmc_indexclass.index_base<IC.index_base>}.
		The parameters go into the first argument.  The definition of the new problem is the second argument.
		The third argument is a random indexer for the new problem.
		(This is provided mostly so you can conveniently find out what parameters are expected.)
	@type arglist: list(str)
	@note: This assumes that an optimization does not require any fixed parameters.
	@note: Don't mix -restart and -Restart.
	@note: The -Restart flag will be treated as -restart if you attempt to adjust parameters.
	"""
	argliste = list(arglist)
	ll_ll = []
	hdrs = {}
	while arglist and arglist[0] in ['-restart', '-Restart']:
		assert len(arglist)>=2
		# Note: this will give you the wrong answer if you mix -Restart and -restart.
		trust_logp = adjust_prms is None and arglist[0]=='-Restart'
		hdr, lll = LGI.read_tail_uid(arglist[1], uid, Nsamp=-1)
		ll_ll.extend(lll)
		hdrs.update(hdr)
		arglist = arglist[2:]
	if 'Argv' in hdrs:
		argliste = cPickle.loads(hdrs['Argv']) + arglist

	ll_ll.reverse()

	pd = pd_factory(list(argliste), hdrs)
	guess = IC.guess(pd.PriorProbDist)
		# This is deferred initialization for problem_def:
	tries = 0
	while True:
		try:
			pd.logp_guts(guess)
		except mcmc.NotGoodPosition, ex:
			die.warn("Initial position is not good: %s" % str(ex))
			tries += 1
		else:
			break
		if tries > 20:
			die.die("Cannot find a good starting parameter in %d tries." % tries)
	die.info("Adjustable parameters: %s" % ' '.join(guess.columns()))
	n = 4 + guess.n()
	pd.set_idxr( IC.guess_to_indexer(guess) )

	def no_adjustment(x, j1, j2):
		return x
	
	if adjust_prms is None:
		adjust_prms = no_adjustment

	def lop_gen(lll, guess):
		"""Yield some initial values to seed the iteration."""
		for ll in lll:
			ic = IC.reindex(adjust_prms(ll.idxr(), pd, pd.idxr), pd.idxr)
			# ...at this point, we have lost the fixed parameters.
			if trust_logp:
				pos = psn_factory(ic.get_prms(), pd, logp=ll.logp())
			else:
				pos = psn_factory(ic.get_prms(), pd)
			yield pos
		while True:
			yield psn_factory(guess.get_prms(), pd)

	# Delete the oldest and the worst.
	while len(ll_ll) > n:
		del ll_ll[-1]
		if trust_logp:
			smi = -1
			smlp = HUGE
			for (i, ll) in enumerate(ll_ll):
				if ll.logp() < smlp:
					smlp = ll.logp()
					smi = i
			if smi >= 0:
				del ll_ll[smi]
	return (argliste, pd, lop_gen(ll_ll, guess), len(ll_ll))

