import theano.tensor as T
import cPickle as pkl
from theano.compat.python2x import OrderedDict
from pylearn2.costs.cost import DefaultDataSpecsMixin, Cost
from pylearn2.models.mlp import Softmax, Sigmoid
from pylearn2.expr.nnet import kl
from pylearn2.utils import sharedX



class KDCost(DefaultDataSpecsMixin, Cost):
    """
    Represents an objective function to be minimized by some
    `TrainingAlgorithm`.
    """
    
    # If True, the data argument to expr and get_gradients must be a
    # (X, Y) pair, and Y cannot be None.
    supervised = True
    
    def __init__(self, teacher_path, temperature=1, lambda_target=1, lambda_teach=1, hints=None):
      self.temperature = temperature
      
      # Load teacher network and change parameters according to temperature.
      fo = open(teacher_path, 'r')
      teacher = pkl.load(fo)
      fo.close()
      
      tparams = teacher.layers[-1].get_param_values()
      tparams = [x/float(self.temperature) for x in tparams]
      teacher.layers[-1].set_param_values(tparams)

      self.teacher = teacher
      self.lambda_target = lambda_target
      self.lambda_teach = sharedX(lambda_teach, 'lambda_teach')
      self.hints = hints
      
    def cost_wrt_target(self, model, data):
        space, sources = self.get_data_specs(model)
        space.validate(data)
        (x, targets) = data
                                
        axes = model.input_space.axes
                        
        # Compute student output
        Ps_y_given_x = model.fprop(x)

        # Compute cross-entropy cost
        if isinstance(model.layers[-1], Sigmoid):
	  total = self.kl(Y=targets, Y_hat=Ps_y_given_x, batch_axis=axes.index('b'))
	  rval = total.mean()
	else:
	  targets = T.argmax(targets, axis=1)
	  rval = -T.log(Ps_y_given_x)[T.arange(targets.shape[0]), targets]
  
        rval = rval.astype('float32')
        
        return rval
        
    def kl(self, Y, Y_hat, batch_axis):
        """
        Computes the KL divergence.


        Parameters
        ----------
        Y : Variable
            targets for the sigmoid outputs. Currently Y must be purely binary.
            If it's not, you'll still get the right gradient, but the
            value in the monitoring channel will be wrong.
        Y_hat : Variable
            predictions made by the sigmoid layer. Y_hat must be generated by
            fprop, i.e., it must be a symbolic sigmoid.

        Returns
        -------
        ave : Variable
            average kl divergence between Y and Y_hat.

        Notes
        -----
        Warning: This function expects a sigmoid nonlinearity in the
        output layer and it uses kl function under pylearn2/expr/nnet/.
        Returns a batch (vector) of mean across units of KL
        divergence for each example,
        KL(P || Q) where P is defined by Y and Q is defined by Y_hat:

        p log p - p log q + (1-p) log (1-p) - (1-p) log (1-q)
        For binary p, some terms drop out:
        - p log q - (1-p) log (1-q)
        - p log sigmoid(z) - (1-p) log sigmoid(-z)
        p softplus(-z) + (1-p) softplus(z)
        """
        div = kl(Y=Y, Y_hat=Y_hat, batch_axis=batch_axis)
        return div

        
    def cost_wrt_teacher(self, model, data):
        space, sources = self.get_data_specs(model)
        space.validate(data)
        (x, targets) = data
        
        axes = model.input_space.axes
                                                
        # Compute teacher softened (relaxed) output
	Pt_y_given_x_relaxed = self.teacher.fprop(x)
	
	# Compute student softened (relaxed) output
	sparams = model.layers[-1].get_param_values()
	sparams_relaxed = [item/float(self.temperature) for item in sparams]
	model.layers[-1].set_param_values(sparams_relaxed) 
        Ps_y_given_x_relaxed = model.fprop(x)	
        model.layers[-1].set_param_values(sparams)
				
                         
	# Compute cost
	if isinstance(model.layers[-1], Sigmoid):
	  term1 = -T.log(Ps_y_given_x_relaxed) * Pt_y_given_x_relaxed
	  term2 = -T.log((1 - Ps_y_given_x_relaxed)) * (1 - Pt_y_given_x_relaxed)
	  rval = term1 + term2
	else:
	  rval = -T.log(Ps_y_given_x_relaxed) * Pt_y_given_x_relaxed 
	  rval = T.sum(rval, axis=1)
        
        return rval  
        
    def expr(self, model, data, ** kwargs):
        """
        Returns a theano expression for the cost function.
        
        Parameters
        ----------
        model : a pylearn2 Model instance
        data : a batch in cost.get_data_specs() form
        kwargs : dict
            Optional extra arguments. Not used by the base class.
        """
        
        if self.lambda_target == 0:
	  cost_wrt_y = 0
	else:
	  cost_wrt_y = self.cost_wrt_target(model,data)
	  
	if self.lambda_teach == 0:
	  cost_wrt_teacher = 0
	else:
	  cost_wrt_teacher = self.cost_wrt_teacher(model,data)
        
	# Compute cost
        cost = self.lambda_target*cost_wrt_y + self.lambda_teach*cost_wrt_teacher 
        
        return T.mean(cost)
        
    def get_monitoring_channels(self, model, data, **kwargs):
        """
        .. todo::

            WRITEME

        .. todo::

            how do you do prereqs in this setup? (I think PL changed
            it, not sure if there still is a way in this context)

        Returns a dictionary mapping channel names to expressions for
        channel values.

        Parameters
        ----------
        model : Model
            the model to use to compute the monitoring channels
        data : batch
            (a member of self.get_data_specs()[0])
            symbolic expressions for the monitoring data
        kwargs : dict
            used so that custom algorithms can use extra variables
            for monitoring.

        Returns
        -------
        rval : dict
            Maps channels names to expressions for channel values.
        """
               
	rval = super(KDCost, self).get_monitoring_channels(model,data)
	                
        value_cost_wrt_target = self.cost_wrt_target(model,data)
        if value_cost_wrt_target is not None:
	   name = 'cost_wrt_target'
	   rval[name] = self.lambda_target*T.mean(value_cost_wrt_target)
                
        value_cost_wrt_teacher = self.cost_wrt_teacher(model,data)
        if value_cost_wrt_teacher is not None:
	   name = 'cost_wrt_teacher'
	   rval[name] = self.lambda_teach*T.mean(value_cost_wrt_teacher)
	   
	rval['lambda_teach'] = self.lambda_teach
	   	
        return rval        






