import copy

import numpy as np
from typing import Callable, Iterable


class SandCat:
    """
    Агент песчаного кота для базового SCSO.
    """

    def __init__(self, x, function: Callable):
        self.x = np.array(x, dtype=float)
        self.function = function
        self.f = function(self.x)

    def update_cat(
            self,
            best_cat,
            r_g: float,
            ub: np.ndarray,
            lb: np.ndarray
    ):
        dim = self.x.size

        R = (2 * r_g) * (2 * np.random.rand(dim) - 1)
        abs_R = np.abs(R)

        S_r = r_g * np.random.rand(dim)

        theta = 2 * np.pi * np.random.rand(dim)

        new_x = self.x.copy()

        if np.any(abs_R > 1):
            rand_point = lb + np.random.rand(dim) * (ub - lb)
            rand_scale = np.random.rand(dim)
            new_x = S_r * (rand_point - rand_scale * self.x)
        else:
            rand_scale = np.random.rand(dim)
            direction = best_cat.x - self.x
            new_x = best_cat.x - S_r * rand_scale * direction * np.cos(theta)

        new_x = np.clip(new_x, lb, ub)
        self.x = new_x
        self.f = self.function(self.x)

    def __lt__(self, other):
        return self.f < other.f

    def __le__(self, other):
        return self.f <= other.f

    def __gt__(self, other):
        return self.f > other.f

    def __ge__(self, other):
        return self.f >= other.f


class SandCatOptimizer:

    def __init__(
            self,
            f: Callable,
            x_min: Iterable[float],
            x_max: Iterable[float],
            n_cats: int = 30,
            max_iter: int = 500,
            r_g_max: float = 2.0,
            r_g_min: float = 0.0,
            printing: bool = False
    ):
        self.objective_func = f
        self.n_cats = n_cats
        self.max_iter = max_iter

        self.lb = np.array(x_min, dtype=float)
        self.ub = np.array(x_max, dtype=float)
        self.r_g_max = r_g_max
        self.r_g_min = r_g_min

        self.printing = printing

    def _first_iteration(self):
        self.t = 0
        self.cats = []
        for _ in range(self.n_cats):
            xi = np.random.uniform(self.lb, self.ub)
            self.cats.append(SandCat(xi, self.objective_func))
        self.cats.sort()
        self.best_cat = copy.copy(self.cats[0])

    def _compute_r_g(self):
        progress = self.t / max(1, self.max_iter)
        return self.r_g_max - progress * (self.r_g_max - self.r_g_min)

    def _iteration(self):
        self.cats.sort()
        if self.cats[0].f < self.best_cat.f:
            self.best_cat = copy.copy(self.cats[0])

        r_g = self._compute_r_g()

        for cat in self.cats:
            cat.update_cat(
                best_cat=self.best_cat,
                r_g=r_g,
                ub=self.ub,
                lb=self.lb
            )

        if self.printing:
            print(self.t, self.best_cat.x, self.best_cat.f)

        self.t += 1

    def start(self):
        self._first_iteration()
        while self.t <= self.max_iter:
            self._iteration()
        return self.best_cat.x, self.best_cat.f, False, self.t

    def testing(self, canonical_x, epsilon):
        self._first_iteration()
        while self.t <= self.max_iter:
            self._iteration()
            if np.linalg.norm(canonical_x - self.best_cat.x) < epsilon:
                return self.best_cat.x, self.best_cat.f, self.t
        return self.best_cat.x, self.best_cat.f, self.t
