# standard library
from typing import Union, List, Optional

# third party imports
import numpy as np

# local imports
from probeye.definition.sensor import Sensor
from probeye.subroutines import make_list, translate_prms_def


class NoiseModelBase:
    def __init__(
        self,
        dist: str,
        prms_def: Union[str, List[Union[str, dict]], dict],
        sensors: Union[Sensor, List[Sensor]],
        name: Optional[str] = None,
        corr: Optional[str] = None,
        corr_model: Optional[str] = None,
        noise_type: str = "additive",
    ):
        """
        Parameters
        ----------
        dist
            A string specifying the probability distribution the noise model is based
            on, e.g. 'normal'.
        prms_def
            A list of parameter names (strings) defining how a noise parameter vector
            given to the loglike_contribution method is interpreted. For example:
            prms_def = ['mu', 'sigma'] means that the noise parameter vector has two
            elements, the first of which gives the value of 'mu' and the second gives
            the value of 'sigma'.
        sensors
            Sensor objects that are required to evaluate the noise model.
        name
            Unique name of the noise model. This name is None, if the user does not
            specify it when adding the noise model to the problem. It is then named
            automatically before starting the inference engine.
        corr
            Defines the correlation model. So far this is just a placeholder. It is not
            clear yet how exactly the correlation should be defined. When it is set to
            None, all sensors/sensor elements are independent.
        corr_model
            Defines the correlation function to be used in case corr isn't None.
        noise_type
            Either 'additive', 'multiplicative' or 'other'. Defines if the error is
            computed via [prediction - measurement] ('additive') or via [prediction/
            measurement-1] ('multiplicative') or in some 'other' i.e.,
            non-standard fashion.
        """
        self.dist = dist
        self.prms_def, self.prms_dim = translate_prms_def(prms_def)
        self.sensors = make_list(sensors)
        self.sensor_names = [sensor.name for sensor in self.sensors]
        self.name = name
        self.corr = corr
        self.corr_model = corr_model
        self.noise_type = noise_type
        # this is a list of experiment names, that relate to the noise model; the list
        # will be filled after experiments have been added to the InferenceProblem and
        # the problem definition is complete; in this case call InferenceProblem.
        # assign_experiments_to_noise_models() to fill it with the corresponding names
        self.experiment_names = []  # type: List[str]
        # as soon as defined, this attribute will be a pointer to the inference
        # problems experiments (it will be used for consistency checks)
        self.problem_experiments = {}  # type: dict

        # set the error_function depending on the noise-type
        if noise_type == "additive":
            self.error_function = self.error_function_additive
        elif noise_type == "multiplicative":
            self.error_function = self.error_function_multiplicative
        elif noise_type == "other":
            self.error_function = self.error_function_other
        else:
            raise ValueError(
                f"Encountered unknown noise_type: '{noise_type}'. The noise_type must "
                f"be either 'additive', 'multiplicative' or 'other'."
            )

    def add_experiments(self, experiment_names_: Union[str, List[str]]):
        """
        Adds experiment names to the noise model. When the noise model is evaluated it
        will only be evaluated for those experiments added here.

        Parameters
        ----------
        experiment_names_
            Names (strings) of experiments from the InferenceProblem that should be
            added to the noise model.
        """
        # check if the given experiments are compatible with the noise model with
        # respect to the sensors
        experiment_names = make_list(experiment_names_)
        forward_models = set()
        for exp_name in experiment_names:
            exp_dict = self.problem_experiments[exp_name]
            forward_models.add(exp_dict["forward_model"])
            sensor_names_exp = [*exp_dict["sensor_values"].keys()]
            for sensor_name in self.sensor_names:
                if sensor_name not in sensor_names_exp:
                    raise RuntimeError(
                        f"Experiment '{exp_name}' does not contain a sensor "
                        f"'{sensor_name}' which is required for the evaluation of the "
                        f"noise model."
                    )
        # check if the given experiments all refer to one forward model
        if len(forward_models) > 1:
            raise RuntimeError(
                f"The given experiments refer to more than one forward model!"
            )
        # check if one of the given experiments have been added before
        for exp_name in experiment_names:
            if exp_name in self.experiment_names:
                raise RuntimeError(
                    f"The experiment '{exp_name}' has already been added to this noise "
                    f"model. Something might be wrong here."
                )
        self.experiment_names += experiment_names

    def error(self, model_response_dict: dict) -> dict:
        """
        Computes the model error for all of the noise model's experiments and returns
        them in a dictionary that is sorted by output sensor_values.

        Parameters
        ----------
        model_response_dict
            The first key is the name of the experiment. The values are dicts which
            contain the forward model's output sensor's names as keys have the
            corresponding model responses as values.

        Returns
        -------
        model_error
            A dictionary with the keys being the noise model's sensor names, and 1D
            numpy arrays representing the model errors as values.
        """
        # prepare the dictionary keys
        model_error_dict = {name: np.array([]) for name in self.sensor_names}

        # fill the dictionary with model error vectors
        for exp_name in self.experiment_names:
            exp_dict = self.problem_experiments[exp_name]
            ym_dict = model_response_dict[exp_name]
            ye_dict = exp_dict["sensor_values"]
            me_dict = self.error_function(ym_dict, ye_dict)
            model_error_dict = {
                name: np.append(model_error_dict[name], me_dict[name])
                for name in self.sensor_names
            }

        return model_error_dict

    def error_function_additive(self, ym_dict: dict, ye_dict: dict) -> dict:
        """
        Evaluates the additive model error for each of the noise model' sensors.

        Parameters
        ----------
        ym_dict
            The computed values for the model's output sensor_values.
        ye_dict
            The measured values for the model's output sensor_values.

        Returns
        -------
        error_dict
            The computed model error for the model's output sensor_values.
        """
        # for each sensor, its own error metric is used to compute the error
        error_dict = {name: ym_dict[name] - ye_dict[name] for name in self.sensor_names}
        return error_dict

    def error_function_multiplicative(self, ym_dict: dict, ye_dict: dict) -> dict:
        """
        Evaluates the multiplicative model error for each of the noise model's sensors.

        Parameters
        ----------
        ym_dict
            The computed values for the model's output sensor_values.
        ye_dict
            The measured values for the model's output sensor_values.

        Returns
        -------
        error_dict
            The computed model error for the model's output sensor_values.
        """
        # for each sensor, its own error metric is used to compute the error
        error_dict = {
            name: ym_dict[name] / ye_dict[name] - 1.0 for name in self.sensor_names
        }
        return error_dict

    def error_function_other(self, ym_dict: dict, ye_dict: dict) -> dict:
        """
        Non-standard error function self.error_function will point to when self.
        noise_type is set to 'other'. See self.error_function for more information.
        """
        raise NotImplementedError(
            "Your model does not have an non-standard error_function-method yet. If "
            "you want to use it, you need to implement it first."
        )

    def loglike_contribution(
        self, model_response_dict: dict, prms: dict, worst_value: float = -np.infty
    ) -> float:
        """
        Evaluates the log-likelihood function for the given model error and the given
        noise parameter vector. This method has to be overwritten.

        Parameters
        ----------
        model_response_dict
            The first key is the name of the experiment. The values are dicts which
            contain the forward model's output sensor's names as keys have the
            corresponding model responses as values.
        prms
            Dictionary containing parameter name:value pairs.
        worst_value
            This value is returned when this method does not result in a numeric value.
            This might happen for example when the given parameters are not valid (for
            example in case of a negative standard deviation). The returned value in
            such cases should represent the worst possible value of the contribution.

        Returns
        -------
        ll
            The evaluated log-likelihood function.
        """
        raise NotImplementedError(
            "Your model does not have a loglike_contribution-method. You need to "
            "define this method so you can evaluate your noise model."
        )


class NormalNoiseModel(NoiseModelBase):
    """
    A general Gaussian (normal) noise model with or without correlations.
    """

    def __init__(
        self,
        prms_def: Union[str, List[Union[str, dict]], dict],
        sensors: Union[Sensor, List[Sensor]],
        name: Optional[str] = None,
        corr: Optional[str] = None,
        corr_model: Optional[str] = None,
        noise_type: str = "additive",
    ):
        """
        See docstring of NoiseModelBase for information on the input arguments.
        """

        # initialize the base class with given input
        super().__init__(
            "normal",
            prms_def,
            sensors,
            name=name,
            corr=corr,
            corr_model=corr_model,
            noise_type=noise_type,
        )

        # check that at the standard deviation is provided (this can be either as a
        # constant or a latent parameter, but it has to be given)
        if "std" not in [*self.prms_def.values()]:
            raise RuntimeError(
                "The standard deviation 'std' was not provided in prms_def!"
            )

        # the mean value(s) do not have to be stated explicitly; if they are not given,
        # the are assumed to be zero
        self.zero_mean = True
        if "mean" in [*self.prms_def.values()]:
            self.zero_mean = False
