
import math
import numpy

import die
import fiatio
import gpkmisc
import mcmc_logger
import mcmc_newlogger as LG
import mcmc_indexclass as IC



FILE_DROP_FAC = 0.2
TRIGGER = 'run_to_bottom finished'


def read_many_files(fnames, uid, Nsamp, tail, trigger):
	"""OBSOLETE
	@return: a dictionary mapping filenames onto "data" and header information.
		Header information is from the last file read.
		The "data" are L{onelog} class instances that encapsulate
		log information from one log file.
	@rtype: tuple(dict(str:onelog), dict(str:str))
	"""
	hdr = {}
	per_fn = {}
	if uid is None:
		for fname in fnames:
			hdr, indexers, logps = mcmc_logger.read_multisample(fname, Nsamp=Nsamp, tail=tail, trigger=trigger)
			per_fn[fname] = onelog(logps, indexers, fname)
	else:
		for fname in fnames:
			try:
				hdr, indexers, logps = LG.read_multi_uid(fname, uid,
									Nsamp=Nsamp, tail=tail,
									trigger=trigger
									)
			except LG.BadFormatError, x:
				try:
					hdr, statedict, logps = LG.read_human_fmt(fname)
				except LG.BadFormatError, y:
					raise LG.BadFormatError, "Unreadable in either format: {%s} or {%s}" % (x, y)
				try:
					indexers = [statedict[uid]]
					logps = [logps[uid]]
				except KeyError:
					raise LG.BadFormatError, "human format: uid=%s not found" % uid
				else:
					per_fn[fname] = onelog(logps, indexers, fname)
	
			except LG.NoDataError, x:
				die.warn("No data in %s: %s" % (fname, x))
			else:
				per_fn[fname] = onelog(logps, indexers, fname)
	return (per_fn, hdr)


def read_uid_many_files(fnames, uid, Nsamp, tail, trigger):
	"""
	@return: a dictionary mapping filenames onto "data" and header information.
		Header information is from the last file read.
		The "data" are L{onelog} class instances that encapsulate
		log information from one log file.
	@rtype: tuple(dict(str:onelog), dict(str:str))
	"""
	hdr = {}
	per_fn = {}
	for fname in fnames:
		try:
			hdr, loglines = LG.read_tail_uid(fname, uid,
							Nsamp=Nsamp, tail=tail,
							trigger=trigger
							)
			per_fn[fname] = loglines
		except LG.NoDataError, ex:
			die.warn("No data from %s: %s" % (fname, str(ex)))
	return (per_fn, hdr)


def get_pmap(per_fn):
	"""
	@type per_fn: dict(str: L{mcmc_newlogger.logline})
	"""
	for ol in per_fn.values():
		for ll in ol:
			return ll.idxr().map
	raise ValueError, "No data - cannot get_pmap."


def list_prm_samples(per_fn, sample_selector, out):
	"""
	@type per_fn: dict(str: L{mcmc_newlogger.logline})
	"""
	w = fiatio.writer(out)
	for ll in sample_selector(per_fn):
		tmp = {'T': ll.T() }
		for (k, v) in ll.idxr().get_all().items():
			tmp[IC.index_base._fmt(k)] = v
		w.datum(tmp)


def indexer_covar(per_fn, sample_selector, weight_by_T=False):
	"""
	@type per_fn: dict(str: L{mcmc_newlogger.logline})
	"""
	idxr_map = get_pmap(per_fn)
	ndim = len(idxr_map)
	sum = numpy.zeros((ndim,))
	sw = 0
	for ll in sample_selector(per_fn):
		idxr = ll.idxr()
		if weight_by_T:
			wt = 1.0/math.sqrt(ll.T())
			numpy.add(sum, idxr.get_prms()*wt, sum)
		else:
			wt = 1.0
			numpy.add(sum, idxr.get_prms(), sum)
		sw += wt
	assert sw > 0
	m = sum/sw

	n = 0
	mxwt = 0.0
	sum = numpy.zeros((ndim, ndim))
	for ll in sample_selector(per_fn):
		idxr = ll.idxr()
		delta = idxr.get_prms() - m
		assert len(delta.shape) == 1
		if weight_by_T:
			wt = 1.0/math.sqrt(ll.T())
			numpy.multiply(delta, wt, delta)
			if wt > mxwt:
				mxwt = wt
		else:
			mxwt = 1.0
		delta2 = numpy.outer(delta, delta)
		assert len(delta2.shape) == 2
		numpy.add(sum, delta2, sum)
		n += 1
	if sw > mxwt:
		var = (1.0-mxwt/sw)*sum/n
	else:
		var = None
	return (m, var, sw, idxr_map)


def logp_stdev(per_fn, sample_selector):
	"""
	@type per_fn: dict(str: L{mcmc_newlogger.logline})
	"""
	lp = []
	T = []
	for ll in sample_selector(per_fn):
		lp.append( ll.logp() )
		try:
			T.append( ll.T() )
		except ValueError:
			pass
	assert len(lp) > 0, "No data in average in logp_stdev"
	rv = [ (('logp',),) + gpkmisc.mean_stdev(lp) ]
	if len(T) > 1:
		rv.append( (('T',),) + gpkmisc.mean_stdev(T) )
	return rv


def all(per_fn, source=None):
	"""This selects all measurements."""
	for (fn, ol) in per_fn.items():
		if source is not None:
			source.add(fn, len(ol))
		for ll in ol:
			yield ll


def some_after_convergence(per_fn, source=None):
	"""This selects which measurements will be used.
	It looks after convergence, then throws out optimizations that haven't
	converged.
	"""
	maxes = []
	for ol in per_fn.values():
		maxes.append( max( q.logp() for q in ol ) )
	ndim = len(get_pmap(per_fn))
	assert ndim > 0
	tol = 6 * math.sqrt(ndim)
	threshold = max(maxes) - tol
	for (fn, ol) in per_fn.items():
		for ll in ol:
			if ll.logp() > threshold:
				if source is not None:
					source.add(fn, 1)
				yield ll


def near_each_max(per_fn, source=None):
	"""This selects which measurements will be used.
	It looks after convergence, then gives you the best few
	results from each run.
	@type per_fn: dict(str: L{mcmc_newlogger.logline})
	"""
	ndim = len(get_pmap(per_fn))
	assert ndim > 0
	tol = 6 * math.sqrt(ndim)
	for (fn,ol) in per_fn.items():
		themax = max( q.logp() for q in ol )
		threshold = themax - tol
		for ll in ol:
			if ll.logp() > threshold:
				if source is not None:
					source.add(fn, 1)
				yield ll



def last(per_fn, source=None):
	"""This selects which measurements will be used.
	@type per_fn: dict(str: L{mcmc_newlogger.logline})
	"""
	for (fn, ol) in per_fn.items():
		if source is not None:
			source.add(fn, 1)
		yield ol[-1]


def each_best(per_fn, source=None):
	"""This selects which measurements will be used."""
	HUGE = 1e30
	for (fn, ol) in per_fn.items():
		bestlogp = -HUGE
		best = None
		for ll in ol:
			if ll.logp() > bestlogp:
				best = ll
				bestlogp = ll.logp()
		if best is not None:
			if source is not None:
				source.add(fn, 1)
			yield best



def overall_best(per_fn, source=None):
	"""This selects which measurements will be used."""
	HUGE = 1e30
	bestlogp = -HUGE
	best = None
	for (fn, ol) in per_fn.items():
		for ll in ol:
			if ll.logp() > bestlogp:
				best = ll
				bestlogp = ll.logp()
	if best is not None:
		if source is not None:
			source.add(fn, 1)
		yield best




def indexer_stdev(per_fn, selector, weight_by_T=False):
	"""Return a summary of the properties of the selected indexers.
	@param selector: A generator that selects some of the C{logline}s in C{per_fn} that meet certain
		selection criteria.
	@type selector: function(dict(filename: list(L{mcmc_newlogger.logline})))
		that yields L{mcmc_newlogger.logline}
	"""
	mean, covar, n, idxr_map = indexer_covar(per_fn, selector, weight_by_T=weight_by_T)
	o = []
	for (nm, i) in idxr_map.items():
		if covar is None:
			std = None
		else:
			std = math.sqrt(covar[i,i])
		o.append( (nm, mean[i], std) )
	return o


def print_index_error(ke):
	print 'Cannot find key= %s in index.' % str(ke.args[0])
	idxr = ke.args[1]
	kl = sorted(idxr.map.keys())
	print 'The index has these keys:'
	for k in kl:
		print 'Key=', idxr._fmt(k)




class onelog:
	def __init__(self, logps, indexers, fname):
		assert len(indexers) == len(logps)
		assert isinstance(logps, numpy.ndarray)
		assert isinstance(indexers, list)
		self.logps = logps
		self.indexers = indexers
		self.fname = fname
		self.convergence = None





def drop_files(per_fn, FileDropFac, Trim=None, Stretch=None):
	"""
	@param per_fn: produced byLT.read_uid_many_files()
	"""
	mxvs = []
	for (k, lls) in per_fn.items():
		mxvs.append( (max([q.logp() for q in lls]), k) )
	mxvs.sort()
	for (lp, kk) in mxvs:
		print '# max(%s) = %g' % (kk, lp)
	ndrop = int(math.floor(FileDropFac * len(per_fn)))
	rv = {}
	for (mxlogp, k) in mxvs[ndrop:]:
		rv[k] = per_fn[k]
	return rv


def ascii_cmp(a, b):
	"""Compares the ASCII form of keys, for sorting purposes.   Does a good
	attempt at ASCII ordering for strings and numeric ordering for numbers.
	"""
	asp = a.split(',')
	bsp = b.split(',')
	return key_cmp(asp, bsp)


def key_cmp(a, b):
	"""Compares the tuple form of keys, for sorting purposes.   Does a good
	attempt at ASCII ordering for strings and numeric ordering for numbers.
	"""
	for (aa, bb) in zip(a, b):
		for fcn in [float, str]:
			try:
				c = cmp(fcn(aa), fcn(bb))
				if c != 0:
					return c
				break	# c==0
			except ValueError:
				pass
	return 0


def P_bayes_list(per_fn, argv, m, arg, selector, hdr):
	"""C{Log(PosteriorMarginalLikelihood)} is the posterior marginal likelihood of this model:
	C{P_posterior[Data|model] = integral( P[D|params,model] * P[params|data,model] * d_params)}.
	It's the average of P[D|params,model] over the posterior distribution
	of P[params|data,model]*Prior(params).

	References are:

		- Rampal S. Etienne, Han Olff (2005) Confronting different models of community structure to species-abundance data: a Bayesian model comparison Ecology Letters 8 (5) , 493-504 doi:10.1111/j.1461-0248.2005.00745.x 

	and that references

		- M. Aitkin 1991: Posterior Bayes Factors, J. of the Royal Statistical Soc. B 53: 111-142

		- P. W. Laud and J. G. Ibrahim 1995: Predictive Model Selection, J. of the Royal Statistical Soc. B 57: 247-262.

		- F. De Santis and F. Spezzaferri 1997: Alternative Bayes Factors for Model Selection, Canadian J. of Statistics 25: 503-515

		- S. K. Upadhyay and M. Peswani 2003: Choice between Weibull and Lognormal Models: a simulation based Bayesian Study Communications in Statistics: Theory and Methods 32: 318-405

		- P. K. Vlachos and A. E. Gelfand 2003: On the calibration of Bayeseian model choice criteria, J. of Statistical Planning and Inference 111: 223-234

		- R. E. Kass and A. E. Rafferty 1995: Bayes Factors, J. of the American Statistical Assoc. 90: 773-795

	The other thing, Log(BayesWeightedBayes)
	is the average of P(D|params,M)*Prior(params) over the posterior distribution
	of P[params|data,model]*Prior(params).
	It has no real statistical backing, but it's a crude approximation
	for the Bayes evidence itself,
	the normalized P[params|data,model]*Prior(params).
	"""
	pd = m.pd_factory(argv, hdr=hdr)
	lplist = []
	lblist = []
	idxrs = []
	for ll in selector(per_fn):
		idxr = ll.idxr()
		idxrs.append(idxr)
		try:
			lpdata = float(pd.logp_data_normalized(idxr))
			lplist.append(lpdata)
			if hasattr(pd, 'logp_prior_normalized'):
				lpprior = float(pd.logp_prior_normalized(idxr))
				lblist.append(lpdata + lpprior)
		except IC.IndexKeyError, ke:
			print_index_error(ke)
			raise
	return (lplist, lblist, idxrs)


def P_bayes(per_fn, argv, m, arg, selector, hdr):
	lplist, lblist, idxrs = P_bayes_list(per_fn, argv, m, arg, selector, hdr)
	rv = []
	mxlp = max(lplist)
	psum = 0.0
	for lp in lplist:
		psum += math.exp(max(-100.0, lp-mxlp))
	rv.append( (('Log(PosteriorMarginalLikelihood)',),
			mxlp + math.log(psum/len(lplist)))
		)
	mxlb = max(lplist)
	bsum = 0.0
	if lblist:
		for lb in lblist:
			bsum += math.exp(max(-100.0, lb-mxlb))
		rv.append( (('Log(BayesWeightedBayes)',),
				mxlb + math.log(bsum/len(lblist)))
			)
	else:
		die.info("You may want to have problem_definition.logp_prior_normalized() to compute the BayesWeightedBayes evidence.")
	return rv
