#!/usr/bin/python

"""Description
"""

import sys
print "Path=", sys.path
from gmisclib import die
from gmisclib import mcmc
from gmisclib import mcmc_helper as mcmcP
import bsr_analysis_guts as AG

mcmcP.Debug = 1
mcmc.Debug = 1

class logger_c(mcmcP.logger_template):
	"""This class is called from the Markov-Chain stepper
	and writes the log file.
	"""
	HUGE = 1e30

	def __init__(self, iterstep=100, ergstep=0.1):
		self.__g_implements__ = ['logger_template']
		assert iterstep >= 0
		assert ergstep >= 0.0
		self.iterstep = iterstep
		self.ergstep = ergstep
		self.BestLogP = -1e30
		self.nlogged = 0
		self.erg = 0.0
		self.last_add_iter = -1
		self.last_logged_erg = -self.HUGE
		self.last_logged_iter = -1
		self.was_reset = False


	def set_logstep(self, iterstep, ergstep):
		assert iterstep > 0
		assert ergstep is None or ergstep >= 0
		self.iterstep = iterstep
		self.ergstep = ergstep


	def add(self, stepperInstance, iter):
		"""Called every step: this decides whether
		data should be logged.
		"""
		if self.was_reset:
			stepperInstance.current().pd.reset_outlist()
			self.was_reset = False
		self.erg += stepperInstance.ergodic() * (iter-self.last_add_iter)
		self.last_add_iter = iter
		if(iter>self.last_logged_iter+self.iterstep
				and (self.ergstep is None or self.erg>self.last_logged_erg+self.ergstep)):
			self.Add(stepperInstance, iter)
			cur = stepperInstance.current()
			if cur.logp() > self.BestLogP:
				cur.pd.max_fcn(cur.prms())
				self.BestLogP = cur.logp()


	def Add(self, stepperInstance, iter):
		"""Actually write a set of parameters into the log file."""
		print 'Logging', stepperInstance.current().logp(), 'iter=', iter
		sys.stdout.flush()
		self.nlogged += 1
		self.last_logged_erg = self.erg
		self.last_logged_iter = iter
		cur = stepperInstance.current()
		cur.pd.output_fcn(cur.prms())


	def close(self):
		pass

	def comment(self, comment):
		pass

	def header(self, k, v):
		pass

	def headers(self, hdict):
		pass


	def reset(self):
		self.nlogged = 0
		self.erg = 0.0
		self.last_logged_erg = -self.HUGE
		self.was_reset = True





def run(arglist):
	froot = arglist[0]
	if len(arglist) != 1:
		die.die("Usage: bsr_analysis.py filename_root")
	ifilename = froot + ".mcmci"
	opt_filename = froot + ".mcmco"
	speedfac = 1
	pd = AG.problem_def(ifilename, opt_filename)
	print "There will be %d samples taken." % pd.requested_samples()

	lov = []
	lop = []
	# Guess some initial values to seed the iteration:
	for x in pd.guesses():
		try:
			lop.append(mcmc.position_repeatable(x, pd))
			lov.append(x)
		except mcmc.NotGoodPosition:
			die.warn("Wproject reports a bad initial position")
	# Initialize the Markov-Chain Monte-Carlo:
	v = mcmc.diag_variance(lov)
	st = mcmc.BootStepper(lop, v)

	logger = logger_c(iterstep=0)
	sh = mcmcP.stepper(st, logger=logger)

	# Find the bottom, roughly:
	acc = mcmcP.step_acceptor(sh.x.np, T0=1000.0, T=1.0)
	sh.run_to_bottom(ns=5+15/speedfac, acceptable_step=acc)
	logger.Add(sh.x, sh.iter)	# Log the (hopeful) minimum.

	# Run for a while to let the temperature equilibriate:
	sh.x.set_sort_strategy(sh.x.SSNEVER)
	for i in range((20*sh.x.np)/(5*speedfac)+1):
		sh.run_to_change(5, acceptable_step=mcmc.T_acceptor())
	sh.run_to_ergodic(max(2.0, 10.0/speedfac), T=1.0)	# Make sure the temperature is cold.
	logger.Add(sh.x, sh.iter)

	# Set up to sample from the probability distribution:
	logstep = 1+5//speedfac
	logger.set_logstep(logstep, 0.1)
	ERGSTEP = 1.0
	nsamp = pd.requested_samples()
	print "Beginning %d samples." % nsamp
	while logger.nlogged < nsamp:
		print "Sampling", logger.nlogged
		sh.run_to_ergodic(ERGSTEP, T=1.0)	# Collect the data.
		print "Sampled", logger.nlogged
	print "Closing pd from python"
	pd.close()



if __name__ == '__main__':
	arglist = sys.argv[1:]
	print 'Arglist:', arglist
	run(arglist)
