import warnings
import numpy as np
import networkx as nx
from itertools import islice, combinations
from collections import deque
from abc import ABC, abstractmethod
from collections.abc import Iterable
from sklearn.cluster import SpectralClustering, AgglomerativeClustering

class Metrics:
    
    @staticmethod
    def hamming_distance(ind1, ind2):
        return sum(b1 != b2 for b1, b2 in zip(ind1, ind2))
    
    @staticmethod
    def np_hamming_distance(ind1, ind2):
        result = np.sum(ind1 != ind2)
        return result

class Archive(ABC):  

    def __init__(self, archive:Iterable, metric) -> None:
        super().__init__()
        self.archive = archive
        self.metric = metric
    
    @abstractmethod
    def novelty(self, ind):
        pass

    def get_graph(self, extract_data=None):
        graph = nx.Graph()

        # Add nodes
        for i, ind in enumerate(self.archive):
            graph.add_node(i, data = ind)

        for node1, node2 in combinations(graph.nodes(), 2):
            distance = self.metric(graph.nodes[node1]['data'], graph.nodes[node2]['data'])
            graph.add_edge(node1, node2, weight=distance)
        
        if extract_data is not None:
            for node in graph.nodes:
                data = graph.nodes[node]['data']
                graph.nodes[node]['data'] = extract_data(data)

        return graph
    
    def num_archived_items(self):
        return len(self.archive)

class SimpleArchive(Archive):
    
    def __init__(self, size, metric, k = 5):
        super().__init__([], metric)
        self.size = size
        self.k = k
        
    def add(self, population):
        self.archive += population
        overflow = len(self.archive)-self.size
        if(overflow>0):
            self.archive=self.archive[-overflow:]
    
    def novelty(self, ind):
        distances = [self.metric(ind, arch_ind) for arch_ind in self.archive]
        return np.mean(np.sort(distances)[:self.k]) if distances else 0

class WeightedArchive(Archive):

    def __init__(self, size, metric, k = 5):
        super().__init__([], metric)
        self.size = size
        self.weights = []
        self.k = k

    def add(self, population):
        self.archive += population
        self.weights += [1]*len(population)
        overflow_size = len(self.archive)-self.size
        if(overflow_size>0):
            overflow_archive = self.archive[:overflow_size]
            self.archive = self.archive[-overflow_size:]
            overflow_weights = self.archive[:overflow_size]
            self.weights = self.weights[-overflow_size:]
            for i in range(overflow_size):
                distances = np.array([1/self.metric(overflow_archive[i], arch_ind) for arch_ind in self.archive])
                idx = np.argsort(distances)
                sum_distances = np.sum(distances[idx])
                for j in idx:
                    self.weights[j] += overflow_weights[i]*distances[j]/sum_distances

    def novelty(self, ind):
        distances = [self.metric(ind, arch_ind) for arch_ind in self.archive]
        return np.mean(np.sort(distances)[:self.k]) if distances else 0
    
class UniqueArchive(Archive):

    def __init__(self, size, metric, eq, k = 5):
        super().__init__([], metric)
        self.size = size
        self.eq = eq
        self.k = k

    def add(self, population):
        self.archive = [ind for ind in self.archive if not np.any([self.eq(ind,new_ind) for new_ind in population])]
        self.archive += population
        overflow = len(self.archive)-self.size
        if(overflow>0):
            self.archive=self.archive[-overflow:]

    def novelty(self, ind):
        distances = [self.metric(ind, arch_ind) for arch_ind in self.archive]
        return np.mean(np.sort(distances)[:self.k]) if distances else 0

class TrestholdArchive(Archive):
    
    def __init__(self, size, metric, k=5, threshold = None):
        super().__init__(deque(), metric)
        self.size = size
        self.k = k
        self.threshold = threshold
        
    def add(self, population):
        pass
    
    def novelty(self, ind):
        distances = [self.metric(ind, arch_ind) for arch_ind in self.archive]
        closest = np.sort(distances)
        novelty = np.mean(closest[:self.k]) if distances else 0
        
        if len(self.archive)==0 or (novelty>self.threshold and closest[0]>0):
            self.archive.append(ind)
            if(len(self.archive)>self.size):
                self.archive.popleft()
        
        return novelty

#REM: heap priority heapdict.heapdict() an have better performance
class BestNeighbourArchive(Archive):
    
    def __init__(self, size, metric, get_fitness, k=5, threshold = None):
        super().__init__(deque(), metric)
        self.size = size
        self.get_fitness = get_fitness
        self.k = k
        self.threshold = threshold
        
    def add(self, population):
        pass
    
    def novelty(self, ind):
        distances = np.array([self.metric(ind, arch_ind) for arch_ind in self.archive])
        sort_idx = np.argsort(distances)
        novelty = np.mean(distances[sort_idx[:self.k]]) if distances.size > 0 else 0
        closest_idx = sort_idx[0] if sort_idx.size != 0 else None
        
        
        if (not self.archive) or (novelty > self.threshold and distances[closest_idx] > 0):
            self.archive.append(ind)
            if(len(self.archive)>self.size):
                self.archive.popleft()
        elif self.get_fitness(ind) > self.get_fitness(self.archive[closest_idx]):
            self.archive.remove(self.archive[closest_idx])
            self.archive.append(ind)
        
        return novelty

class LocalCompetitionArchive(Archive):
    
    def __init__(self, size, metric, get_fitness, k=5, threshold = None):
        super().__init__(deque(), metric)
        self.size = size
        self.get_fitness = get_fitness
        self.k = k
        self.threshold = threshold
        
    def add(self, population):
        pass
    
    def _local_rank(self, ind, idxs, distances):
        rank = 0
        worst_idx = None
        worst_fitness = None
        
        for idx in idxs:
            if distances[idx]<=self.threshold:
                fit_ind = self.get_fitness(ind)
                fit_arch_ind = self.get_fitness(self.archive[idx])
                if fit_ind > fit_arch_ind:
                    rank += 1
                    if worst_fitness is None or fit_arch_ind < worst_fitness:
                        worst_idx = idx
        
        return rank,worst_idx
    
    def novelty(self, ind):
        distances = np.array([self.metric(ind, arch_ind) for arch_ind in self.archive])
        sort_idx = np.argsort(distances)
        closest_idx = sort_idx[0] if sort_idx.size != 0 else None
        
        novelty = np.sum(distances[sort_idx[:self.k]] > self.threshold) if distances.size > 0 else self.k
        rank, replace_idx = self._local_rank( ind, sort_idx[:self.k], distances)
                    
        if novelty>=rank and (closest_idx==None or distances[closest_idx] > 0):
            self.archive.append(ind)
            if(len(self.archive)>self.size):
                self.archive.popleft()
        elif (distances[closest_idx] > 0 or 
                (replace_idx==closest_idx and 
                self.get_fitness(self.archive[replace_idx])<self.get_fitness(ind))
            ):
            self.archive.remove(self.archive[replace_idx])
            self.archive.append(ind)
        
        return max(novelty,rank)

class Buffer:
    def __init__(self, size, func) -> None:
        self.size = size
        self.func = func
        self.buffer = []
    
    def add(self, data):
        self.buffer.append(data)
        if len(self.buffer)>=self.size:
            self.func(self.buffer)
            self.buffer = []

def get_diversity_subgraph(sorted_nodes, graph, n_rep, compare, diversity_weight):
    diversity_count = int((len(sorted_nodes)-n_rep)*diversity_weight)
    
    def filter_node(node):
        if diversity_count<=0:
            return True
        ind = graph.nodes[node].get("data",None)
        return compare(ind, sorted_nodes[n_rep + diversity_count])
    
    return nx.subgraph_view(graph, filter_node=filter_node)
            

def get_representatives(graph, n_rep, compare, clust_alg='agglomerative', debug=False):
    population = graph.nodes()
    adjacenc_matrix = nx.to_numpy_array(graph)
    
    n_rep = min(n_rep, len(population))
    clustering = None
    if clust_alg == 'spectral' :
        clustering = SpectralClustering(n_clusters=n_rep, affinity='precomputed', assign_labels='discretize').fit(adjacenc_matrix)
    elif clust_alg == 'agglomerative' :  
        clustering = AgglomerativeClustering(n_clusters=n_rep, metric='precomputed', linkage='average').fit(adjacenc_matrix)
    else:
        warnings.warn("Unsuported clustering algorithm chosen. Agglomerative Clustering used instead.")
        clustering = AgglomerativeClustering(n_clusters=n_rep, metric='precomputed', linkage='average').fit(adjacenc_matrix)
    
    labels = clustering.labels_
    results = [None] * n_rep
    for node, label in zip(population, labels):
        ind = graph.nodes[node]["data"]
        if (results[label]==None or 
            compare(ind, results[label])):
            results[label]=ind
    
    if debug:
        _, counts = np.unique(labels,return_counts=True)
        print(f"Archive cluster counts: {counts}")
        return results, labels
    else:
        return results