"""
Model class that unites theory with data.
"""

import logging
logger = logging.getLogger('Model_mod')

import copy

import scipy

import SloppyCell
import SloppyCell.Residuals as Residuals
import SloppyCell.Collections as Collections
import SloppyCell.Utility as Utility
from . import KeyedList_mod as KeyedList_mod
KeyedList = KeyedList_mod.KeyedList

_double_epsilon_ = scipy.finfo(scipy.float_).eps
    
class Model:
    """
    A Model object connects a set of experimental data with the objects used to
    model that data.

    Most importantly, a Model can calculate a cost for a given set of 
    parameters, characterizing how well those parameters fit the data contained
    within the model.
    """

    imag_cutoff = 1e-8

    def __init__(self, expts, calcs):
        """
        expts  A sequence of Experiments to be fit to.
        calcs  A sequence of calculation objects referred to by the 
               Experiments.
        """

        self.calcVals = {}
        self.calcSensitivityVals = {}
        self.internalVars = {}
        self.internalVarsDerivs = {}
        self.residuals = KeyedList()

        if isinstance(expts, list):
            expts = Collections.ExperimentCollection(expts)
        elif isinstance(expts, dict):
            expts = Collections.ExperimentCollection(expts.values())
        self.SetExperimentCollection(expts)

        if isinstance(calcs, list):
            calcs = Collections.CalculationCollection(calcs)
        elif isinstance(calcs, dict):
            calcs = Collections.CalculationCollection(calcs.values())
        self.SetCalculationCollection(calcs)

        self.observers = KeyedList()
        self.parameter_bounds = {}

    def compile(self):
        """
        Compile all the calculations contained within the Model.
        """
        for calc in self.get_calcs().values():
            calc.compile()

    def copy(self):
        return copy.deepcopy(self)

    def get_params(self):
        """
        Return a copy of the current model parameters
        """
        return self.calcColl.GetParameters()

    def get_ICs(self):
        """
        Get the initial conditions currently present in a model
        for dynamic variables that are not assigned variables.

        Outputs:
          KeyedList with keys (calcName,varName) --> initialValue
        """
        ics=KeyedList()
        for calcName, calc in self.calcColl.items():
            for varName in calc.dynamicVars.keys():
                if varName in calc.assignedVars.keys(): continue
                ics.set( (calcName,varName), calc.get_var_ic(varName))
        return ics

    def set_ICs(self, ics):
        """
        Sets the initial conditions into the model. Uses the input
        format defined by 'getICs'.

        Inputs:
         ics -- Initial conditions to set in KeyedList form:
                  keys: (calcName, varName) --> intialValue

        Outputs:
         None
        """
        for (calcName, varName), initialValue in ics.items():
            self.calcColl.get(calcName).set_var_ic(varName, initialValue)

    def _evaluate(self, params, T=1):
        """
        Evaluate the cost for the model, returning the intermediate residuals,
        and chi-squared.

        (Summing up the residuals is a negligible amount of work. This 
         arrangment makes notification of observers much simpler.)
        """
        self.params.update(params)
        self.check_parameter_bounds(params)
        self.CalculateForAllDataPoints(params)
        self.ComputeInternalVariables(T)

       
 resvals = [res.GetValue(self.calcVals, self.internalVars, self.params)
                   for res in self.residuals.values()]

        # Occ
asionally it's useful to use residuals with a sqrt(-1) in them,
        #  to get negative squares. Then, however, we might get small imaginary
        #  parts in our results, which this shaves off.
        chisq = scipy.real_if_close(scipy.sum(scipy.asarray(resvals)**2), 
                                    tol=self.imag_cutoff)
        if scipy.isnan(chisq):
            logger.warn('Chi^2 is NaN, converting to Infinity.')
            chisq = scipy.inf
        cost = 0.5 * chisq

        entropy = 0
        for expt, sf_ents in self.internalVars['scaleFactor_entropies'].items():
            for group, ent in sf_ents.items():
                entropy += ent

        self._notify(event = 'evaluation', 
                     resvals = resvals,
                     chisq = chisq,
                     cost = cost, 
                     free_energy = cost-T*entropy,
                     entropy = entropy,
                     params = self.params)

        return resvals, chisq, cost, entropy

    def res(self, params):
        """
        Return the residual values of the model fit given a set of parameters
        """

        return self._evaluate(params)[0]

    def res_log_params(self, log_params):
        """
        Return the residual values given the logarithm of the parameters
        """
        return self.res(scipy.exp(log_params))

    def res_dict(self, params):
        """
        Return the residual values of the model fit given a set of parameters
        in dictionary form.
        """
        return dict(zip(self.residuals.keys(), self.res(params)))

    def chisq(self, params):
        """
        Return the sum of the squares of the residuals for the model
        """
        return self._evaluate(params)[1]

    def redchisq(self, params):
        """
        Return chi-squared divided by the number of degrees of freedom

        Question: Are priors to be included in the N data points?
                  How do scale factors change the number of d.o.f.?
        """
        return self.chisq(params)/(len(self.residuals) - len(self.params))

    def cost(self, params):
        """
        Return the cost (1/2 chisq) of the model
        """
        return self._evaluate(params)[2]

    def cost_log_params(self, log_params):
        """
        Return the cost given the logarithm of the input parameters
        """
        return self.cost(scipy.exp(log_params))

    def free_energy(self, params, T):
        temp, temp, c, entropy = self._evaluate(params, T=T)
        return c - T * entropy

    def _notify(self, **args):
        """
        Call all observers with the given arguments.
        """
        for obs in self.observers:
            obs(**args)

    def attach_observer(self, obs_key, observer):
        """
        Add an observer to be notified by this Model.
        """
        self.observers.set(obs_key, observer)

    def detach_observer(self, obs_key):
        """
        Remove an observer from the Model.
        """
        self.observers.remove_by_key(obs_key)

    def get_observers(self):
        """
        Return the KeyedList of observers for this model.
        """
        return self.observers

    def reset_observers(self):
        """
        Call reset() for all attached observers.
        """
        for obs in self.observers:
            if hasattr(obs, 'reset'):
                obs.reset()

    resDict = res_dict
    # ...


    def AddResidual(self, res):
        self.residuals.setByKey(res.key, res)

    def Force(self, params, epsf, relativeScale=False, stepSizeCutoff=None):
        """
        Force(parameters, epsilon factor) -> list

        Returns a list containing the numerical gradient of the cost with
        respect to each parameter (in the parameter order of the
        CalculationCollection). Each element of the gradient is:
            cost(param + eps) - cost(param - eps)/(2 * eps).
        If relativeScale is False then epsf is the stepsize used (it should
        already be multiplied by typicalValues before Jacobian is called)
        If relativeScale is True then epsf is multiplied by params.
        The two previous statements hold for both scalar and vector valued
        epsf.
        """

        force = []
        params = scipy.array(params)
        
        if stepSizeCutoff==None:
            stepSizeCutoff = scipy.sqrt(_double_epsilon_)
            
        if relativeScale is True:
            eps = epsf *