import copy

import numpy as np
from typing import Callable, Iterable


class Bee:
    def __init__(self, x, function: Callable):
        self.x = np.array(x)
        self.function = function
        self.f = function(self.x)

    def evaluate(self):
        self.f = self.function(self.x)

    def clone_with_perturbation(self, radius, lb, ub):
        x_new = self.x + np.random.uniform(-radius, radius, size=self.x.size)
        x_new = np.clip(x_new, lb, ub)
        return Bee(x_new, self.function)

    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 BeesOptimizer:

    def __init__(
            self,
            f: Callable,
            x_min: Iterable[float],
            x_max: Iterable[float],
            n_scouts: int = 30,
            n_selected: int = 10,
            n_elite: int = 3,
            n_recruited_elite: int = 10,
            n_recruited_selected: int = 5,
            max_iter: int = 500,
            radius: float = 0.1,
            radius_shrink: float = 0.95,
            printing: bool = False
    ):
        self.objective_func = f
        self.lb = np.array(x_min, dtype=float)
        self.ub = np.array(x_max, dtype=float)

        self.n_scouts = n_scouts
        self.n_selected = min(n_selected, n_scouts)
        self.n_elite = min(n_elite, self.n_selected)
        self.n_recruited_elite = n_recruited_elite
        self.n_recruited_selected = n_recruited_selected

        self.max_iter = max_iter
        self.radius = radius
        self.radius_shrink = radius_shrink
        self.printing = printing

    def _random_bee(self):
        x = np.random.uniform(self.lb, self.ub)
        return Bee(x, self.objective_func)

    def _first_iteration(self):
        self.t = 0
        self.bees = [self._random_bee() for _ in range(self.n_scouts)]
        self.bees.sort()
        self.best = copy.copy(self.bees[0])

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

        selected_sites = self.bees[:self.n_selected]
        elite_sites = selected_sites[:self.n_elite]
        other_selected = selected_sites[self.n_elite:]

        new_bees = []

        for site in elite_sites:
            local_bees = [site]
            for _ in range(self.n_recruited_elite - 1):
                local_bees.append(site.clone_with_perturbation(
                    self.radius, self.lb, self.ub
                ))
            local_bees.sort()
            new_bees.append(local_bees[0])

        for site in other_selected:
            local_bees = [site]
            for _ in range(self.n_recruited_selected - 1):
                local_bees.append(site.clone_with_perturbation(
                    self.radius, self.lb, self.ub
                ))
            local_bees.sort()
            new_bees.append(local_bees[0])

        n_new = self.n_scouts - len(new_bees)
        for _ in range(n_new):
            new_bees.append(self._random_bee())

        self.bees = new_bees
        if self.printing:
            print(self.t, self.best.x, self.best.f)

        self.radius *= self.radius_shrink

        self.t += 1

    def start(self):
        self._first_iteration()
        while self.t <= self.max_iter:
            self._iteration()
        return self.best.x, self.best.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.x) < epsilon:
                return self.best.x, self.best.f, self.t
        return self.best.x, self.best.f, self.t
