#!/usr/bin/env 

__author__ = 'Florian Hase'

#========================================================================

import numpy as np 

#========================================================================

def sigmoid(x):
	return 1. / (1. + np.exp( - x))

#========================================================================

class NumpyGraph(object):

	def __init__(self, config, model_details):
		self.config = config
#		Logger.__init__(self, 'NumpyGraph', self.config.get('verbosity'))

		self.model_details = model_details
		for key, value in self.model_details.items():
			setattr(self, '_%s' % str(key), value)

		self.feature_size    = len(self.config.kernel_names)
		self.bnn_output_size = len(self.config.kernel_names)
		self.target_size     = len(self.config.kernel_names)

	
	def declare_training_data(self, features):
		
		self.num_obs  = len(features)
		self.features = features


	def compute_kernels(self, posteriors):
	
		tau_rescaling = np.zeros((self.num_obs, self.bnn_output_size))
		kernel_ranges = self.config.kernel_ranges
		for obs_index in range(self.num_obs):
			tau_rescaling[obs_index] += kernel_ranges
		tau_rescaling = tau_rescaling**2
	
		# sample from BNN
		activations = [np.tanh, np.tanh, lambda x: x]
		post_layer_outputs = [np.array([self.features for _ in range(self._num_draws)])]
		for layer_index in range(self._num_layers):

			weight     = posteriors['weight_%d' % layer_index]
			bias       = posteriors['bias_%d'   % layer_index]
			activation = activations[layer_index]

			outputs = []
			for sample_index in range(len(weight)):

				single_weight = weight[sample_index]
				single_bias   = bias[sample_index]

				output = activation( np.matmul( post_layer_outputs[-1][sample_index], single_weight) + single_bias)
				outputs.append(output)

#			print(np.array(outputs).shape)
#
#			quit()

			post_layer_output = np.array(outputs)
			post_layer_outputs.append(post_layer_output)

#			print('DONE')

		post_bnn_output = post_layer_outputs[-1]

		# note: np.random.gamma is parametrized with k and theta, while ed.models.Gamma is parametrized with alpha and beta
		post_tau_normed = np.random.gamma( self.num_obs**2 + np.zeros(post_bnn_output.shape), np.ones(post_bnn_output.shape))
		post_tau        = post_tau_normed / tau_rescaling
		post_sqrt_tau   = np.sqrt(post_tau)
		post_scale      = 1. / post_sqrt_tau

#		print('BNN OUTPUT', post_bnn_output.shape)
#		print('BNN TAU', post_sqrt_tau.shape)
#		quit()
		

		# map BNN output to predictions

		post_kernels = {}
		
		target_element_index = 0
		kernel_element_index = 0

		while kernel_element_index < len(self.config.kernel_names):

			kernel_type = self.config.kernel_types[kernel_element_index]
			kernel_size = self.config.kernel_sizes[kernel_element_index]

			feature_begin, feature_end = target_element_index, target_element_index + 1
			kernel_begin,  kernel_end  = kernel_element_index, kernel_element_index + kernel_size

#			print('KERNEL', kernel_begin, kernel_end)

			post_relevant = post_bnn_output[:, :, kernel_begin : kernel_end]

			if kernel_type == 'continuous':

				lowers       = self.config.kernel_lowers[kernel_begin : kernel_end]
				uppers       = self.config.kernel_uppers[kernel_begin : kernel_end]
				post_support = (uppers - lowers) * (1.2 * sigmoid(post_relevant) - 0.1) + lowers
				
				post_kernels['param_%d' % target_element_index] = {'loc':       post_support,
																   'sqrt_prec': post_sqrt_tau[:, :, kernel_begin : kernel_end],
																   'scale':     post_scale[:, :, kernel_begin : kernel_end]}


			elif kernel_type == 'categorical':

				post_temperature = 0.5 + 10.0 / self.num_obs
#				post_temperature = 0.4
				post_support     = post_relevant

				post_probs        = 1. / (1. + np.exp( - post_support))
				post_probs_normed = post_probs / np.sum(post_probs)

				# sample from relaxed one hot categorical
				gumbel_samples       = np.random.gumbel(loc = 0., scale = 1., size = post_support.shape)
				exp_samples          = np.exp( ( post_support + gumbel_samples) / post_temperature )
				exp_samples_sums     = np.sum(exp_samples, axis = 2)

				post_predict_relaxed = np.zeros(exp_samples.shape)
				for sample_index in range(exp_samples.shape[0]):
					for param_index in range(exp_samples.shape[1]):
						post_predict_relaxed[sample_index, param_index] = exp_samples[sample_index, param_index] / exp_samples_sums[sample_index, param_index]

				post_kernels['param_%d' % target_element_index] = {'probs': post_predict_relaxed}

			elif kernel_type == 'discrete':
				
				post_temperature = 0.5 + 1.0 / self.num_obs
	#			post_temperature = 0.4
				post_support     = post_relevant

				post_probs        = 1. / (1. + np.exp( - post_support))
				post_probs_normed = post_probs / np.sum(post_probs)

				# sample from relaxed one hot categorical
				gumbel_samples = np.random.gumbel(loc = 0., scale = 1., size = post_support.shape)
				
				exp_samples          = np.exp( ( post_support + gumbel_samples) / post_temperature )
				exp_samples_sums     = np.sum(exp_samples, axis = 2)

				post_predict_relaxed = np.zeros(exp_samples.shape)
				for sample_index in range(exp_samples.shape[0]):
					for param_index in range(exp_samples.shape[1]):
						post_predict_relaxed[sample_index, param_index] = exp_samples[sample_index, param_index] / exp_samples_sums[sample_index, param_index]

				post_kernels['param_%d' % target_element_index] = {'probs': post_predict_relaxed}



			else:
				raise NotImplementedError

			target_element_index += 1
			kernel_element_index += kernel_size


		return post_kernels



