"""An implementation of an box-constrained n-dimensional continuous space."""

from math import isfinite
from typing import Callable, Final, Iterable, cast

import numpy as np
from numpy import clip
from pycommons.types import type_error

from moptipy.spaces.nparrayspace import NPArraySpace
from moptipy.utils.logger import KeyValueLogSection
from moptipy.utils.nputils import DEFAULT_FLOAT, array_to_str, is_np_float

#: the log key for the lower bound, i.e., the minimum permitted value
KEY_LOWER_BOUND: Final[str] = "lb"
#: the log key for the upper bound, i.e., the maximum permitted value
KEY_UPPER_BOUND: Final[str] = "ub"


class VectorSpace(NPArraySpace):
    """
    A vector space where each element is an n-dimensional real vector.

    Such spaces are useful for continuous optimization. The vectors are
    implemented as one-dimensional :class:`numpy.ndarray` of length `n`.
    A vector space is constraint by a box which defines the minimum and
    maximum permitted value for each of its `n` elements.

    >>> s = VectorSpace(3)
    >>> print(s.dimension)
    3
    >>> print(s.dtype)
    float64
    >>> print(s.lower_bound)
    [0. 0. 0.]
    >>> print(s.upper_bound)
    [1. 1. 1.]
    >>> print(s.lower_bound_all_same)
    True
    >>> print(s.upper_bound_all_same)
    True
    >>> s = VectorSpace(2, -1.0, 5.0)
    >>> print(s.lower_bound)
    [-1. -1.]
    >>> print(s.upper_bound)
    [5. 5.]
    >>> s = VectorSpace(2, [-1.0, -2.0], 5.0)
    >>> print(s.lower_bound)
    [-1. -2.]
    >>> print(s.upper_bound)
    [5. 5.]
    >>> print(s.lower_bound_all_same)
    False
    >>> print(s.upper_bound_all_same)
    True
    """

    def __init__(self, dimension: int,
                 lower_bound: float | Iterable[float] = 0.0,
                 upper_bound: float | Iterable[float] = 1.0,
                 dtype: np.dtype = DEFAULT_FLOAT) -> None:
        """
        Create the vector-based search space.

        :param dimension: The dimension of the search space,
            i.e., the number of decision variables.
        :param dtype: The basic data type of the vector space,
            i.e., the type of the decision variables
        :param lower_bound: the optional minimum value(s)
        :param upper_bound: the optional maximum value(s)
        """
        super().__init__(dimension, dtype)
        if not is_np_float(dtype):
            raise TypeError(f"Invalid data type {dtype}.")

        # first, we process the lower bounds
        lower_bound_all_same: bool = True
        if isinstance(lower_bound, float):
            # only a single value is given
            if not isfinite(lower_bound):
                raise ValueError(
                    f"invalid lower bound {lower_bound}.")
            # if a single value is given, then we expand it a vector
            lower_bound = np.full(dimension, lower_bound, dtype)
        elif isinstance(lower_bound, Iterable):
            # lower bounds are given as vector or iterable
            lb = np.array(lower_bound, dtype)
            if len(lb) != dimension:
                raise ValueError(f"wrong length {lb} of lower "
                                 f"bound iterable {lower_bound}")
            if lb.shape != (dimension, ):
                raise ValueError(f"invalid shape={lb.shape} of "
                                 f"lower bound {lower_bound}")
            first = lb[0]
            for index, item in enumerate(lb):
                if first != item:
                    lower_bound_all_same = False
                if not np.isfinite(item):
                    raise ValueError(f"{index}th lower bound={item}")
            lower_bound = lb
        else:
            raise type_error(lower_bound, "lower_bound", (
                float, Iterable, None))

        # now, we process the upper bounds
        upper_bound_all_same: bool = True
        if isinstance(upper_bound, float):
            # only a single value is given
            if not isfinite(upper_bound):
                raise ValueError(
                    f"invalid upper bound {upper_bound}.")
            # if a single value is given, then we expand it a vector
            upper_bound = np.full(dimension, upper_bound, dtype)
        elif isinstance(upper_bound, Iterable):
            # upper bounds are given as vector or iterable
            lb = np.array(upper_bound, dtype)
            if len(lb) != dimension:
                raise ValueError(f"wrong length {lb} of upper "
                                 f"bound iterable {upper_bound}")
            if lb.shape != (dimension,):
                raise ValueError(f"invalid shape={lb.shape} of "
                                 f"upper bound {upper_bound}")
            first = lb[0]
            for index, item in enumerate(lb):
                if first != item:
                    upper_bound_all_same = False
                if not np.isfinite(item):
                    raise ValueError(f"{index}th upper bound={item}")
            upper_bound = lb
        else:
            raise type_error(upper_bound, "upper_bound", (
                float, Iterable, None))

        # check that the bounds are consistent
        for idx, ll in enumerate(lower_bound):
            if not ll < upper_bound[idx]:
                raise ValueError(f"lower_bound[{idx}]={ll} >= "
                                 f"upper_bound[{idx}]={upper_bound[idx]}")

        #: the lower bounds for all variables
        self.lower_bound: Final[np.ndarray] = lower_bound
        #: all dimensions have the same lower bound
        self.lower_bound_all_same: Final[bool] = lower_bound_all_same
        #: the upper bounds for all variables
        self.upper_bound: Final[np.ndarray] = upper_bound
        #: all dimensions have the same upper bound
        self.upper_bound_all_same: Final[bool] = upper_bound_all_same

    def clipped(self, func: Callable[[np.ndarray], int | float]) \
            -> Callable[[np.ndarray], int | float]:
        """
        Wrap a function ensuring that all vectors are clipped to the bounds.

        This function is useful to ensure that only valid vectors are passed
        to :meth:`~moptipy.api.process.Process.evaluate`.

        :param func: the function to wrap
        :returns: the wrapped function
        """
        return cast(Callable[[np.ndarray], int | float],
                    lambda x, lb=self.lower_bound, ub=self.upper_bound,
                    ff=func: ff(clip(x, lb, ub, x)))

    def validate(self, x: np.ndarray) -> None:
        """
        Validate a vector.

        :param x: the real vector
        :raises TypeError: if the string is not an :class:`numpy.ndarray`.
        :raises ValueError: if the shape of the vector is wrong or any of its
            element is not finite.
        """
        super().validate(x)

        mib: Final[np.ndarray] = self.lower_bound
        mab: Final[np.ndarray] = self.upper_bound
        for index, item in enumerate(x):
            miv = mib[index]
            mav = mab[index]
            if not np.isfinite(item):
                raise ValueError(f"x[{index}]={item}, which is not finite")
            if not (miv <= item <= mav):
                raise ValueError(
                    f"x[{index}]={item}, but should be in [{miv},{mav}].")

    def n_points(self) -> int:
        """
        Get an upper bound for the number of different values in this space.

        :return: We return the approximate number of finite floating point
          numbers while ignoring the box constraint. This value here therefore
          is an upper bound.

        >>> import numpy as npx
        >>> print(VectorSpace(3, dtype=npx.dtype(npx.float64)).n_points())
        6267911251143764491534102180507836301813760039183993274367
        """
        if self.dtype.char == "e":
            exponent = 5
            mantissa = 10
        elif self.dtype.char == "f":
            exponent = 8
            mantissa = 23
        elif self.dtype == "d":
            exponent = 11
            mantissa = 52
        elif self.dtype == "g":
            exponent = 15
            mantissa = 112
        else:
            raise ValueError(f"Invalid dtype {self.dtype}.")

        base = 2 * ((2 ** exponent) - 1) * (2 ** mantissa) - 1
        return base ** self.dimension

    def __str__(self) -> str:
        """
        Get the name of this space.

        :return: "r" + dimension + dtype.char

        >>> import numpy as npx
        >>> print(VectorSpace(3, dtype=npx.dtype(npx.float64)))
        r3d
        """
        return f"r{self.dimension}{self.dtype.char}"

    def log_bounds(self, logger: KeyValueLogSection) -> None:
        """
        Log the bounds of this space to the given logger.

        :param logger: the logger for the parameters

        >>> from moptipy.utils.logger import InMemoryLogger
        >>> import numpy as npx
        >>> space = VectorSpace(2, -5.0, [2.0, 3.0])
        >>> with InMemoryLogger() as l:
        ...     with l.key_values("C") as kv:
        ...         space.log_bounds(kv)
        ...     text = l.get_log()
        >>> text[-2]
        'ub: 2;3'
        >>> text[-3]
        'lb: -5'
        >>> len(text)
        4
        """
        if self.lower_bound_all_same:
            logger.key_value(
                KEY_LOWER_BOUND, self.lower_bound[0], also_hex=False)
        else:
            logger.key_value(KEY_LOWER_BOUND, array_to_str(self.lower_bound))
        if self.upper_bound_all_same:
            logger.key_value(
                KEY_UPPER_BOUND, self.upper_bound[0], also_hex=False)
        else:
            logger.key_value(KEY_UPPER_BOUND, array_to_str(self.upper_bound))

    def log_parameters_to(self, logger: KeyValueLogSection) -> None:
        """
        Log the parameters of this space to the given logger.

        :param logger: the logger for the parameters

        >>> from moptipy.utils.logger import InMemoryLogger
        >>> import numpy as npx
        >>> space = VectorSpace(2, -5.0, [2.0, 3.0])
        >>> space.dimension
        2
        >>> space.dtype.char
        'd'
        >>> with InMemoryLogger() as l:
        ...     with l.key_values("C") as kv:
        ...         space.log_parameters_to(kv)
        ...     text = l.get_log()
        >>> text[-2]
        'ub: 2;3'
        >>> text[-3]
        'lb: -5'
        >>> text[-4]
        'dtype: d'
        >>> text[-5]
        'nvars: 2'
        >>> text[-6]
        'class: moptipy.spaces.vectorspace.VectorSpace'
        >>> text[-7]
        'name: r2d'
        >>> len(text)
        8
        """
        super().log_parameters_to(logger)
        self.log_bounds(logger)
