import networkx as nx
import osmnx as ox
from src.helper.debug_printer import debug_print
from sklearn.cluster import KMeans, BisectingKMeans

# j'adore la duplication de code

# TODO : à commenter et debug, pour l'instant go dodo


def display_partitions(G, partitions):
    """
        Affiche sur les graphe 'G' (souvent Montreal), chaque partitions d'une couleur unique
    Parameters:
        G (nx.DiGraph ou MultiDiGraph) : le graphe qui contient les partitions
        partitions (list[nx.DiGraph]) : des sous graphes incluent dans 'G'
    Return:
        void
    """
    if len(partitions) > 50:
        raise ValueError(
            "src/helper/partitions; display_partition: ne peux pas afficher plus de 50 couleurs différentes, la partition est probablement valide mais inaffichable")

    colors = [
        "#e6194b", "#3cb44b", "#ffe119", "#0082c8", "#f58231", "#911eb4", "#46f0f0", "#f032e6",
        "#d2f53c", "#fabebe", "#008080", "#e6beff", "#aa6e28", "#fffac8", "#800000", "#aaffc3",
        "#808000", "#ffd8b1", "#000080", "#808080", "#FFFFFF", "#000000", "#e41a1c", "#377eb8",
        "#4daf4a", "#984ea3", "#ff7f00", "#ffff33", "#a65628", "#f781bf", "#999999", "#66c2a5",
        "#fc8d62", "#8da0cb", "#e78ac3", "#a6d854", "#ffd92f", "#e5c494", "#b3b3b3", "#1b9e77",
        "#d95f02", "#7570b3", "#e7298a", "#66a61e", "#e6ab02", "#a6761d", "#666666", "#7fc97f",
        "#beaed4", "#fdc086"
    ]

    associated_color = {}

    i = 0
    for subg in partitions:
        for u, v in subg.edges:
            if (v, u) in G.edges(keys=False):
                associated_color[(v, u)] = colors[i]
            associated_color[(u, v)] = colors[i]

        i += 1

    count = 0
    edge_colors = []
    for u, v in G.edges(keys=False):
        if (u, v) in associated_color:
            edge_colors.append(associated_color[(u, v)])
            count += 1
        else:
            edge_colors.append("grey")

    print(f"{count} arêtes colorées (quartiers + connexions)")

    fig, ax = ox.plot_graph(G, edge_color=edge_colors,
                            node_size=0, edge_linewidth=1)


def partition(G, K, method="KMeans", debug_mode=True):
    """
        Partitionne le graphe 'G', en 'K' partitions, selon la méthode 'method', par zone géographique,
        G pour le cas des déneigeuses est l'union des graphes de chaque quartiers
    Parameters:
        G (nx.DiGraph ou MultiDiGraph) : le graphe qui contient les partitions
        K (int) : le nombre de partitions voulues
        method (string: default = "KMeans") : choix de méthode de partitions, soit "KMeans" soit "BisectingKMeans"
        debug_mode (bool: default = True) : activation du mode debug ou pas
    Return:
        list(nx.DiGraph) : des sous graphes incluent dans 'G', 
    """

    # WARNING : G à très certainement été créé à partir de generate_quartier, qui crée un graphe à partir de l'output du drone qui NE contient PAS la data initiale du graphe (length, gps et autre)

    debug_print("debut de l'étape de partitionnement", debug_mode)

    centroids = []
    edges = []
    for u, v, data in G.edges(data=True):
        x1, y1 = G.nodes[u]["x"], G.nodes[u]["y"]
        x2, y2 = G.nodes[v]["x"], G.nodes[v]["y"]
        cx, cy = (x1 + x2)/2, (y1 + y2)/2
        centroids.append([cx, cy])
        edges.append((u, v, data))

    part = None
    if (method == "KMeans"):
        part = KMeans(n_clusters=K, random_state=0).fit(centroids)
    elif (method == "BisectingKMeans"):
        part = BisectingKMeans(n_clusters=K, random_state=0).fit(centroids)
    else:
        raise ValueError(
            "src/helper/partitions; partition: method invalide, choisir entre 'KMeans' et 'BisectingKMeans'")

    labels = part.labels_

    partitioned_subgraphs = [nx.DiGraph() for _ in range(K)]
    for (u, v, data), label in zip(edges, labels):
        partitioned_subgraphs[label].add_edge(u, v, **data)
        partitioned_subgraphs[label].add_node(u, **G.nodes[u])
        partitioned_subgraphs[label].add_node(v, **G.nodes[v])

    debug_print("fin de l'étape de partitionnement", debug_mode)

    if (debug_mode):
        tot_edges = 0
        tot_weight = 0
        lst_edges = []
        lst_weight = []

        for partitioned in partitioned_subgraphs:
            tot_edges += len(partitioned.edges)
            lst_edges.append(len(partitioned.edges))
            tmp_w = 0
            for u, v, data in partitioned.edges(data=True):
                tot_weight += data["length"]
                tmp_w += data["length"]
            lst_weight.append(tmp_w)

        debug_print("List du nb d'edges", debug_mode)
        print(lst_edges)
        debug_print("List des poids", debug_mode)
        print(lst_weight)
        debug_print("Nombre total d'edges", debug_mode)
        print(tot_edges)
        debug_print("Poids total", debug_mode)
        print(tot_weight)

    return partitioned_subgraphs