import numpy as np
import theano as th
import theano.tensor as T


def RiskCohort(Model, Features):
	"""
	Generates partial derivative weights of features and risk for the given
	model and profiles. Generates mean and standard deviation of these partial
	derivative weights.

	Parameters:
	----------
	Model : theano deep learning model
	a theano Neuralnetwork  model containing theano functions to feed the
	model with a profile of feature values.

	Profiles : array_like
	a matrix of profiles containing features value to be used as a input
	for feeding the model.

	Output :
	--------
	Gradients: numpy matrix
	A [N*D] matrix  contains feature weights.(N = number of profiles and D =
	number of features(dimention of input to the model )). each row contains
	feature weights of Correspondence profile in profiles matrix.

	Gradients_mean: numpy matrix
	 a [ 1 * D] matrix contains mean value of feautre weights.

	Gradients_std: numpy matrix
	a [ 1 * D] matrix contains standrad deviation of feautre weights.
	"""

	# initialize container for risk gradient profiles
	Gradients = np.zeros(Features.shape)

	# copy input to matrix for Theano
	Matrix = np.matrix(Features)

	# iterate through samples, calculating risk gradient profile for each
	for i in np.arange(Features.shape[0]):
		Gradients[i, :] = _RiskBackpropagate(Model, Matrix[i, :])

	return Gradients


def _RiskBackpropagate(Model, Features):
	"""
	Generates partial derivatives of input features in a neural network model.
	These represent the rate of change of risk with respect to each input
	feature.

	Parameters:
	----------
	Model : class
	A fine tuned model generated by training.

	Features : array_like
	A 1 x P numpy array of features corresponding to one sample.

	Output:
	----------
	Gradient : array_like
	A 1
	an array of the feature weights.
	"""

	# define partial derivative
	X = T.matrix('X')
	AtRisk = T.ivector('AtRisk')
	Observed = T.ivector('Observed')
	Is_train = T.scalar('Is_train', dtype='int32')
	masks = T.lmatrix('mask_' + str(0))
	partial_derivative = th.function(on_unused_input='ignore',
			inputs=[X, AtRisk, Observed, Is_train, masks],
			outputs=T.grad(Model.risk_layer.output[0],
				Model.x),
			givens={Model.x: X, Model.o: AtRisk,
				Model.at_risk: Observed,
				Model.is_train: Is_train,
				Model.masks[0]: masks},
			name='partial_derivative')

	# define parameters for risk and calculate partial
	sample_O = np.array([0]).astype(np.int32)
	sample_T = np.array([0]).astype(np.int32)
	# Create dummy masks for graph
	dummy_masks = np.ones((1, Model.n_hidden), dtype='int64')
	Gradient = partial_derivative(Features, sample_O, sample_T, 0, dummy_masks)

	return Gradient
