from collections.abc import Sequence, Callable, Iterable
from random import random as builtin_random
from typing import Protocol, Self, cast

from altbacken.core.annealing import RandomNumberGenerator, AnnealingState
from altbacken.core.neighbourhood import no_temperature_scaling, TemperatureScaling
from altbacken.internal.neighbourhood.adaptive import Epsilon, ConstantEpsilon


class VectorNeighbourhood:
    """
    Defines a class for generating a neighborhood of vectors.

    The VectorNeighbourhood class generates new vectors by applying a small
    random perturbation to each component of an input vector.

    Attributes:
        epsilon (float): Magnitude of the perturbation applied to the input
            vector during neighborhood generation.
        random (RandomNumberGenerator): Random number generator used for
            creating perturbations.
        vector_type (Callable[[Iterable[float]], Sequence[float]]): Function
            used to define the type or structure of the output vector.
    """
    def __init__(
            self,
            epsilon: float | Epsilon[Sequence[float], float] = 0.1,
            random: RandomNumberGenerator = builtin_random,
            vector_type: Callable[[Iterable[float]], Sequence[float]] = tuple,
            temperature_scaling: TemperatureScaling = no_temperature_scaling
    ):
        self._epsilon: Epsilon[Sequence[float], float] = cast(
            Epsilon[Sequence[float], float],
            ConstantEpsilon(epsilon) if isinstance(epsilon, (int, float)) else epsilon
        )
        self._random: RandomNumberGenerator = random
        self._vector_type: Callable[[Iterable[float]], Sequence[float]] = vector_type
        self._temperature_scaling: TemperatureScaling = temperature_scaling

    def __call__(self, vector: AnnealingState[Sequence[float]]) -> Sequence[float]:
        return self._vector_type(
            (2.0*self._random()-1) * self._epsilon(vector) * self._temperature_scaling(vector.temperature) + x
            for x in vector.current_solution
        )


class IntegerVectorNeighbourhood:
    """
    Represents a neighborhood function for integer vectors, used in simulated annealing
    to generate neighboring solutions.

    This class is designed to provide a way to generate new candidate solutions by
    perturbing an integer vector. The perturbation is controlled by parameters like
    epsilon, randomness, and temperature scaling.

    Attributes:
        epsilon (int | Epsilon[Sequence[int], int]): The magnitude of the
            perturbations. If an integer is provided, a constant epsilon is used.
            Otherwise, a custom Epsilon instance can be specified.
        random (RandomNumberGenerator): A random number generator used for
            creating perturbations.
        vector_type (Callable[[Iterable[int]], Sequence[int]]): A callable that
            determines the type of vector to be returned, such as tuple or list.
        temperature_scaling (TemperatureScaling): A function or callable used
            to scale perturbations based on the current temperature during the
            annealing process.
    """
    def __init__(
            self,
            epsilon: int | Epsilon[Sequence[int], int] = 1,
            random: RandomNumberGenerator = builtin_random,
            vector_type: Callable[[Iterable[int]], Sequence[int]] = tuple,
            temperature_scaling: TemperatureScaling = no_temperature_scaling
    ):
        self._epsilon: Epsilon[Sequence[int], int] = cast(
            Epsilon[Sequence[int], int],
            ConstantEpsilon(abs(epsilon)) if isinstance(epsilon, (int, float)) else epsilon
        )
        self._random: RandomNumberGenerator = random
        self._vector_type: Callable[[Iterable[int]], Sequence[int]] = vector_type
        self._temperature_scaling: TemperatureScaling = temperature_scaling

    def __call__(self, vector: AnnealingState[Sequence[int]]) -> Sequence[int]:
        return self._vector_type(
            (round(2.0*self._random()-1*self._epsilon(vector) * self._temperature_scaling(vector.temperature)) + x)
            for x in vector.current_solution
        )


class SupportsArrayOps(Protocol):
    shape: tuple[int, ...]
    def __add__(self, other) -> Self: ...


class ArrayNeighbourhood[A: SupportsArrayOps]:
    """
    Represents a neighborhood for array-based solutions in optimization algorithms.

    This class is designed to generate neighboring solutions in an array-based
    optimization context, such as simulated annealing. It utilizes a specified
    epsilon value for scaling the neighborhood bounds and applies temperature
    scaling where applicable. The class is parameterized to work with data types
    that support array operations.

    Attributes:
        epsilon (float | Epsilon[A, float]): Scaling factor for determining the
            range of the neighborhood bounds. Can either be a constant float or
            a custom epsilon function.
        seed (int): Seed for the random number generator to ensure reproducibility
            of the generated neighbors.
        temperature_scaling (TemperatureScaling): A function or object used for
            scaling the neighborhood size based on the current temperature.
    """
    def __init__(
        self,
        epsilon: float | Epsilon[A, float] = 0.1,
        seed: int = 0,
        temperature_scaling: TemperatureScaling = no_temperature_scaling
    ):
        self._epsilon: Epsilon[A, float] = cast(
            Epsilon[A, float],
            ConstantEpsilon(abs(epsilon)) if isinstance(epsilon, (float, int)) else epsilon
        )
        from altbacken.optional.numpy.random import default_rng
        self._rng = default_rng(seed)
        self._temperature_scaling = temperature_scaling

    def __call__(self, state: AnnealingState[A]) -> A:
        array: A = state.current_solution
        bounds = self._temperature_scaling(state.temperature) * self._epsilon(state)
        return array + self._rng.uniform(low=-bounds, high=bounds, size=array.shape)


