"""Function definitions with their gradients for optimization."""

import math
from abc import ABC, abstractmethod
from typing import Tuple

import numpy as np


class Function1D(ABC):
    """Abstract base class for 1D functions."""

    name: str = "Abstract 1D Function"

    @abstractmethod
    def __call__(self, x: float) -> float:
        """Evaluate function at x."""
        pass

    @abstractmethod
    def gradient(self, x: float) -> float:
        """Compute gradient (derivative) at x."""
        pass

    @property
    @abstractmethod
    def domain(self) -> Tuple[float, float]:
        """Return the domain [a, b] for this function."""
        pass


class Function2D(ABC):
    """Abstract base class for 2D functions."""

    name: str = "Abstract 2D Function"

    @abstractmethod
    def __call__(self, x: np.ndarray) -> float:
        """Evaluate function at point x = [x1, x2]."""
        pass

    @abstractmethod
    def gradient(self, x: np.ndarray) -> np.ndarray:
        """Compute gradient at point x = [x1, x2]."""
        pass

    @property
    @abstractmethod
    def plot_bounds(self) -> Tuple[Tuple[float, float], Tuple[float, float]]:
        """Return bounds ((x1_min, x1_max), (x2_min, x2_max)) for plotting."""
        pass


class TaskFunction1D(Function1D):
    """
    f(x) = sqrt(x^2 + 9) / 4 + (5 - x) / 5

    Derivative: f'(x) = x / (4 * sqrt(x^2 + 9)) - 1/5
    """

    name = "f(x) = √(x² + 9)/4 + (5 - x)/5"

    def __call__(self, x: float) -> float:
        return math.sqrt(x**2 + 9) / 4 + (5 - x) / 5

    def gradient(self, x: float) -> float:
        return x / (4 * math.sqrt(x**2 + 9)) - 1 / 5

    @property
    def domain(self) -> Tuple[float, float]:
        return (-3.0, 8.0)


class HimmelblauFunction(Function2D):
    """
    Himmelblau's function:
    f(x, y) = (x^2 + y - 11)^2 + (x + y^2 - 7)^2

    Has 4 identical local minima at:
    - (3.0, 2.0)
    - (-2.805118, 3.131312)
    - (-3.779310, -3.283186)
    - (3.584428, -1.848126)

    Gradient:
    ∂f/∂x = 4x(x² + y - 11) + 2(x + y² - 7)
    ∂f/∂y = 2(x² + y - 11) + 4y(x + y² - 7)
    """

    name = "Himmelblau: (x² + y - 11)² + (x + y² - 7)²"

    def __call__(self, x: np.ndarray) -> float:
        x1, x2 = x[0], x[1]
        return (x1**2 + x2 - 11) ** 2 + (x1 + x2**2 - 7) ** 2

    def gradient(self, x: np.ndarray) -> np.ndarray:
        x1, x2 = x[0], x[1]
        df_dx1 = 4 * x1 * (x1**2 + x2 - 11) + 2 * (x1 + x2**2 - 7)
        df_dx2 = 2 * (x1**2 + x2 - 11) + 4 * x2 * (x1 + x2**2 - 7)
        return np.array([df_dx1, df_dx2])

    @property
    def plot_bounds(self) -> Tuple[Tuple[float, float], Tuple[float, float]]:
        return ((-5.0, 5.0), (-5.0, 5.0))


class RavineFunction(Function2D):
    """
    Овражная функция (эллиптический параболоид):
    f(x, y) = x² + 20y²

    Минимум в (0, 0), f(0,0) = 0

    Демонстрирует "эффект оврага" - градиент почти перпендикулярен
    направлению к минимуму, что замедляет сходимость.

    Gradient:
    ∂f/∂x = 2x
    ∂f/∂y = 40y
    """

    name = "Овраг: f(x,y) = x² + 20y²"

    def __call__(self, x: np.ndarray) -> float:
        x1, x2 = x[0], x[1]
        return x1**2 + 20 * x2**2

    def gradient(self, x: np.ndarray) -> np.ndarray:
        x1, x2 = x[0], x[1]
        df_dx1 = 2 * x1
        df_dx2 = 40 * x2
        return np.array([df_dx1, df_dx2])

    @property
    def plot_bounds(self) -> Tuple[Tuple[float, float], Tuple[float, float]]:
        return ((-2.0, 2.0), (-0.5, 0.5))


# Registry of available functions
FUNCTIONS_1D = {
    "task": TaskFunction1D,
}

FUNCTIONS_2D = {
    "himmelblau": HimmelblauFunction,
    "ravine": RavineFunction,
}