import numpy as np
import math
import random
import time


def euclidean_distance(pt1, pt2):
    return math.sqrt((pt1[0] - pt2[0]) ** 2 + (pt1[1] - pt2[1]) ** 2)


class ContinuousFunction:
    """
    This is a class that represents the function we are trying to optimize.
    eval_function is a lambda function representing the function we're
       optimizing, and bounds represents the min and max allowed value in
       each dimension (so, this assumes a rectangular boundary)
    """

    def __init__(self, eval_function, bounds):
        self.eval_function = eval_function
        self.bounds = bounds

    """
    Return a random point within the bounds
    """
    def random(self):
        pt = []
        for (_min, _max) in self.bounds:
            diff = _max - _min
            pt.append(random.random() * diff + _min)
        return pt

    """
    Check if a point is within the bounds
    """
    def in_bounds(self, point):
        return all(p >= bnd[0] and p <= bnd[1] for (p, bnd) in zip(point, self.bounds))

    """
    Plug the point into the function to find its value at this point
    """
    def score(self, point):
        return self.eval_function(*point)

# It is pretty silly to have TWO classes for this, but for this demo I thought it
#   made the code easier to read and the theory easier to understand.
class Nest:
    """
    A nest holds an egg. That's it!
    """
    def __init__(self, egg):
        self.egg = egg

class Egg:
    """
    An egg has a vector for it's position. That's it!
    """
    def __init__(self, position):
        self.position = np.array(position)


class CuckooManager:
    """
    This is a class that handles the main routine, including managing
    all of the nests.
    """
    def __init__(
        self,
        space,
        dimension,
        N,
        levy_parameter,
        alpha,
        percent_to_fly,
        maximization=True,
    ):
        """
        space = object representing the problem in this case a ContinuousFunction objec
        dimension = dimension of the problem
        N = number of particles
        levy_parameter = which levy flight parameter to use -> bigger = less frequent large jumps
        alpha = a constant multiplied by the levy flight to account for scaling
        maximization = true if maximizing, false if minimizing
        """
        self.space = space
        self.dimension = dimension
        self.N = N
        self.maximization = maximization

        self.levy_parameter = levy_parameter
        self.alpha = alpha
        self.percent_to_fly = percent_to_fly

        self.nests = [self.random_nest() for _ in range(N)]

        self.global_best_sol = None
        self.global_best_score = None
        self.set_global_best()

    def is_better(self, new_score, old_score):
        """
        return if new_score is better than old_score (depends if we're
        maximizing or minimizing)
        """
        if self.maximization:
            return new_score > old_score
        return new_score < old_score

    def set_global_best(self):
        """
        Loops over each firefly and computes its score. If it's the best
        we've ever seen, remember it.
        """
        for nest in self.nests:
            score = self.space.score(nest.egg.position)
            if self.global_best_score is None or self.is_better(
                score, self.global_best_score
            ):
                self.global_best_score = score
                self.global_best_sol = nest.egg.position.copy()

    def random_nest(self):
        return Nest(Egg(self.space.random()))

    def levy(self):
        """
        Compute a levy jump
        """
        return np.array(
            [
                self.alpha * random.choice([-1, 1]) * (random.random() ** (-1 / self.levy_parameter) - 1)
                for _ in range(self.dimension)
            ]
        )

    def advance_one_nest(self):
        """
        pick a random nest, move its egg with a levy flight to get sol S
        pick a new random nest, if S is better than the egg in that nest, replace
          that egg with S
        """
        nest1, nest2 = random.sample(self.nests, 2)
        
        new_egg = Egg(nest1.egg.position + self.levy())

        if self.is_better(
            self.space.score(new_egg.position),
            self.space.score(nest2.egg.position),
        ) and self.space.in_bounds(new_egg.position):
            nest2.egg = new_egg

    def drop_lowest_percentage(self):
        """
        take the [self.percent_to_fly] lowest scoring eggs, and for each
        of them, do a levy jump and make that their new position as long
        as it's in bounds
        """
        self.nests.sort(key=lambda nest: self.space.score(nest.egg.position))
        num_to_drop = math.floor(self.percent_to_fly * len(self.nests))

        for nest in self.nests[:num_to_drop]:
            new_pos = nest.egg.position + self.levy()
            if self.space.in_bounds(new_pos):
                nest.egg = Egg(new_pos)


eval_function = lambda x,y: -(math.sin(x)*math.sin(x**2/math.pi)**20 + math.sin(y)*math.sin(2*y**2/math.pi)**20)
bounds = [[0, 4]]*2
cns_func = ContinuousFunction(eval_function, bounds)

CM = CuckooManager(cns_func, 2, 10, 2, alpha=0.1, percent_to_fly=0.25, maximization=True)
gen = 0
last_best = 0

while True:
    gen += 1
    current_scores = [cns_func.score(nest.egg.position) for nest in CM.nests]
    CM.set_global_best()

    if CM.global_best_score != last_best:
        last_best = CM.global_best_score
        print(f"Gen {gen}: best score = {CM.global_best_score} at point {CM.global_best_sol}")

    CM.advance_one_nest()
    CM.drop_lowest_percentage()

# plt.show(block=True)