"""
Model runner class
"""
from __future__ import annotations

from functools import partial
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Dict,
    Iterable,
    Optional,
    Protocol,
    Tuple,
    Union,
)

import numpy as np
import openscm_units
from attrs import define

if TYPE_CHECKING:
    import numpy.typing as nptype
    import pint
    import scmdata.run


class XToNamedPintConvertor(Protocol):  # pylint: disable=too-few-public-methods
    """
    Callable that supports converting the x-array to Pint quantities
    """

    def __call__(
        self,
        x: nptype.NDArray[Union[np.float_, np.int_]],  # pylint: disable=invalid-name
    ) -> Dict[str, Union[pint.Quantity[np.float_], np.float_]]:
        """
        Convert x to pint quantities
        """


class ModelRunsInputGenerator(Protocol):  # pylint: disable=too-few-public-methods
    """
    Callable that supports generating model run inputs
    """

    def __call__(
        self, **kwargs: Union[pint.Quantity[np.float_], np.float_]
    ) -> Dict[str, Any]:
        """
        Generate model run inputs
        """


class ModelRunner(Protocol):  # pylint: disable=too-few-public-methods
    """
    Callable that supports running the model
    """

    def __call__(self, **kwargs: Any) -> scmdata.run.BaseScmRun:
        """
        Run the model
        """


@define
class OptModelRunner:  # pylint: disable=too-few-public-methods
    """
    Model runner during optimisation
    """

    convert_x_to_names_with_units: XToNamedPintConvertor
    """
    Callable used to translate the x-array into input for ``do_model_runs_input_generator``

    This translates from the x-array used internally by e.g. scipy and
    emcee into a dictionary with meaningful keys and quantities with units (as
    needed). It must produce named output that can be passed directly to
    ``self.do_model_runs_input_generator``.
    """

    do_model_runs_input_generator: ModelRunsInputGenerator
    """
    Generator of inputs for ``do_model_runs``

    More specifically, the callable used to translate the parameters (already
    converted to :obj:`pint.Quantity`) into the keyword arguments required by
    ``self.do_model_runs``
    """

    do_model_runs: ModelRunner
    """
    Function that runs the model

    Runs the desired experiments based on inputs generated by
    ``self.do_model_runs_input_generator``
    """

    @classmethod
    def from_parameters(
        cls,
        params: Iterable[Tuple[str, Union[str, pint.Unit, None]]],
        do_model_runs_input_generator: ModelRunsInputGenerator,
        do_model_runs: ModelRunner,
    ) -> OptModelRunner:
        """
        Initialise from list of parameters

        This is a convenience method

        Parameters
        ----------
        params
            List of parameters

        do_model_runs_input_generator
            Generator of input for ``do_model_runs``. See docstring of
            ``self`` for more details

        do_model_runs
            Callable which does the model runs. See docstring of  ``self`` for
            more details

        Returns
        -------
            Initialised :obj:`OptModelRunner`
        """
        convert_x_to_names_with_units = partial(
            x_and_parameters_to_named_with_units,
            params=params,
        )

        return OptModelRunner(
            convert_x_to_names_with_units=convert_x_to_names_with_units,
            do_model_runs_input_generator=do_model_runs_input_generator,
            do_model_runs=do_model_runs,
        )

    def run_model(
        self,
        x: nptype.NDArray[Union[np.float_, np.int_]],  # pylint: disable=invalid-name
    ) -> scmdata.run.BaseScmRun:
        """
        Run the model

        Parameters
        ----------
        x
            Vector of calibration parameter values

        Returns
        -------
            Results of run
        """
        x_converted_name = self.convert_x_to_names_with_units(x)

        do_model_runs_inputs = self.do_model_runs_input_generator(**x_converted_name)

        res = self.do_model_runs(**do_model_runs_inputs)

        return res


def x_and_parameters_to_named_with_units(
    x: nptype.NDArray[Union[np.float_, np.int_]],  # pylint: disable=invalid-name
    params: Iterable[Tuple[str, Union[str, pint.Unit, None]]],
    get_unit_registry: Optional[Callable[[], pint.UnitRegistry]] = None,
) -> Dict[str, Union[pint.Quantity[np.float_], np.float_]]:
    """
    Convert x array and parameters to a dictionary and add units

    Parameters
    ----------
    x
        Vector of calibration parameter values

    params
        params to be calibrated

        This defines both the names and units of the parameters. If the units
        are ``None``, the values are assumed to be plain numpy quantities.

    get_unit_registry
        Function to get unit registry. This allows the user to do a delayed
        import of the unit registry which is important because pint's unit
        registries don't parallelise well.

        If not provided, ``openscm_units.unit_registry`` is used

    Returns
    -------
        Parameters, named and converted to :obj:`pint.Quantity` where
        appropriate

    Examples
    --------
    It also possible to inject a different registry as needed
    >>> import pint
    >>> ur_plus_pop = pint.UnitRegistry()
    >>> ur_plus_pop.define("thousands = [population]")
    >>> def get_ur_with_pop():
    ...     return ur_plus_pop
    >>> # Withoout the injection, an error is raised
    >>> x_and_parameters_to_named_with_units(
    ...     [1.1, 3.2],
    ...     [("para_a", "m"), ("pop_weight", "thousands")],
    ... )
    Traceback (most recent call last):
    ...
    pint.errors.UndefinedUnitError: 'thousands' is not defined in the unit registry
    >>> # With the injection, this works nicely
    >>> x_and_parameters_to_named_with_units(
    ...     [1.1, 3.2, 4.0],
    ...     [("para_a", "m"), ("pop_weight", "thousands"), ("factor", None)],
    ...     get_ur_with_pop,
    ... )
    {'para_a': <Quantity(1.1, 'meter')>, 'pop_weight': <Quantity(3.2, 'thousands')>, 'factor': 4.0}
    """
    unit_reg = get_unit_registry() if get_unit_registry else openscm_units.unit_registry

    out = {}
    for val, (key, unit) in zip(x, params):
        if unit is not None:
            val_out = unit_reg.Quantity(val, unit)
        else:
            val_out = val

        out[key] = val_out

    return out
