from typing import Dict, List, Tuple

from ... import Model, ModelParameter
from ...common import TX, TY


class ReparametrizedModel(Model):
    """Model formed by reparametrizing an existing :class Model:.

    Reparametrizing a model involves replacing "bound" parameters with "new" parameters,
    whose values the bound parameter values are calculated from.

    All non-bound parameters of the original model as well as all "new" parameters are
    parameters of the :class ReparametrizedModel: (bound parameters are internal
    parameters of the new model, but are not directly exposed to the user). All derived
    results from the original model are derived results from the
    :class ReparametrizedModel: (override :meth calculate_derived_params: to change this
    behaviour).

    Values and uncertainties for the bound parameters are exposed as derived results.

    Subclasses must override :meth bound_param_values:, :meth model_uncertainty_func:
    and :meth reparametrized_value_func: to specify the mapping between "new" and
    "bound" parameters.
    """

    def __init__(
        self,
        model: Model,
        new_params: Dict[str, ModelParameter],
        bound_params: List[str],
    ):
        """
        :param model: The model to be reparametrized. This model is considered "owned"
          by the :class ReparametrizedModel: and should not be used / modified
          elsewhere.
        :param new_params: dictionary of new parameters of the
          :class ReparametrizedModel:
        :param bound_params: list of parameters of the :class Model: to be
          reparametrized. These parameters are not exposed as parameters of the
          :class ReparametrizedModel:.
        """
        self.model = model
        self.bound_params = bound_params
        self.unbound_model_params = [
            param_name
            for param_name in self.model.parameters.keys()
            if param_name not in bound_params
        ]

        internal_parameters = self.model.internal_parameters
        parameters = dict(self.model.parameters)

        if not set(bound_params).issubset(set(parameters.keys())):
            raise ValueError("Bound parameters must all be parameters of the model")

        for param_name in bound_params:
            internal_parameters.append(parameters.pop(param_name))

        duplicates = set(new_params.keys()).intersection(parameters.keys())
        if duplicates:
            raise ValueError(
                "New parameter names must not duplicate names of unbound model "
                f"parametes. Duplicates are {duplicates}."
            )

        parameters.update(new_params)

        super().__init__(parameters=parameters, internal_parameters=internal_parameters)

    @staticmethod
    def bound_param_values(param_values: Dict[str, float]) -> Dict[str, float]:
        """Returns a dictionary of values of the model's bound parameters.

        This method must be overridden to specify the mapping from parameters of the
        :class ReparameterizedModel: to values of the bound parameters.

        :param new_param_values: dictionary of parameter values for the
          :class ReparameterizedModel:.
        :returns: dictionary of values for the bound parameters of the original model.
        """
        raise NotImplementedError

    @staticmethod
    def bound_param_uncertainties(
        param_values: Dict[str, float], param_uncertainties: Dict[str, float]
    ) -> Dict[str, float]:
        """Returns a dictionary of uncertainties for the model's bound parameters.

        This method must be overridden to specify the mapping from parameter
        uncertainties for the :class ReparameterizedModel: to bound parameter
        uncertainties.

        :param param_values: dictionary of values for parameters of the
          :class ReparameterizedModel:.
        :param param_uncertainties: dictionary of uncertainties for parameters of the
          :class ReparameterizedModel:.
        :returns: dictionary of values for the bound parameters of the original model.
        """
        raise NotImplementedError

    @staticmethod
    def new_param_values(model_param_values: Dict[str, float]) -> Dict[str, float]:
        """Returns a dictionary of values of the model's "new" parameters.

        This method must be overridden to specify the mapping from values of the
        original model to values of the :class ReparametrizedModel:'s "new" parameters.

        This is used to find estimates for the new parameters from on the original
        model's parameter estimates.

        :param model_param_values: dictionary of parameter values for the original model
        :returns: dictionary of values for the "new" parameters of the
          :class ReparametrizedModel:
        """
        raise NotImplementedError

    def get_num_x_axes(self) -> int:
        return self.model.get_num_x_axes()

    def get_num_y_axes(self) -> int:
        return self.model.get_num_y_axes()

    def can_rescale(self) -> Tuple[List[bool], List[bool]]:
        return self.model.can_rescale()

    def func(self, x: TX, param_values: Dict[str, float]) -> TY:
        model_values = self.bound_param_values(param_values)
        model_values.update(
            {
                param_name: param_values[param_name]
                for param_name in self.unbound_model_params
            }
        )
        return self.model.func(x, model_values)

    def estimate_parameters(self, x: TX, y: TX):
        self.model.estimate_parameters(x=x, y=y)
        model_heuristics = {
            param_name: param_data.get_initial_value()
            for param_name, param_data in self.model.parameters.items()
        }
        new_heuristics = self.new_param_values(model_heuristics)

        for param_name in self.unbound_model_params:
            self.parameters[param_name].heuristic = model_heuristics[param_name]

        for param_name, param_heuristic in new_heuristics.items():
            self.parameters[param_name].heuristic = param_heuristic

    def calculate_derived_params(
        self,
        x: TX,
        y: TY,
        fitted_params: Dict[str, float],
        fit_uncertainties: Dict[str, float],
    ) -> Tuple[Dict[str, float], Dict[str, float]]:

        bound_values = self.bound_param_values(fitted_params)
        bound_uncertainties = self.bound_param_uncertainties(
            fitted_params, fit_uncertainties
        )

        model_fitted_params = {
            param_name: fitted_params[param_name]
            for param_name in self.unbound_model_params
        }
        model_fit_uncertainties = {
            param_name: fit_uncertainties[param_name]
            for param_name in self.unbound_model_params
        }

        model_fitted_params.update(bound_values)
        model_fit_uncertainties.update(bound_uncertainties)

        derived_values, derived_uncertainties = self.model.calculate_derived_params(
            x=y,
            y=y,
            fitted_params=model_fitted_params,
            fit_uncertainties=model_fit_uncertainties,
        )

        derived_values.update(bound_values)
        derived_uncertainties.update(bound_uncertainties)

        return derived_values, derived_uncertainties
