import copy
import numpy as np
from typing import Callable, Iterable


class Wolf:

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

    def update_wolf(self, a: float, alpha, beta, delta, ub, lb):
        r1 = np.random.random(self.x.size)
        r2 = np.random.random(self.x.size)

        A1 = 2 * a * r1 - a
        C1 = 2 * r2
        D_alpha = np.abs(C1 * alpha.x - self.x)
        X1 = alpha.x - A1 * D_alpha

        r1 = np.random.random(self.x.size)
        r2 = np.random.random(self.x.size)
        A2 = 2 * a * r1 - a
        C2 = 2 * r2
        D_beta = np.abs(C2 * beta.x - self.x)
        X2 = beta.x - A2 * D_beta

        r1 = np.random.random(self.x.size)
        r2 = np.random.random(self.x.size)
        A3 = 2 * a * r1 - a
        C3 = 2 * r2
        D_delta = np.abs(C3 * delta.x - self.x)
        X3 = delta.x - A3 * D_delta
        self.x = (X1 + X2 + X3) / 3
        self.x = np.clip(self.x, lb, ub)
        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 GreyWolfOptimizer:

    def __init__(
            self,
            f: Callable,
            x_min: Iterable[float],
            x_max: Iterable[float],
            n_wolves: int = 30,
            max_iter: int = 500,
            printing=False
    ):
        self.objective_func = f
        self.n_wolves = n_wolves
        self.max_iter = max_iter

        self.lb = np.array(x_min)
        self.ub = np.array(x_max)

        self.printing = printing

    def _first_iteration(self):
        self.t = 0
        self.wolves = []
        for i in range(self.n_wolves):
            xi = np.random.uniform(self.ub, self.lb)
            self.wolves.append(Wolf(xi, self.objective_func))
        self.wolves.sort()

    def _iteration(self):
        alpha, beta, delta = tuple(map(copy.copy, self.wolves[:3]))
        a = 2 - self.t * (2 / self.max_iter)
        for wolf in self.wolves:
            wolf.update_wolf(a, alpha, beta, delta, self.ub, self.lb)
        self.wolves.sort()
        if self.printing:
            print(self.t, self.wolves[0].x, self.wolves[0].f)
        self.t += 1

    def start(self):
        self._first_iteration()
        while self.t <= self.max_iter:
            stop = self._iteration()
            if stop:
                break
        return (self.wolves[0].x, self.wolves[0].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 - np.array(self.wolves[0].x)) < epsilon:
                return self.wolves[0].x, self.wolves[0].f, self.t
        return self.wolves[0].x, self.wolves[0].f, self.t
