import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
import tabulate

from genome import Genome, NodeType, mutate
from node import NodeType


def _find_layer(g: nx.DiGraph, hidden_node: int, inputs: list[int]) -> int:
    paths = []
    for input_node in inputs:
        paths += list(nx.all_simple_paths(g, input_node, hidden_node))

    path_lengths = [len(path) for path in paths]
    return 2 if len(path_lengths) == 0 else max(path_lengths)


def genome_graph(genome: Genome):
    graph = nx.DiGraph()

    # Add nodes
    for node in genome.nodes.keys():
        graph.add_node(node)

    # Add edges
    for connection in genome.connections.values():
        if connection.disabled:
            continue

        from_node, to_node = connection.nodes
        graph.add_edge(from_node, to_node, weight=connection.weight)

    inputs = [node.id for node in genome.nodes.values() if node.type == NodeType.INPUT]
    hidden = [node.id for node in genome.nodes.values() if node.type == NodeType.HIDDEN]
    outputs = [node.id for node in genome.nodes.values() if node.type == NodeType.OUTPUT]

    for input_node in inputs:
        graph.nodes[input_node]["layer"] = 0

    max_layer = 1
    for hidden_node in hidden:
        layer = _find_layer(graph, hidden_node, inputs)
        max_layer = max(layer, max_layer)
        graph.nodes[hidden_node]["layer"] = layer

    for output_node in outputs:
        graph.nodes[output_node]["layer"] = max_layer + 1

    plt.subplot()
    pos = nx.multipartite_layout(graph, subset_key="layer")
    nx.draw_networkx_nodes(graph, pos, nodelist=inputs, node_color="#ff0000")
    nx.draw_networkx_nodes(graph, pos, nodelist=hidden, node_color="#00ff00")
    nx.draw_networkx_nodes(graph, pos, nodelist=outputs, node_color="#0000ff")
    nx.draw_networkx_labels(graph, pos)
    nx.draw_networkx_edges(graph, pos)


def genome_table(genome: Genome):
    table = [
        (conn.innovation_no, "->".join([str(n) for n in conn.nodes]), "DIS" if conn.disabled else "")
        for conn in genome.connections.values()
    ]
    table.sort(key=lambda c: c[0])
    table = zip(*table)

    print(tabulate.tabulate(table, tablefmt="psql"))


if __name__ == "__main__":
    g1 = Genome.new(3, 2)
    g1.add_connection(0, 4, 0.5)
    mutate(g1)
    mutate(g1)
    mutate(g1)
    mutate(g1)
    mutate(g1)
    mutate(g1)

    # mutate(g1)
    # genome_graph(g1)
    # plt.show()
    genome_table(g1)
