"""
.. module:: parameterization

:Synopsis: Class managing the possibly different parameterizations
           used by sampler and likelihoods
:Author: Jesus Torrado

"""
# Global
import numpy as np
from numbers import Real
from itertools import chain
from copy import deepcopy
from typing import Mapping, Sequence, Dict, Set, List, Tuple, Any, Callable, Union

# Local
from cobaya.typing import ParamsDict, ParamDict, ParamInput, \
    ExpandedParamsDict, ParamValuesDict, partags
from cobaya.tools import get_external_function, ensure_nolatex, is_valid_variable_name, \
    getfullargspec, deepcopy_where_possible, invert_dict, str_to_list
from cobaya.log import LoggedError, HasLogger


def is_fixed_or_function_param(info_param: ParamInput) -> bool:
    """
    Returns True if the parameter has been fixed to a value or through a function.
    """
    return expand_info_param(info_param).get("value") is not None


def is_sampled_param(info_param: ParamInput) -> bool:
    """
    Returns True if the parameter has a prior.
    """
    return "prior" in expand_info_param(info_param)


def is_derived_param(info_param: ParamInput) -> bool:
    """
    Returns True if the parameter is saved as a derived one.
    """
    return expand_info_param(info_param).get("derived", False) is not False


def expand_info_param(info_param: ParamInput, default_derived=True) -> ParamDict:
    """
    Expands the info of a parameter, from the user-friendly, shorter format
    to a more unambiguous one.
    """
    info_param = deepcopy_where_possible(info_param)
    if not isinstance(info_param, Mapping):
        if info_param is None:
            info_param = {}
        elif isinstance(info_param, Sequence) and not isinstance(info_param, str):
            values = list(info_param)
            allowed_lengths = [2, 4, 5]
            if len(values) not in allowed_lengths:
                raise LoggedError(__name__,
                                  "Parameter info length not valid: %d. "
                                  "The allowed lengths are %r. See documentation.",
                                  len(values), allowed_lengths)
            info_param = {"prior": [values[0], values[1]]}
            if len(values) >= 4:
                info_param["ref"] = [values[2], values[3]]
                if len(values) == 5:
                    info_param["proposal"] = values[4]
        else:
            info_param = {"value": info_param}
    if all(f not in info_param for f in ["prior", "value", "derived"]):
        info_param["derived"] = default_derived
    # Dynamical input parameters: save as derived by default
    value = info_param.get("value")
    if isinstance(value, str) or callable(value):
        info_param["derived"] = info_param.get("derived", True)
    # noinspection PyTypeChecker
    return info_param


def reduce_info_param(info_param: ParamDict) -> ParamInput:
    """
    Compresses the info of a parameter, suppressing default values.
    This is the opposite of :func:`~input.expand_info_param`.
    """
    info_param = deepcopy_where_possible(info_param)
    if not isinstance(info_param, dict):
        return None
    # All parameters without a prior are derived parameters unless otherwise specified
    if info_param.get("derived") is True:
        info_param.pop("derived")
    # Fixed parameters with single "value" key
    if list(info_param) == ["value"] and not callable(info_param["value"]):
        return info_param["value"]
    return info_param


_WrappedFunc = Tuple[Callable, Dict[str, Any], List[str]]


class Parameterization(HasLogger):
    """
    Class managing parameterization.
    Translates parameter between sampler+prior and likelihood
    """

    def __init__(self, info_params: Union[ParamsDict, ExpandedParamsDict],
                 allow_renames=True,
                 ignore_unused_sampled=False):
        self.set_logger()
        self.allow_renames = allow_renames
        # First, we load the parameters,
        # not caring about whether they are understood by any likelihood.
        # `input` contains the parameters (expected to be) understood by the likelihood,
        #   with its fixed value, its fixing function, or None if their value is given
        #   directly by the sampler.
        self._infos = {}
        self._input: ParamValuesDict = {}
        self._input_funcs = {}
        self._input_args = {}
        self._input_dependencies: Dict[str, Set[str]] = {}
        self._dropped: Set[str] = set()
        self._output: ParamValuesDict = {}
        self._constant: ParamValuesDict = {}
        self._sampled: ParamValuesDict = {}
        self._sampled_renames: Dict[str, List[str]] = {}
        self._derived: ParamValuesDict = {}
        self._derived_inputs = []
        self._derived_funcs = {}
        self._derived_args = {}
        self._derived_dependencies: Dict[str, Set[str]] = {}
        # Notice here that expand_info_param *always* adds a "derived":True tag
        # to infos without "prior" or "value", and a "value" field
        # to fixed params
        for p, info in info_params.items():
            if isinstance(info, Mapping) and not set(info).issubset(partags):
                raise LoggedError(self.log, "Parameter '%s' has unknown options %s",
                                  p, set(info).difference(partags))
            info = expand_info_param(info)
            self._infos[p] = info
            if is_fixed_or_function_param(info):
                if isinstance(info["value"], Real):
                    self._constant[p] = float(info["value"])
                    self._input[p] = self._constant[p]
                    if info.get("drop"):
                        self._dropped.add(p)
                else:
                    self._input[p] = np.nan
                    self._input_funcs[p] = get_external_function(info["value"])
                    self._input_args[p] = getfullargspec(self._input_funcs[p]).args
            if is_sampled_param(info):
                self._sampled[p] = np.nan
                self._input[p] = np.nan
                if info.get("drop"):
                    self._dropped.add(p)
                self._sampled_renames[p] = str_to_list(info.get("renames") or [])
            if is_derived_param(info):
                self._derived[p] = np.nan
                # Dynamical parameters whose value we want to save
                if info["derived"] is True and is_fixed_or_function_param(info):
                    # parameters that are already known or computed by input funcs
                    self._derived_inputs.append(p)
                elif info["derived"] is True:
                    self._output[p] = np.nan
                else:
                    self._derived_funcs[p] = get_external_function(info["derived"])
                    self._derived_args[p] = getfullargspec(self._derived_funcs[p]).args
        # Check that the sampled and derived params are all valid python variable names
        for p in chain(self._sampled, self._derived):
            if not is_valid_variable_name(p):
                is_in = p in self._sampled
                eg_in = "  p_prime:\n    prior: ...\n  %s: " \
                        "'lambda p_prime: p_prime'\n" % p
                eg_out = "  p_prime: 'lambda %s: %s'\n" % (p, p)
                raise LoggedError(
                    self.log, "Parameter name '%s' is not a valid Python variable name "
                              "(it needs to start with a letter or '_').\n"
                              "If this is an %s parameter of a likelihood or theory, "
                              "whose name you cannot change,%s define an associated "
                              "%s one with a valid name 'p_prime' as: \n\n%s",
                    p, "input" if is_in else "output",
                    "" if is_in else " remove it and",
                    "sampled" if is_in else "derived",
                    eg_in if is_in else eg_out)

        # input params depend on input and sampled only,
        # never on output/derived unless constant
        known_input = set(self._input)
        all_input_arguments = set(chain(*self._input_args.values()))
        bad_input_dependencies = all_input_arguments - known_input
        if bad_input_dependencies:
            raise LoggedError(
                self.log,
                "Input parameters defined as functions can only depend on other "
                "input parameters. In particular, an input parameter cannot depend on %r."
                " Use an explicit Theory calculator for more complex dependencies.\n"
                "If you intended to define a derived output parameter use derived: "
                "instead of value:", list(bad_input_dependencies))

        # Assume that the *un*known function arguments are likelihood/theory
        # output parameters
        for arg in (all_input_arguments.union(*self._derived_args.values())
                .difference(known_input).difference(self._derived)):
            self._output[arg] = np.nan

        # Useful set: directly "output-ed" derived
        self._directly_output = [p for p in self._derived if p in self._output]

        self._wrapped_input_funcs, self._wrapped_derived_funcs = \
            self._get_wrapped_functions_evaluation_order()

        # Useful mapping: input params that vary if each sample is varied
        self._sampled_input_dependence = {s: [i for i in self._input
                                              if s in self._input_dependencies.get(i, {})]
                                          for s in self._sampled}
        # From here on, some error control.
        # Only actually raise error after checking if used by prior.
        if not ignore_unused_sampled:
            self._dropped_not_directly_used = self._dropped.intersection(
                p for p, v in self._sampled_input_dependence.items() if not v)
        else:
            self._dropped_not_directly_used = set()

        # warn if repeated labels
        labels_inv_repeated = invert_dict(self.labels())
        labels_inv_repeated = {k: v for k, v in labels_inv_repeated.items() if len(v) > 1}
        if labels_inv_repeated:
            self.mpi_warning(
                "There are repeated parameter labels: %r", labels_inv_repeated)

    def dropped_param_set(self) -> Set[str]:
        return self._dropped.copy()

    def input_params(self) -> ParamValuesDict:
        return self._input.copy()

    def output_params(self) -> ParamValuesDict:
        return self._output.copy()

    def constant_params(self) -> ParamValuesDict:
        return self._constant.copy()

    def sampled_params(self) -> ParamValuesDict:
        return self._sampled.copy()

    def sampled_params_info(self) -> ExpandedParamsDict:
        return {p: deepcopy_where_possible(info) for p, info
                in self._infos.items() if p in self._sampled}

    def sampled_params_renames(self) -> Dict[str, List[str]]:
        return deepcopy(self._sampled_renames)

    def derived_params(self) -> ParamValuesDict:
        return self._derived.copy()

    def derived_params_info(self) -> ExpandedParamsDict:
        return {p: deepcopy_where_possible(info) for p, info
                in self._infos.items() if p in self._derived}

    def sampled_input_dependence(self) -> Dict[str, List[str]]:
        return deepcopy(self._sampled_input_dependence)

    def get_input_func(self, p, **params_values):
        func = self._input_funcs[p]
        args = self._input_args[p]
        return func(*[params_values.get(arg) for arg in args])

    @property
    def input_dependencies(self) -> Dict[str, Set[str]]:
        return self._input_dependencies

    def to_input(self, sampled_params_values) -> ParamValuesDict:
        # Gets all current sampled and input derived parameters as a dictionary,
        # including dropped parameters. Result is not a copy and must not be modified.

        # Store sampled params, so that derived can depend on them
        if not isinstance(sampled_params_values, dict):
            sampled_params_values = dict(zip(self._sampled, sampled_params_values))
        else:
            sampled_params_values = sampled_params_values.copy()
        self._sampled = sampled_params_values

        # First include all sampled input parameters,
        self._input.update(sampled_params_values)

        if self._wrapped_input_funcs:
            # Then evaluate the functions
            for p, (func, args, to_set) in self._wrapped_input_funcs.items():
                for arg in to_set:
                    args[arg] = self._input.get(arg, sampled_params_values.get(arg))
                self._input[p] = self._call_param_func(p, func, args)
        return self._input

    def to_derived(self, output_params_values) -> ParamValuesDict:
        if not isinstance(output_params_values, dict):
            output_params_values = dict(zip(self._output, output_params_values))
        # Fill first derived parameters which are direct output parameters
        for p in self._directly_output:
            self._derived[p] = output_params_values[p]
        for p in self._derived_inputs:
            self._derived[p] = self._input[p]
        # Then evaluate the functions
        if self._wrapped_derived_funcs:
            for p, (func, args, to_set) in self._wrapped_derived_funcs.items():
                for arg in to_set:
                    if (val := self._input.get(arg)) is None:
                        if (val := output_params_values.get(arg)) is None:
                            val = self._derived.get(arg)
                    args[arg] = val
                self._derived[p] = self._call_param_func(p, func, args)
        return self._derived

    def check_sampled(self, sampled_params: Union[Sequence[float], Dict[str, float]]
                      ) -> Union[Sequence[float], Dict[str, float]]:
        """
        Performs some checks on the given sampled params.

        If an array is passed, the only test performed is for the right amount of
        parameters, and the same array is returned if successful.

        If a dictionary is passed, it checks that it contains all the sampled parameters,
        and just them. This function is aware of known renamings. Returns dict of
        parameters (model's naming, not remanes) and their values.
        """
        if sampled_params is None:  # only works if there are no sampled params
            sampled_params = []
        if hasattr(sampled_params, "keys"):
            return self.check_sampled_dict(**sampled_params)
        else:
            if len(sampled_params) != len(self._sampled):
                raise LoggedError(self.log, "Wrong number of sampled parameters passed: "
                                            "%d given vs %d expected",
                                  len(sampled_params),
                                  len(self._sampled))
            return sampled_params

    def check_sampled_dict(self, **sampled_params) -> ParamValuesDict:
        """
        Check that the input dictionary contains all the sampled parameters,
        and just them.

        This function is aware of known renamings.

        Returns dict of parameters (model's naming, not renames) and their values.
        """
        sampled_output: ParamValuesDict = {}
        for p, renames in self._sampled_renames.items():
            for pprime in sampled_params:
                if pprime == p or (pprime in renames if self.allow_renames else False):
                    sampled_output[p] = sampled_params.pop(pprime)
                    break
        if len(sampled_output) < len(self._sampled):
            not_found = set(self._sampled).difference(sampled_output)
            if self.allow_renames:
                msg = ("The following expected sampled parameters " +
                       ("(or their aliases) " if self.allow_renames else "") +
                       "were not found : %r",
                       ({p: self._sampled_renames[p] for p in not_found}
                        if self.allow_renames else not_found))
            else:
                msg = ("The following expected sampled parameters "
                       "were not found : %r",
                       {p: self._sampled_renames[p] for p in not_found})
            raise LoggedError(self.log, *msg)
        # Ignore fixed input parameters if they have the correct value
        not_used = set(sampled_params)
        for p, value in sampled_params.items():
            known_value = self._constant.get(p)
            if known_value is None:
                raise LoggedError(self.log, "Unknown parameter %r.", p)
            elif np.allclose(value, known_value):
                not_used.remove(p)
                self.log.debug("Fixed parameter %r ignored.", p)
            else:
                raise LoggedError(
                    self.log, "Cannot change value of constant parameter: "
                              "%s = %g (new) vs %g (old).", p, value, known_value)
        if not_used:
            duplicated = not_used.intersection(
                chain(
                    *[list(chain(*[[k], v])) for k, v in self._sampled_renames.items()]))
            not_used = not_used.difference(duplicated)
            derived = not_used.intersection(self._derived)
            input_ = not_used.intersection(self._input)
            unknown = not_used.difference(derived).difference(input_)
            msg_text = ("Incorrect parameters! " +
                        ("\n   Duplicated entries (using their aliases): %r" %
                         list(duplicated) if duplicated else "") +
                        ("\n   Not known: %r" % list(unknown) if unknown else "") +
                        ("\n   Cannot be fixed: %r " % list(input_) +
                         "--> instead, fix sampled parameters that depend on them!"
                         if input_ else "") +
                        ("\n   Cannot be fixed because are derived parameters: %r " %
                         list(derived) if derived else ""))
            for line in msg_text.split("\n"):
                self.log.error(line)
            raise LoggedError
        return sampled_output

    def check_dropped(self, external_dependence):
        # some error control, given external_dependence from prior
        # only raise error after checking not used by prior
        if self._dropped_not_directly_used.difference(external_dependence):
            raise LoggedError(
                self.log,
                "Parameters %r are sampled but not passed to a likelihood or theory "
                "code, and never used as arguments for any prior or parameter "
                "functions. Check that you are not using "
                "the '%s' tag unintentionally.",
                list(self._dropped_not_directly_used), "drop")

    def labels(self) -> Dict[str, str]:
        """
        Returns a dictionary of LaTeX labels of the sampled and derived parameters.

        Uses the parameter name if no label has been given.
        """

        def get_label(p, info):
            return ensure_nolatex(getattr(info, "get", lambda x, y: y)
                                  ("latex", p.replace("_", r"\ ")))

        return {p: get_label(p, info) for p, info in self._infos.items()}

    def _call_param_func(self, p, func, kwargs):
        try:
            return func(**kwargs)
        except NameError as exception:
            unknown = str(exception).split("'")[1]
            raise LoggedError(
                self.log, "Unknown variable '%s' was referenced in the definition of "
                          "the parameter '%s', with arguments %r.", unknown, p,
                list(kwargs))
        except:
            self.log.error("Function for parameter '%s' failed at evaluation "
                           "and threw the following exception:", p)
            raise

    def _get_wrapped_functions_evaluation_order(self):
        # get evaluation order for input and derived parameter function
        # and pre-prepare argument dicts

        wrapped_funcs: Tuple[Dict[str, _WrappedFunc], Dict[str, _WrappedFunc]] = ({}, {})
        known = set(chain(self._constant, self._sampled))

        for derived, wrapped_func in zip((False, True), wrapped_funcs):
            if derived:
                inputs = self._derived_funcs.copy()
                input_args = self._derived_args
                known.update(self._output)
                output = self._derived
                dependencies = self._derived_dependencies
            else:
                inputs = self._input_funcs.copy()
                input_args = self._input_args
                output = self._input
                dependencies = self._input_dependencies

            while inputs:
                for p, func in inputs.items():
                    args = input_args[p]
                    if not known.issuperset(args):
                        continue
                    known.add(p)
                    dependencies[p] = set(
                        chain(args, *(dependencies.get(arg, []) for arg in args)))

                    if set(args).issubset(self._constant):
                        # all inputs are constant, so output is constant and precomputed
                        self._constant[p] = \
                            self._call_param_func(p, func,
                                                  {arg: self._constant[arg] for arg in
                                                   args})
                        output[p] = self._constant[p]
                    else:
                        # Store function, argument dict with constants pre-filled,
                        # and unset args as tuple
                        wrapped_func[p] = \
                            (func, {arg: self._constant.get(arg) for arg in args},
                             [arg for arg in args if arg not in self._constant])
                    del inputs[p]
                    break
                else:
                    raise LoggedError(
                        self.log, "Could not resolve arguments for parameters %s. "
                                  "Maybe there is a circular dependency between derived "
                                  "parameters?", list(inputs))
        return wrapped_funcs
