import random
import numpy as np


class Plant:

    def __init__(self, x, function):
        self.x = np.array(x)
        self.f = function(self.x)
        self.function = function
    
    def calculate_n(self, f_best, f_worst, n_min, n_max):
        return round(self.f * (n_max - n_min) / (f_best - f_worst) + (f_best * n_min - f_worst * n_max) / (f_best - f_worst))

    def create_new_plant(self, sigma):
        new_x = self.x + np.random.normal(0, sigma, self.x.size)
        return Plant(new_x, 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 IWO:

    def __init__(self, s0, s_max, n_min, n_max, sigma_b, sigma_e, x_min, x_max, m, t_max, f, printing=False):
        self.s0 = s0
        self.s_max = s_max
        self.n_min = n_min
        self.n_max = n_max
        self.sigma_b = sigma_b
        self.sigma_e = sigma_e
        self.x_min = x_min
        self.x_max = x_max
        self.m = m
        self.t_max = t_max
        self.f = f
        self.printing = printing

    def _first_iteration(self):
        self.t = 0
        self.plants = []
        for i in range(self.s0):
            xi = []
            for j in range(len(self.x_min)):
                xi.append(self.x_min[j] + random.random() * (self.x_max[j] - self.x_min[j]))
            self.plants.append(Plant(xi, self.f))
        self.plants.sort()

    def _iteration(self):

        n = []
        f_best = self.plants[0].f
        f_worst = self.plants[-1].f
        for plant in self.plants:
            n.append(plant.calculate_n(f_best, f_worst, self.n_min, self.n_max))
        new_plants = []
        sigma = (((self.t_max - self.t) / self.t_max) ** self.m) * (self.sigma_b - self.sigma_e) + self.sigma_e
        for i, plant in enumerate(self.plants):
            for j in range(n[i]):
                new_plants.append(plant.create_new_plant(sigma))
        self.plants += new_plants
        self.plants.sort()
        if len(self.plants) > self.s_max:
            self.plants = self.plants[:self.s_max]
        if self.printing:
            print(self.t, self.plants[0].x, self.plants[0].f)
        self.t += 1

    def start(self):
        self._first_iteration()
        while self.t <= self.t_max:
            stop = self._iteration()
            if stop:
                break
        return (self.plants[0].x, self.plants[0].f, False, self.t)

    def testing(self, canonical_x, epsilon):
        self._first_iteration()
        while self.t <= self.t_max:
            self._iteration()
            if np.linalg.norm(canonical_x - np.array(self.plants[0].x)) < epsilon:
                return self.plants[0].x, self.plants[0].f, self.t
        return self.plants[0].x, self.plants[0].f, self.t