"""This is a helper module to make use of mcmc.py and mcmc_big.py.
It allows you to conveniently run a Monte-Carlo simulation of any
kind until it converges (L{stepper.run_to_bottom}) or until it
has explored a large chunk of parameter space (L{stepper.run_to_ergodic}).

It also helps you with logging the process.
"""

import sys
from mpi4py import MPI		# This uses mpi4py
import random
from gmisclib import die
from gmisclib import Num
from gmisclib import mcmc
from gmisclib import mcmc_helper as MCH
MCH.Debug = True
from gmisclib.mcmc_helper import TooManyLoops, warnevery, logger_template, test

mpi = MPI.COMM_WORLD

class _counter(object):
	def __init__(self):
		self.count = 0

	def update(self, accepted):
		self.count += accepted

	def get(self):
		return self.count



class _accum_erg(object):
	def __init__(self):
		self.erg = 0.0
		self.last_resetid = -1

	def update(self, x):
		"""Increment an accumulator, but reset it to
		zero if the underlying stepper has called reset().
		The increment is by the ergodicity estimate, so
		C{self.erg} should approximate how many times the
		MCMC process had wandered across the probability
		distribution.
		"""
		tmp = x.reset_id()
		if tmp != self.last_resetid:
			self.erg = 0.0
			self.last_resetid = tmp
		else:
			self.erg += x.ergodic()


	def get(self):
		return self.erg


class stepper(MCH.stepper):
	def __init__(self, x, maxloops=-1, logger=None):
		die.info('# mpi stepper rank=%d size=%d' % (mpi.rank, mpi.size))
		assert maxloops == -1
		self.maxloops = -1
		MCH.stepper.__init__(self, x, maxloops, logger)


	def reset_loops(self, maxloops=-1):
		assert maxloops == -1
		self.maxloops = -1


	def communicate_hook(self):
		self.note('chook iter=%d' % self.iter)
		if mpi.size > 1:
			self.note('chook active iter=%d' % self.iter)
			c = self.x.current()
			v = c.vec()
			lp = c.logp()

			if mpi.rank%2 == 0:
				mpi.Send((v, lp), (mpi.rank+1)%mpi.size, tag=self.MPIID)
			nv, nlp = mpi.Recv(tag=self.MPIID)
			if mpi.rank%2 == 1:
				mpi.Send((v, lp), (mpi.rank+1)%mpi.size, tag=self.MPIID)

			if nlp > lp - self.x.T*random.expovariate(1.0):
				self.x._set_current(c.new(nv-v, logp=nlp))
				self.note('communicate succeeded')
			else:
				self.note('communicate not accepted')
			sys.stdout.flush()


	MPIID = 1241

	# def _not_enough_changes(self, n, ncw):
		# mytmp = n < ncw
		# self.note('_nyc pre')
		# ntrue = mpi.Allreduce(int(mytmp), MPI.SUM)
		# self.note('_nyc post')
		# return ntrue*3 >= mpi.size


	def _not_yet_ergodic(self, nc, ncw):
		mytmp = nc.get() < ncw
		self.note('_nye pre')
		ntrue = mpi.Allreduce(int(mytmp), MPI.SUM)
		self.note('_nye post')
		return ntrue*2 >= mpi.size


	def _not_at_bottom(self, xchanged, nchg, es, dotchanged, ndot, update_T):
		mytmp = (Num.sometrue(Num.less(xchanged,nchg))
					or es<1.0 or dotchanged<ndot
					or (update_T and self.x.T>1.5)
				)
		self.note('_nab pre')
		ntrue = mpi.Allreduce(int(mytmp), MPI.SUM)
		self.note('_nab post')
		return ntrue*4 >= mpi.size


	def join(self, id):
		self.note('pre join %s' % id)
		rootid = mpi.Bcast(id)
		assert rootid == id
		self.note('post join %s' % id)


	def note(self, s):
		if MCH.Debug:
			print s, "rank=", mpi.rank
			sys.stdout.flush()


def precompute_logp(lop):
	"""Does a parallel evaluation of logp for all items in lop.
	"""
	nper = len(lop)//mpi.size
	r = mpi.rank
	mychunk = lop[r*nper:(r+1)*nper]
	for p in mychunk:
		p.logp()
	for r in range(mpi.size):
		nc = mpi.Bcast(mychunk, r)
		lop[r*nper:(r+1)*nper] = nc
	mpi.Barrier(0)




def test():
	def test_logp(x, c):
		# print '#', x[0], x[1]
		return -(x[0]-x[1]**2)**2 + 0.001*x[1]**2
	x = mcmc.bootstepper(test_logp, Num.array([0.0,2.0]),
				Num.array([[1.0,0],[0,2.0]]))
	print 'rank=', mpi.rank
	thr = stepper(x)
	# nsteps = thr.run_to_bottom(x)
	# print '#nsteps', nsteps
	# assert nsteps < 100
	for i in range(2):
		print 'RTC'
		thr.run_to_change(2)
		print 'RTE'
		thr.run_to_ergodic(1.0)
		print 'DONE'
	thr.close()


if __name__ == '__main__':
	test()
