import numpy as np
from sklearn.neighbors import NearestNeighbors


class Archive:
    """Implements an archive for storing solutions and their informations"""
    def __init__(self, capacity, k):
        self.capacity = capacity
        self.container = {"params": [], "behaviors": [], "fitnesses": [], "gen": [], "from_novelty": []}
        self.position = 0
        self.k = k
        self.neigh = NearestNeighbors(n_neighbors=k)
        self.init = 0

        self.gen = 1

    def add(self, param, behavior, fitness, from_novelty):
        """Add a solution, its associated behavior descriptor and fitness to the archive"""
        # update the container with (params, results in behavior space, and fitness)
        if len(self.container["params"]) == self.capacity:
            self.container["params"][self.position] = param
            self.container["behaviors"][self.position] = behavior
            self.container["fitnesses"][self.position] = fitness
            self.container["gen"][self.position] = self.gen
            self.container["from_novelty"][self.position] = from_novelty
        else:
            self.container["params"].append(param)
            self.container["behaviors"].append(behavior)
            self.container["fitnesses"].append(fitness)
            self.container["gen"].append(self.gen)
            self.container["from_novelty"].append(from_novelty)
        self.position = int((self.position + 1) % self.capacity)

    def update_nearest_neighbors(self):
        """Update the knn model used to compute distances"""
        self.init = 1
        behaviors = np.array(self.container["behaviors"])
        if behaviors.ndim == 1:
            behaviors = behaviors.reshape(-1, 1)
        self.neigh.fit(behaviors)
        self.gen += 1

    def _replace_in_container(self, ind, param, behavior, fitness, from_novelty):
        """Replace a solution in the archive at a given index"""
        self.container["params"][ind] = param
        self.container["behaviors"][ind] = behavior
        self.container["fitnesses"][ind] = fitness
        self.container["gen"][ind] = self.gen
        self.container["from_novelty"][ind] = from_novelty
