from abc import ABC, abstractmethod
from typing import Callable

import numpy as np
from numpy.typing import NDArray

from .chromosome import Chromosome

type FitnessFn = Callable[
    [
        Chromosome,
        NDArray[np.float64],
        Callable[[NDArray[np.float64]], NDArray[np.float64]],
    ],
    float,
]
type TargetFunction = Callable[[NDArray[np.float64]], NDArray[np.float64]]
type TestPointsFn = Callable[[], NDArray[np.float64]]


class BaseFitness(ABC):
    def __init__(self, target_fn: TargetFunction, test_points_fn: TestPointsFn):
        self.target_function = target_fn
        self.test_points_fn = test_points_fn

    @abstractmethod
    def fitness_fn(
        self,
        chromosome: Chromosome,
        predicted: NDArray[np.float64],
        true_values: NDArray[np.float64],
    ) -> float: ...

    def __call__(self, chromosome: Chromosome) -> float:
        test_points = self.test_points_fn()

        context = {t: test_points[:, i] for i, t in enumerate(chromosome.terminals)}
        predicted = chromosome.root.eval(context)
        true_values = self.target_function(test_points)

        return self.fitness_fn(chromosome, predicted, true_values)


class MSEFitness(BaseFitness):
    """Среднеквадратичная ошибка"""

    def fitness_fn(
        self,
        chromosome: Chromosome,
        predicted: NDArray[np.float64],
        true_values: NDArray[np.float64],
    ) -> float:
        return float(np.mean((predicted - true_values) ** 2))


class RMSEFitness(BaseFitness):
    """Корень из среднеквадратичной ошибки"""

    def fitness_fn(
        self,
        chromosome: Chromosome,
        predicted: NDArray[np.float64],
        true_values: NDArray[np.float64],
    ) -> float:
        return float(np.sqrt(np.mean((predicted - true_values) ** 2)))


class MAEFitness(BaseFitness):
    """Средняя абсолютная ошибка"""

    def fitness_fn(
        self,
        chromosome: Chromosome,
        predicted: NDArray[np.float64],
        true_values: NDArray[np.float64],
    ) -> float:
        return float(np.mean(np.abs(predicted - true_values)))


class HuberFitness(BaseFitness):
    """Huber Loss (компромисс между MSE и MAE)"""

    def __init__(
        self,
        target_fn: TargetFunction,
        test_points_fn: TestPointsFn,
        delta: float = 1.0,
    ):
        super().__init__(target_fn, test_points_fn)
        self.delta = delta

    def fitness_fn(
        self,
        chromosome: Chromosome,
        predicted: NDArray[np.float64],
        true_values: NDArray[np.float64],
    ) -> float:
        error = predicted - true_values
        mask = np.abs(error) <= self.delta
        squared = 0.5 * (error[mask] ** 2)
        linear = self.delta * (np.abs(error[~mask]) - 0.5 * self.delta)
        huber = np.concatenate([squared, linear])
        return float(np.mean(huber))


class NRMSEFitness(BaseFitness):
    """Нормализованный RMSE (масштаб-инвариантен)"""

    def fitness_fn(
        self,
        chromosome: Chromosome,
        predicted: NDArray[np.float64],
        true_values: NDArray[np.float64],
    ) -> float:
        denom = np.std(true_values)
        if denom == 0:
            return 1e6
        return float(np.sqrt(np.mean((predicted - true_values) ** 2)) / denom)


class PenalizedFitness(BaseFitness):
    """Фитнес со штрафом за размер и глубину дерева: ошибка + λ * (размер + depth_weight * глубина)"""

    def __init__(
        self,
        target_fn: TargetFunction,
        test_points_fn: TestPointsFn,
        base_fitness: BaseFitness,
        lambda_: float = 0.001,
        depth_weight: float = 0.2,
    ):
        super().__init__(target_fn, test_points_fn)
        self.base_fitness = base_fitness
        self.lambda_ = lambda_
        self.depth_weight = depth_weight

    def fitness_fn(
        self,
        chromosome: Chromosome,
        predicted: NDArray[np.float64],
        true_values: NDArray[np.float64],
    ) -> float:
        base = self.base_fitness.fitness_fn(chromosome, predicted, true_values)

        size = chromosome.root.get_size()
        depth = chromosome.root.get_depth()

        penalty = self.lambda_ * (size + self.depth_weight * depth)
        return float(base + penalty)