import collections
import itertools
import warnings

import pandas as pd

from scmcallib.distributions import Distribution


class ParameterSet(object):
    """Hold all the data used to describe how the emulation or calibration is constrained.

    Tune parameters
        These parameters are updated to find an optimal set/distribution of parameters

    Config parameters
        Other parameters which are passed to the SCM. Typically these parameters are used to set the models into a specific
        reproducible configuration.

    Another feature of the ParameterSet is that each run get's a globally unique run number (for the current python
    execution).

    Attributes
    ----------
    state: dict
        The current state. The key-value pairs of state is passed as arguments to any callable config parameters when they are evaluated

    """

    run_counter = itertools.count()

    def __init__(self):
        self.tune_parameters = collections.OrderedDict()
        self._config_parameters = collections.OrderedDict()
        self.state = {}

    def __str__(self):
        s = "Config Parameters:\n"
        for p in self._config_parameters:
            s += "{} = {}\n".format(p, self._config_parameters[p])
        s += "Tuning Parameters:\n"
        for p in self.tune_parameters:
            s += "{} = {}\n".format(p, self.tune_parameters[p])
        return s

    def set_tune(self, name, distribution, x0=None):
        """Define a tune parameter

        These are also know as priors, the prior knowledge of the parameter space. These parameters are allowed to change
        when calibrating or emulating.

        When a parameter set is used to emulate another model, the distribution used may not influence the find
        parameters found depending on how the model is setup.

        Parameters
        ----------
        name :
            Name of the parameter. This is the name which is passed to the SCM instance.
        distribution :
            An instance of a subclass of `scmcallib.distributions.Distribution`. This variable defines
            how the parameter can change which emulating or calibration.
        """
        # TODO: handle x0
        if not isinstance(distribution, Distribution):
            raise ValueError(
                "distributions must inherit scmcallib.distributions.Distribution"
            )

        if name in self.config_parameters:
            raise ValueError("{} is already defined as a config parameter")

        self.tune_parameters[name] = distribution

    def set_config(self, name, value):
        """Define a config parameter

        Config parameters are used to set the model into a specific setup for finding point estimates/distributions. They remain
        (mostly) fixed throughout the model runs. These are typically scalars or strings, but functions can also be passed. These
        functions are evaluated on calls to the `config_parameters` property and are passed the contents of `state`.

        Parameters
        ----------
        name :
            Name of the config parameter. This is the name which is used when the value is passed to the initialiser
            of a SCM instance.
        value :
            string or callable. The callables are passed the contents of `ParameterSet.state` as kwargs when they are
            evaluated in `config_parameters`.
        """
        if name in self.tune_parameters:
            raise ValueError("{} is already defined as a config parameter")
        self._config_parameters[name] = value

    def evaluate(self, count, include_config=False):
        """Evaluate the parameters

        Parameters
        ----------
        count :
            Number of parameter sets to sample
        include_config :
            If True include the fixed parameters in the parameter dataframe (Default value = False)

        Returns
        -------
        pd.DataFrame
            DataFrame which contains a set of parameters generated by drawing from the parameter sets
            specified.

        """
        params = collections.OrderedDict(
            {k: self.tune_parameters[k].evaluate(count) for k in self.tune_parameters}
        )
        if include_config:
            params.update(self.config_parameters)

        return pd.DataFrame(
            params, index=[next(self.run_counter) for _ in range(count)]
        )

    @property
    def names(self):
        return self.tune_parameters.keys()

    @property
    def config_parameters(self):
        """Evaluate the config_parameters

        Returns
        -------
        collections.OrderedDict
            If the config parameters are functions they are also evaluated. Any functions that cannot be evaulated then they are excluded
            from the result and a warning is printed.
        """
        res = collections.OrderedDict()

        for k in self._config_parameters:
            v = self._config_parameters[k]
            if callable(v):
                try:
                    v = v(**self.state)
                except TypeError:
                    warnings.warn(
                        "could not evaluate config parameter `{}`. Check `ParameterSet.state`".format(
                            k
                        )
                    )
                    continue
            res[k] = v
        return res

    @property
    def model_specific_parameters(self):
        """Evaluate the config_parameters

        Returns
        -------
        collections.OrderedDict
            If the config parameters are functions they are also evaluated. Any functions that cannot be evaulated then they are excluded
            from the result and a warning is printed.
        """
        res = collections.OrderedDict()

        for k in self._config_parameters:
            v = self._config_parameters[k]
            if callable(v):
                try:
                    v = v(**self.state)
                except TypeError:
                    warnings.warn(
                        "could not evaluate config parameter `{}`. Check `ParameterSet.state`".format(
                            k
                        )
                    )
                    continue
                res[k] = v
        return res


class KeyedParameterSet(ParameterSet):
    """
    ParameterSet where the items being iterated over are already known

    set_config can also handle dicts which are keyed by the values

    This parameter set should be used in preference to ParameterSet if state is changing. That functionality will be removed
    from ParameterSet in future.
    """

    key = None

    def __init__(self, items, key=None):
        super().__init__()
        if key is not None:
            self.key = key
        self.items = items
        self.state = {self.key: self.items[0]}

    def __str__(self):
        s = "Config Parameters:\n"
        vals = {k: self.get_config_parameters(k) for k in self.items}
        df = pd.DataFrame(vals)
        df.columns.name = self.key
        s += str(df)
        s += "\nTuning Parameters:\n"
        for p in self.tune_parameters:
            s += "{} = {}\n".format(p, self.tune_parameters[p])
        return s

    def set_config(self, name, value):
        if callable(value):
            warnings.warn("Deprecating callable config parameters", DeprecationWarning)
        else:
            super().set_config(name, value)

    set_config.__doc__ = ParameterSet.set_config.__doc__

    def get_config_parameters(self, key=None):
        if key is None:
            key = self.state[self.key]
        res = collections.OrderedDict()

        for k in self._config_parameters:
            v = self._config_parameters[k]
            if isinstance(v, dict):
                try:
                    v = v[key]
                except KeyError:
                    warnings.warn(
                        "could not evaluate config parameter `{}`. Check `ParameterSet.state`".format(
                            k
                        )
                    )
                    continue
            elif callable(v):
                warnings.warn(
                    "Deprecating callable config parameters", DeprecationWarning
                )
                continue
            res[k] = v
        return res

    @property
    def config_parameters(self):
        return self.get_config_parameters()

    config_parameters.__doc__ = ParameterSet.config_parameters.__doc__


class ScenarioParameterSet(KeyedParameterSet):
    key = "scenario"


class ModelParameterSet(KeyedParameterSet):
    key = "model"
