from functools import cached_property
from itertools import chain
from numbers import Integral
from typing import Any, Dict, Generic, Iterable, Optional, Tuple, TypeVar, Union

import numpy as np
import numpy.typing as npt
from csnlp.core.cache import invalidate_cache
from csnlp.util.io import SupportsDeepcopyAndPickle

from mpcrl.util.math import summarize_array

SymType = TypeVar("SymType")  # most likely, T is cs.SX or MX


class LearnableParameter(SupportsDeepcopyAndPickle, Generic[SymType]):
    """A parameter that is learnable, that is, it can be adjusted via RL. This class
    is useful for managing symbols, bounds and value of learnable parameters."""

    __slots__ = ("name", "shape", "value", "lb", "ub", "sym")

    def __init__(
        self,
        name: str,
        shape: Union[int, Tuple[int, ...]],
        value: npt.ArrayLike,
        lb: npt.ArrayLike = -np.inf,
        ub: npt.ArrayLike = +np.inf,
        sym: Optional[SymType] = None,
    ) -> None:
        """Instantiates a learnable parameter.

        Parameters
        ----------
        name : str
            Name of the learnable parameter.
        shape : int or tuple of ints
            Shape of the parameter.
        value : array_like
            Starting value of the parameter. This can then be updated via `update`.
        lb : array_like, optional
            Lower bound of the parameter values. If not specified, it is unbounded.
        ub : array_like, optional
            Upper bound of the parameter values. If not specified, it is unbounded.
        sym : T, optional
            An optional reference to a symbolic variable representing this parameter.

        Raises
        ------
        ValueError
            Raises if `value`, `lb` or `ub` cannot be broadcasted to a 1D vector with
            shape equal to `shape`; or if the shape of the symbolic variable `sym` does
            not match the shape of the parameter.
        """
        super().__init__()
        self.name = name
        self.shape: Tuple[int, ...] = (
            (shape,) if isinstance(shape, Integral) else shape  # type: ignore
        )
        self.sym = sym
        self.lb: npt.NDArray[np.floating] = np.broadcast_to(lb, shape)
        self.ub: npt.NDArray[np.floating] = np.broadcast_to(ub, shape)
        self._check_sym_shape()
        self._update_value(value)

    @property
    def size(self) -> int:
        """Gets the number of elements in the parameter."""
        return np.prod(self.shape, dtype=int).item()

    def _update_value(self, v: npt.ArrayLike, **is_close_kwargs: Any) -> None:
        """Internal utility for updating the parameter value with a new value.

        Parameters
        ----------
        new_value : array_like
            New value of the parameter.
        is_close_kwargs
            Additional kwargs for `np.isclose`, e.g., `rtol` and `atol`, for checking
            numerical values close to a bound.

        Raises
        ------
        ValueError
            Raises if `new_value` cannot be broadcasted to a 1D vector with shape equal
            to `shape`; or if it does not lie inside the upper and lower bounds within
            the specified tolerances.
        """
        v = np.broadcast_to(v, self.shape)
        lb = self.lb
        ub = self.ub
        if ((v < lb) & ~np.isclose(v, lb, **is_close_kwargs)).any() or (
            (v > ub) & ~np.isclose(v, ub, **is_close_kwargs)
        ).any():
            raise ValueError(f"Updated parameter {self.name} is outside bounds.")
        self.value: npt.NDArray[np.floating] = np.clip(v, lb, ub)

    def _check_sym_shape(self) -> None:
        """Internal utility for checking that the shape of the symbolic variable matches
        the shape of the parameter."""
        if self.sym is None or not hasattr(self.sym, "shape"):
            return
        sym_shape = tuple(self.sym.shape)
        shape = self.shape + tuple(1 for _ in range(len(sym_shape) - len(self.shape)))
        if sym_shape != shape:
            raise ValueError("Shape of `sym` does not match `shape`.")

    def __str__(self) -> str:
        return f"<{self.name}(shape={self.shape})>"

    def __repr__(self) -> str:
        return f"<{self.__class__.__name__}(name={self.name},shape={self.shape})>"


class LearnableParametersDict(
    Dict[str, LearnableParameter[SymType]], SupportsDeepcopyAndPickle
):
    """Dict-based collection of `LearnableParameter` instances that simplifies the
    process of managing and updating these. The dict contains pairs of parameter's name
    vs parameter's instance.

    Note: to speed up computations, properties of this class are often cached for faster
    calls to the same methods. However, these are cleared when the underlying dict is
    modified."""

    def __init__(self, pars: Optional[Iterable[LearnableParameter[SymType]]] = None):
        """Initializes the collection of learnable parameters.

        Parameters
        ----------
        pars : iterable of LearnableParameter, optional
            An optional iterable of parameters to insert into the dict by their names.
        """
        if pars is None:
            dict.__init__(self)
        else:
            dict.__init__(self, map(lambda p: (p.name, p), chain(pars)))
        SupportsDeepcopyAndPickle.__init__(self)

    @cached_property
    def size(self) -> int:
        """Gets the overall size of all the learnable parameters."""
        return sum(p.size for p in self.values())

    @cached_property
    def lb(self) -> npt.NDArray[np.floating]:
        """Gets the lower bound of all the learnable parameters, concatenated."""
        return (
            np.concatenate([p.lb.reshape(-1, order="F") for p in self.values()])
            if self
            else np.empty(0)
        )

    @cached_property
    def ub(self) -> npt.NDArray[np.floating]:
        """Gets the upper bound of all the learnable parameters, concatenated."""
        return (
            np.concatenate([p.ub.reshape(-1, order="F") for p in self.values()])
            if self
            else np.empty(0)
        )

    @cached_property
    def value(self) -> npt.NDArray[np.floating]:
        """Gets the values of all the learnable parameters, concatenated."""
        return (
            np.concatenate([p.value.reshape(-1, order="F") for p in self.values()])
            if self
            else np.empty(0)
        )

    @cached_property
    def value_as_dict(self) -> Dict[str, npt.NDArray[np.floating]]:
        """Gets the values of all the learnable parameters, in a dict."""
        return {p.name: p.value for p in self.values()}

    @cached_property
    def sym(self) -> Dict[str, Optional[SymType]]:
        """Gets symbols of all the learnable parameters, in a dict. If one parameter
        does not possess the symbol, `None` is put."""
        return {
            parname: None if par.sym is None else par.sym
            for parname, par in self.items()
        }

    @invalidate_cache(value, value_as_dict)
    def update_values(
        self,
        new_values: Union[npt.ArrayLike, Dict[str, npt.ArrayLike]],
        **is_close_kwargs: Any,
    ) -> None:
        """Updates the value of each parameter

        Parameters
        ----------
        new_values : array_like or dict[str, array_like]
            The parameters' new values, either as a single concatenated array (which
            will be splitted according to the sizes and each piece sequentially assigned
            to each parameter), or as a dict of parameter's name vs parameter's new
            value.
        is_close_kwargs
            Additional kwargs for `np.isclose`, e.g., `rtol` and `atol`, for checking
            numerical values of parameters close to a bound.

        Raises
        ------
        ValueError
            In case of array-like, raises if `new_values` cannot be split according to
            the sizes of parameters; or if the new values cannot be broadcasted to 1D
            vectors according to each parameter's size; or if the new values lie outside
            either the lower or upper bounds of each parameter.
        """
        if isinstance(new_values, dict):
            for parname, new_value in new_values.items():
                self[parname]._update_value(new_value, **is_close_kwargs)
        else:
            cumsizes = np.cumsum([p.size for p in self.values()])[:-1]
            values_ = np.split(new_values, cumsizes)
            for par, val in zip(self.values(), values_):
                par._update_value(val.reshape(par.shape, order="F"), **is_close_kwargs)

    __cache_decorator = invalidate_cache(size, lb, ub, value, value_as_dict, sym)

    @__cache_decorator
    def __setitem__(self, name: str, par: LearnableParameter[SymType]) -> None:
        assert name == par.name, f"Key '{name}' must match parameter name '{par.name}'."
        return super().__setitem__(name, par)

    @__cache_decorator
    def update(
        self,
        pars: Iterable[LearnableParameter[SymType]],
        *args: LearnableParameter[SymType],
    ) -> None:
        return super().update(map(lambda p: (p.name, p), chain(pars, args)))

    @__cache_decorator
    def setdefault(
        self,
        par: LearnableParameter[SymType],
    ) -> LearnableParameter[SymType]:
        return super().setdefault(par.name, par)

    __delitem__ = __cache_decorator(dict.__delitem__)
    pop = __cache_decorator(dict.pop)
    popitem = __cache_decorator(dict.popitem)
    clear = __cache_decorator(dict.clear)

    def copy(
        self, deep: bool = False, invalidate_caches: bool = True
    ) -> "LearnableParametersDict[SymType]":
        """Creates a shallow or deep copy of the dict of learnable parameters.

        Parameters
        ----------
        deep : bool, optional
            If `True`, a deepcopy of the dict and its parameters is returned; otherwise,
            the copy is only shallow.
        invalidate_caches : bool, optional
            If `True`, methods decorated with `csnlp.util.funcs.invalidate_cache` are
            called to clear cached properties/lru caches in the copied instance.
            Otherwise, caches in the copy are not invalidated. By default, `True`.
            Only relevant when `deep=True`.

        Returns
        -------
        LearnableParametersDict[T]
            A copy of the dict of learnable parameters.
        """
        return (
            SupportsDeepcopyAndPickle.copy(self, invalidate_caches)
            if deep
            else LearnableParametersDict[SymType](self.values())
        )

    def stringify(
        self, summarize: bool = True, precision: int = 3, ddof: int = 0
    ) -> str:
        """Returns a string representing the dict of learnable parameters.

        Parameters
        ----------
        summarize : bool, optional
            If `True` (default), array parameters are summarized; otherwise, the entire
            array is printed.
        precision : int, optional
            The printing precision of floating point numbers.
        ddof : int, optional
            Degrees of freedom for computing standard deviations (see `numpy.std`).

        Returns
        -------
        str
            A string representing the dict and its parameters.
        """

        def p2s(p: LearnableParameter) -> str:
            if p.size == 1:
                return f"{p.name}={p.value.item():.{precision}f}"
            if summarize:
                return f"{p.name}: {summarize_array(p.value, precision, ddof)}"
            return np.array2string(p.value, precision=precision)

        return "; ".join(p2s(p) for p in self.values())

    def __str__(self) -> str:
        return self.stringify()

    def __repr__(self) -> str:
        return f"<{self.__class__.__name__}: {super().__repr__()}>"
