#!python

from __future__ import division

"""Description Here.
"""

from gmisclib import mcmc

import sys
import time
import numpy as NP
cimport numpy as NP
DTYPE = NP.float
ctypedef NP.float_t DTYPE_t

cdef extern from "math.h":
	double pow(double x, double y)
	double log(double x)
	double sqrt(double x)
	double fabs(double x)
	int isnan(double)

cimport libc.stdlib

cdef extern from "gpklib.h":
	double **double2alloc(int nr, int nc)
	double *double1alloc(int nr)
	void double1free(double *q)
	void double2free(double **q)
	int double1test(double *q)
	int double2test(double **q)
	double quiet_nan(long)

cdef extern from "malloc.h":
	void *calloc(unsigned, unsigned)
	void free(void*)


cdef extern from "lawcjcrk.h":
	struct initialization_s:
		int nunconstrained
		int gf_samps

	ctypedef char* const_char_p "const char*"
	ctypedef double* const_double_p "const double*"
	ctypedef initialization_s* CSIP
	ctypedef initialization_s* CCSIP

	void lawcjc_init (const_char_p ifilename, const_char_p ofilename, initialization_s *out )
	void lawcjc_start (int nguesses, double **prmvecs, CSIP isp )
	double wsdo_logp(const_double_p prmvec, CCSIP isp, int do_not_add_fudge)
	int wprojct(double *p, CCSIP isp )
	void wsdo_mcmc_output ( int mcmc_num, double **mcmc_vecs, CSIP isp)
	void wsdo_max_output ( const_double_p prmvec, CCSIP isp)




cdef class problem_def:
	cdef public object __g_implements__
	cdef initialization_s c
	cdef object opt_filename, mcmc_filename, outlist
	cdef double maxlogp

	def __init__(self, ifilename, opt_filename):
		"""This initializes the description of the overall problem
		to be solved.
		"""
		lawcjc_init(ifilename, opt_filename, &self.c)
		self.maxlogp = -1e38
		self.__g_implements__ = [ 'problem_definition' ]
		self.outlist = []

	cdef int nprm(self) except -1:
		assert self.c.nunconstrained > 0
		return self.c.nunconstrained

	def guesses(self):
		cdef int nguesses = self.nprm() + 5
		cdef double**x = double2alloc(nguesses, self.nprm())
		cdef object out = []
		lawcjc_start(nguesses, x, &self.c)
		cdef int i, j
		cdef NP.ndarray[DTYPE_t, ndim=1] a
		for j in range(nguesses):
			a = NP.zeros((self.nprm(),), dtype=DTYPE)
			for i in range(self.nprm()):
				a[i] = x[j][i]
			out.append(a)
		double2free(x)
		return out

	def logp(self, NP.ndarray[DTYPE_t, ndim=1] x):
		"""This function is the key interface to the MCMC module.
		This should return the log of the probability of seeing the
		actual data, for a given set of parameters x.

		In actual fact, it often just pulls the value from the cache,
		because a prior call to fixer() already calculated it.
		"""
		assert x.shape[0] == self.nprm()
		cdef double *xtmp = double1alloc(x.shape[0])
		cdef int i
		for i in range(x.shape[0]):
			xtmp[i] = x[i]
		assert(double1test(xtmp)==0)
		cdef double tmp = wsdo_logp(xtmp, &self.c, 0)
		# TEST:
		# for i in range(x.shape[0]):
			# assert fabs(xtmp[i]-x[i]) < 0.001
		# END TEST:
		# cdef double tmp2 = wsdo_logp(xtmp, &self.c, 0)
		# assert fabs(tmp - tmp2) < 0.01, "wsdo_logp() output changes from %s to %s" % (tmp, tmp2)

		double1free(xtmp)
		print 'logP=%.2f at t=%.3f' % (tmp, time.clock())
		sys.stdout.flush()
		if isnan(tmp):
			return None
		return tmp

	def fixer(self, NP.ndarray[DTYPE_t, ndim=1] x):
		"""This method is called in the MCMC stepper to constrain
		parameters.   For instance, if a parameter is required to be
		positive, this provides an opportunity for self.logp_guts()
		to change its value, if needed.
		"""
		assert x.shape[0] == self.nprm()
		print "Allocating for wprojct, size=", x.shape[0]
		cdef double *xtmp = double1alloc(x.shape[0])
		cdef int i
		for i in range(x.shape[0]):
			xtmp[i] = x[i]
		assert(double1test(xtmp)==0)
		if not wprojct(xtmp, &self.c):
			double1free(xtmp)
			raise mcmc.NotGoodPosition, 'wproject returns false.'
		assert(double1test(xtmp)==0)

		#TEST:
		# cdef double *xtmp2 = double1alloc(x.shape[0])
		# for i in range(x.shape[0]):
			# xtmp2[i] = xtmp[i]
		# cdef bint ok2 = wprojct(xtmp, &self.c)
		# assert(ok2)
		# for i in range(x.shape[0]):
			# assert fabs(xtmp2[i] - xtmp[i]) < 0.001
		#END TEST

		for i in range(self.nprm()):
			x[i] = xtmp[i]
		double1free(xtmp)
		return x

	def output_fcn(self, NP.ndarray[DTYPE_t, ndim=1] x):
		self.outlist.append(x)

	def reset_outlist(self):
		self.outlist = []

	def max_fcn(self, NP.ndarray[DTYPE_t, ndim=1] x):
		cdef double *xtmp
		cdef int i
		assert x.shape[0] == self.nprm()
		xtmp = double1alloc(x.shape[0])
		for i in range(x.shape[0]):
			xtmp[i] = x[i]
		assert(double1test(xtmp)==0)
		wsdo_max_output ( xtmp, &self.c)
		double1free(xtmp)


	def close(self):
		print "Entering bsr_analysis_guts.close()"
		cdef int i, j
		cdef int n = len(self.outlist)
		print "Nprm=", self.nprm(), "n_samples=", n
		cdef double **tmp = double2alloc(self.nprm(), n)
		cdef NP.ndarray[DTYPE_t, ndim=1] x
		for (j, x) in enumerate(self.outlist):
			assert x.shape[0] == self.nprm()
			for i in range(self.nprm()):
				tmp[i][j] = x[i]
		self.outlist = None
		print "Just about to call wsdo_mcmc_output with", n, "samples"
		wsdo_mcmc_output (n, tmp, &self.c)
		double2free(tmp)
		print '\nCLOSE done'

	def requested_samples(self):
		return self.c.gf_samps
