import inference as ifr
import output as op
import tests as tst
import distributions as dst
import numpy as np
from math import exp, sqrt, log, tanh
from scipy import optimize
import copy
import matplotlib.pyplot as plt
import matplotlib.mlab as mlab
import seaborn as sns
import numdifftools as nd
import covariances as cov

class TSM(object):

	def __init__(self,model_type):

		# Holding variables for model output
		self.params = []
		self.ses = []
		self.ihessian = []
		self.chains = []

		# Specify parameter description
		self.param_desc = []
		self.model_type = model_type

	# returns unnormalized posterior
	def posterior(self,beta):
		post = self.likelihood(beta)
		for k in range(self.param_no):
			post += -self.param_desc[k]['prior'].logpdf(beta[k])
		return post

	# Allows for quicker changes to prior specification
	def adjust_prior(self,index,prior):
		if isinstance(index, list):
			for item in index:
				if item < 0 or item > (self.param_no-1) or not isinstance(item, int):
					raise Exception("Oops - the parameter index " + str(item) + " you have entered is invalid!")
				else:
					self.param_desc[item]['prior'] = prior				
		else:
			if index < 0 or index > (self.param_no-1) or not isinstance(index, int):
				raise Exception("Oops - the parameter index " + str(index) + " you have entered is invalid!")
			else:
				self.param_desc[index]['prior'] = prior

	# Wrapper function for inference options
	def fit(self,method=None,printer=True,nsims=100000,cov_matrix=None,step=0.001,iterations=30000):
		if method is None:
			method = self.default_method

		if method not in self.supported_methods:
			raise ValueError("Method not supported!")

		if method == 'MLE':
			self.optimize_fit(self.likelihood,printer)
		elif method == 'MAP':
			self.optimize_fit(self.posterior,printer)	
		elif method == 'M-H':
			self.mcmc_fit(nsims=nsims,method=method,cov_matrix=cov_matrix)
		elif method == "Laplace":
			self.laplace_fit(self.posterior,printer) 
		elif method == "BBVI":
			self.bbvi_fit(self.posterior,printer,step=step,iterations=iterations)
		elif method == "OLS":
			self.ols_fit(printer)
		else:
			raise ValueError("Method not recognized!")

	# Outputs tabular summary of prior assumptions
	def list_priors(self):

		prior_list = []
		prior_desc = []
		param_trans = []

		for i in range(self.param_no):
			param_trans.append(self.param_desc[i]['prior'].transform_name)
			x = self.param_desc[i]['prior']
			if isinstance(x, ifr.Normal):
				prior_list.append('Normal')
				prior_desc.append('mu0: ' + str(self.param_desc[i]['prior'].mu0) + ', sigma0: ' + str(self.param_desc[i]['prior'].sigma0))
			elif isinstance(x, ifr.InverseGamma):
				prior_list.append('Inverse Gamma')
				prior_desc.append('alpha: ' + str(self.param_desc[i]['prior'].alpha) + ', beta: ' + str(self.param_desc[i]['prior'].beta))
			elif isinstance(x, ifr.Uniform):
				prior_list.append('Uniform')
				prior_desc.append('n/a (non-informative)')
			else:
				raise Exception("Error - prior distribution not detected!")

		data = []

		for i in range(self.param_no):
			data.append({'param_index': len(data), 'param_name': self.param_desc[i]['name'], 'param_prior': prior_list[i], 'param_prior_params': prior_desc[i], 'param_trans': param_trans[i]})

		fmt = [
			('Index','param_index',6),		
			('Parameter','param_name',25),
			('Prior Type','param_prior',10),
			('Prior Parameters', 'param_prior_params', 25),
			('Transformation','param_trans',20)]
		
		print( op.TablePrinter(fmt, ul='=')(data) )

	# Approximate distributions used for mean-field variational inference
	def list_q(self):

		q_list = []

		for i in range(self.param_no):
			x = self.param_desc[i]['q']
			if isinstance(x, dst.Normal):
				q_list.append('Normal')
			elif isinstance(x, dst.InverseGamma):
				q_list.append('Inverse Gamma')
			elif isinstance(x, dst.Uniform):
				q_list.append('Uniform')
			else:
				raise Exception("Error - prior distribution not detected!")

		data = []

		for i in range(self.param_no):
			data.append({'param_index': len(data), 'param_name': self.param_desc[i]['name'], 'param_q': q_list[i]})

		fmt = [
			('Index','param_index',6),		
			('Parameter','param_name',25),
			('q(z)','param_q',10)]
		print "Approximate distributions for mean-field variational inference"
		print ""
		print( op.TablePrinter(fmt, ul='=')(data) )

	def normal_posterior_sim(self,parameters,cov_matrix,method_name,printer):
		chain, mean_est, median_est, upper_95_est, lower_95_est = ifr.norm_post_sim(parameters,cov_matrix)

		# Transform parameters - replace if statement with transform information directly in future version
		for k in range(len(chain)):
			if self.param_desc[k]['prior'].transform == np.exp:
				chain[k] = np.exp(chain[k])
				mean_est[k] = exp(mean_est[k])
				median_est[k] = exp(median_est[k])
				upper_95_est[k] = exp(upper_95_est[k])
				lower_95_est[k] = exp(lower_95_est[k])				
			elif self.param_desc[k]['prior'].transform == np.tanh: 
				chain[k] = np.tanh(chain[k])
				mean_est[k] = tanh(mean_est[k])
				median_est[k] = tanh(median_est[k])
				upper_95_est[k] = tanh(upper_95_est[k])
				lower_95_est[k] = tanh(lower_95_est[k])	

		self.chains = chain

		if printer is True:

			data = []

			for i in range(len(self.param_desc)):
				data.append({'param_name': self.param_desc[i]['name'], 'param_mean': round(mean_est[i],4), 'param_median': round(median_est[i],4), 'ci':  "(" + str(round(lower_95_est[i],4)) + " | " + str(round(upper_95_est[i],4)) + ")"})

			fmt = [
				('Parameter','param_name',20),
				('Median','param_median',10),
				('Mean', 'param_mean', 15),
				('95% Credibility Interval','ci',25)]

			if self.model_type == 'ARIMA':
				self.model_name = "ARIMA(" + str(self.ar) + "," + str(self.integ) + "," + str(self.ma) + ") regression"
				self.cutoff = max(self.ar,self.ma)
			elif self.model_type == 'GAS':
				self.model_name = "GAS(" + str(self.ar) + "," + str(self.integ) + "," + str(self.sc) + ") regression"
				self.cutoff = max(self.ar,self.sc)
			elif self.model_type == 'VAR':
				self.model_name = "VAR(" + str(self.lags) + ") regression"
				self.cutoff = self.lags			

			print self.model_name 
			print "=================="
			print "Method: " + method_name + " Approximation"
			print "Number of observations: " + str(len(self.data)-self.cutoff)				
			print "Unnormalized Log Posterior: " + str(round(-self.posterior(mean_est),4))
			print ""
			print( op.TablePrinter(fmt, ul='=')(data) )
		
		# Plot densities
		for j in range(len(self.params)):
			fig = plt.figure()
			a = sns.distplot(chain[j], rug=False, hist=False, label=self.param_desc[j]['name'])
			a.set_title(self.param_desc[j]['name'])
			a.set_ylabel('Density')
				
		sns.plt.show()		

	# Performs a Laplace Approximation using MAP point estimates and Inverse Hessian
	def laplace_fit(self,obj_type,printer=True):

		# Get Mode and Inverse Hessian information
		self.fit(method='MAP',printer=False)

		if len(self.ses) == 0: # no errors, no Laplace party
			raise Exception("No Hessian information - Laplace approximation cannot be performed")
		else:
			self.normal_posterior_sim(self.params,self.ihessian,"Laplace",printer)

	# Stores point estimates for MAP and MLE estimation; prints summary output
	def optimize_fit(self,obj_type,printer=True):

		# Starting parameters
		phi = np.zeros(self.param_no)

		# Starting values
		if self.model_type == 'GAS':
			if self.dist in ['Laplace','Normal']:
				phi[0] = np.mean(self.data)
			elif self.dist == 'Poisson':
				phi[0] = np.log(np.mean(self.data))
		elif self.model_type == 'VAR':
			phi = self.create_B().flatten()
			cov = self.estimate_cov()

			# Inelegant - needs refactoring
			for i in range(self.ylen):
				for k in range(self.ylen):
					if i == k:
						phi = np.append(phi,self.param_desc[len(phi)]['prior'].itransform(cov[i,k]))
					elif i > k:
						phi = np.append(phi,self.param_desc[len(phi)]['prior'].itransform(cov[i,k]))

		# Optimize using L-BFGS-B
		p = optimize.minimize(obj_type,phi,method='L-BFGS-B')
		self.params = p.x

		# Vector for transformed parameters
		t_params = copy.deepcopy(p.x)			

		# Check that matrix is non-singular; act accordingly
		try:
			self.ihessian = np.linalg.inv(nd.Hessian(obj_type)(self.params))
			self.ses = np.abs(np.diag(self.ihessian))**0.5

			t_p_std = copy.deepcopy(self.ses) # vector for transformed standard errors
			
			# Create transformed variables
			for k in range(len(t_params)-self.param_hide):
				z_temp = (p.x[k]/float(self.ses[k]))
				t_params[k] = self.param_desc[k]['prior'].transform(t_params[k])
				t_p_std[k] = t_params[k] / z_temp

			# Replace with something more elegant in future versions
			if printer is True:

				data = []

				for i in range(len(self.param_desc)-self.param_hide):
					data.append({'param_name': self.param_desc[i]['name'], 'param_value':round(self.param_desc[i]['prior'].transform(p.x[i]),4), 'param_std': round(t_p_std[i],4),'param_z': round(t_params[i]/float(t_p_std[i]),4),'param_p': round(tst.find_p_value(t_params[i]/float(t_p_std[i])),4),'ci': "(" + str(round(t_params[i] - t_p_std[i]*1.96,4)) + " | " + str(round(t_params[i] + t_p_std[i]*1.96,4)) + ")"})

				fmt = [
				    ('Parameter',       'param_name',   40),
				    ('Estimate',          'param_value',       10),
				    ('Standard Error', 'param_std', 15),
				    ('z',          'param_z',       10),
				    ('P>|z|',          'param_p',       10),
				    ('95% Confidence Interval',          'ci',       25)
				]

		except: # If Hessian is not invertible...
			print "Hessian not invertible! Consider a different model specification."
			print ""

			# Reset iHessian and standard errors, reflecting non-invertibility
			self.ihessian = []
			self.ses = []

			# Transform parameters
			for k in range(len(t_params)):
				t_params[k] = self.param_desc[k]['prior'].transform(t_params[k])

			if printer is True:

				data = []

				for i in range(len(self.param_desc)):
					data.append({'param_name': self.param_desc[i]['name'], 'param_value':round(self.param_desc[i]['prior'].transform(p.x[i]),4)})

				fmt = [
				    ('Parameter',       'param_name',   40),
				    ('Estimate',          'param_value',       10)
				]

		# Final printed output
		if printer is True:

			if self.model_type == 'ARIMA':
				self.model_name = "ARIMA(" + str(self.ar) + "," + str(self.integ) + "," + str(self.ma) + ") regression"
				self.cutoff = max(self.ar,self.ma)
				self.data_length = self.data
			elif self.model_type == 'GAS':
				self.model_name = "GAS(" + str(self.ar) + "," + str(self.integ) + "," + str(self.sc) + ") regression"
				self.cutoff = max(self.ar,self.sc)
				self.data_length = self.data				
			elif self.model_type == 'VAR':
				self.model_name = "VAR(" + str(self.lags) + ") regression"
				self.cutoff = self.lags			
				self.data_length = self.data[0]

			print self.model_name 
			print "=================="
			if obj_type == self.likelihood:
				print "Method: MLE"
			elif obj_type == self.posterior:
				print "Method: MAP"
			print "Number of observations: " + str(len(self.data_length)-self.cutoff)
			print "Log Likelihood: " + str(round(-self.likelihood(p.x),4))
			if obj_type == self.posterior:
				print "Unnormalized Log Posterior: " + str(round(-self.posterior(p.x),4))
			print "AIC: " + str(round(2*len(p.x)+2*self.likelihood(p.x),4))
			print "BIC: " + str(round(2*self.likelihood(p.x) + len(p.x)*log(len(self.data_length)-self.cutoff),4))
			print ""
			print( op.TablePrinter(fmt, ul='=')(data) )

	def ols_fit(self,printer):
		self.params = self.create_B().flatten()
		cov = self.estimate_cov()

		# Inelegant - needs refactoring
		for i in range(self.ylen):
			for k in range(self.ylen):
				if i == k:
					self.params = np.append(self.params,self.param_desc[len(self.params)]['prior'].itransform(cov[i,k]))
				elif i > k:
					self.params = np.append(self.params,self.param_desc[len(self.params)]['prior'].itransform(cov[i,k]))

		self.ihessian = self.estimator_cov('OLS')
		self.ses = np.abs(np.diag(self.ihessian))**0.5
		t_params = copy.deepcopy(self.params)			
		t_p_std = copy.deepcopy(self.ses) # vector for transformed standard errors
			
		# Create transformed variables
		for k in range(len(t_params)-self.param_hide):
			z_temp = (self.params[k]/float(self.ses[k]))
			t_params[k] = self.param_desc[k]['prior'].transform(t_params[k])
			t_p_std[k] = t_params[k] / z_temp

		# Replace with something more elegant in future versions
		if printer is True:

			data = []

			for i in range(len(self.param_desc)-self.param_hide):
				data.append({'param_name': self.param_desc[i]['name'], 'param_value':round(self.param_desc[i]['prior'].transform(self.params[i]),4), 'param_std': round(t_p_std[i],4),'param_z': round(t_params[i]/float(t_p_std[i]),4),'param_p': round(tst.find_p_value(t_params[i]/float(t_p_std[i])),4),'ci': "(" + str(round(t_params[i] - t_p_std[i]*1.96,4)) + " | " + str(round(t_params[i] + t_p_std[i]*1.96,4)) + ")"})

			fmt = [
			    ('Parameter',       'param_name',   40),
			    ('Estimate',          'param_value',       10),
			    ('Standard Error', 'param_std', 15),
			    ('z',          'param_z',       10),
			    ('P>|z|',          'param_p',       10),
			    ('95% Confidence Interval',          'ci',       25)
				]		

			if self.model_type == 'ARIMA':
				self.model_name = "ARIMA(" + str(self.ar) + "," + str(self.integ) + "," + str(self.ma) + ") regression"
				self.cutoff = max(self.ar,self.ma)
				self.data_length = self.data
			elif self.model_type == 'GAS':
				self.model_name = "GAS(" + str(self.ar) + "," + str(self.integ) + "," + str(self.sc) + ") regression"
				self.cutoff = max(self.ar,self.sc)
				self.data_length = self.data				
			elif self.model_type == 'VAR':
				self.model_name = "VAR(" + str(self.lags) + ") regression"
				self.cutoff = self.lags			
				self.data_length = self.data[0]

			print self.model_name 
			print "=================="
			print "Method: OLS"
			print "Number of observations: " + str(len(self.data_length)-self.cutoff)
			print "Log Likelihood: " + str(round(-self.likelihood(self.params),4))
			print "AIC: " + str(round(2*len(self.params)+2*self.likelihood(self.params),4))
			print "BIC: " + str(round(2*self.likelihood(self.params) + len(self.params)*log(len(self.data_length)-self.cutoff),4))
			print ""
			print( op.TablePrinter(fmt, ul='=')(data) )

	def mcmc_fit(self,scale=(2.38/sqrt(10000)),nsims=100000,printer=True,method="M-H",cov_matrix=None):

		self.fit(method='MAP',printer=False)
		
		if method == "M-H":
			chain, mean_est, median_est, upper_95_est, lower_95_est = ifr.metropolis_hastings(self.posterior,scale,nsims,self.params,cov_matrix=cov_matrix)
		else:
			raise Exception("Method not recognized!")

		self.params = np.asarray(mean_est)

		for k in range(len(chain)):
			if self.param_desc[k]['prior'].transform == np.exp:
				chain[k] = np.exp(chain[k])
				mean_est[k] = exp(mean_est[k])
				median_est[k] = exp(median_est[k])
				upper_95_est[k] = exp(upper_95_est[k])
				lower_95_est[k] = exp(lower_95_est[k])				
			elif self.param_desc[k]['prior'].transform == np.tanh: 
				chain[k] = np.tanh(chain[k])
				mean_est[k] = tanh(mean_est[k])
				median_est[k] = tanh(median_est[k])				
				upper_95_est[k] = tanh(upper_95_est[k])
				lower_95_est[k] = tanh(lower_95_est[k])	

		self.chains = chain

		if printer == True:

			data = []

			for i in range(len(self.param_desc)):
				data.append({'param_name': self.param_desc[i]['name'], 'parm_mean':round(mean_est[i],4), 'parm_median':round(median_est[i],4), 'ci': "(" + str(round(lower_95_est[i],4)) + " | " + str(round(upper_95_est[i],4)) + ")"})

		fmt = [
			('Parameter','parm_name',20),
			('Median','parm_median',10),
			('Mean', 'parm_mean', 15),
			('95% Credibility Interval','ci',25)]

		if self.model_type == 'ARIMA':
			self.model_name = "ARIMA(" + str(self.ar) + "," + str(self.integ) + "," + str(self.ma) + ") regression"
			self.cutoff = max(self.ar,self.ma)
		elif self.model_type == 'GAS':
			self.model_name = "GAS(" + str(self.ar) + "," + str(self.integ) + "," + str(self.sc) + ") regression"
			self.cutoff = max(self.ar,self.sc)
		elif self.model_type == 'VAR':
			self.model_name = "VAR(" + str(self.lags) + ") regression"
			self.cutoff = self.lags						

		print self.model_name 
		print "=================="
		print "Method: Metropolis-Hastings"
		print "Number of simulations: " + str(nsims)
		print "Number of observations: " + str(len(self.data)-self.cutoff)
		print "Unnormalized Log Posterior: " + str(round(-self.posterior(mean_est),4))
		print ""
		print( op.TablePrinter(fmt, ul='=')(data) )

		fig = plt.figure()

		for j in range(len(self.params)):

			for k in range(4):
				iteration = j*4 + k + 1
				ax = fig.add_subplot(len(self.params),4,iteration)

				if iteration in range(1,len(self.params)*4 + 1,4):
					a = sns.distplot(chain[j], rug=False, hist=False)
					a.set_ylabel(self.param_desc[j]['name'])
					if iteration == 1:
						a.set_title('Density Estimate')
				elif iteration in range(2,len(self.params)*4 + 1,4):
					a = plt.plot(chain[j])
					if iteration == 2:
						plt.title('Trace Plot')
				elif iteration in range(3,len(self.params)*4 + 1,4): 
					plt.plot(np.cumsum(chain[j])/np.array(range(1,len(chain[j])+1)))
					if iteration == 3:
						plt.title('Cumulative Average')					
				elif iteration in range(4,len(self.params)*4 + 1,4):
					plt.bar(range(1,10),[cov.acf(chain[j],lag) for lag in range(1,10)])
					if iteration == 4:
						plt.title('ACF Plot')						
		sns.plt.show()		


	def plot_fit(self):
		if len(self.params) == 0:
			raise Exception("No parameters estimated!")
		else:
			if self.model_type == 'GAS':
				plt.figure()
				date_index = self.index[max(self.ar,self.sc):len(self.data)]
				mu, Y, scores = self.model(self.params,self.data)
				plt.plot(date_index,Y,label='Data')
				plt.plot(date_index,self.link(mu),label='Filter',c='black')
				plt.title(self.data_name)
			elif self.model_type == 'VAR':
				date_index = self.index[self.lags:len(self.data[0])]
				mu, Y = self.model(self.params,self.data)
				for series in range(len(Y)):
					plt.figure()
					plt.plot(date_index,Y[series],label='Data ' + str(series))
					plt.plot(date_index,mu[series],label='Filter' + str(series),c='black')	
					plt.title(self.data_name[series])
			else:
				plt.figure()
				date_index = self.index[max(self.ar,self.ma):len(self.data)]
				mu, Y = self.model(self.params,self.data)
				plt.plot(date_index,Y,label='Data')
				plt.plot(date_index,mu,label='Filter',c='black')
				plt.title(self.data_name)
			plt.show()				

	# Performs Black Box Variational Inference
	def bbvi_fit(self,posterior,printer=True,step=0.001,iterations=30000):

		self.fit("MAP",printer=False)

		# Starting values for approximate distribution
		for i in range(len(self.param_desc)):
			approx_dist = self.param_desc[i]['q']
			if isinstance(approx_dist, dst.Normal):
				self.param_desc[i]['q'].loc = self.params[i]
				if len(self.ses) == 0:
					self.param_desc[i]['q'].scale = -3.0
				else:
					self.param_desc[i]['q'].scale = log(self.ses[i])
		q_list = [k['q'] for k in self.param_desc]
		
		bbvi_obj = ifr.BBVI(posterior,q_list,12,step,iterations)
		q, q_params, q_ses = bbvi_obj.lambda_update()
		self.params = q_params

		for k in range(len(self.param_desc)):
			self.param_desc[k]['q'] = q[k]

		self.normal_posterior_sim(self.params,np.diag(np.exp(q_ses)),"Black Box Variational Inference",printer)