import random
from statistics import mean

import matplotlib.pyplot as plt
import networkx as nx
from matplotlib import animation


class Node:
    def __init__(self, is_infected=False):
        self.id = random.randint(1, 2000000)
        self.is_infected = is_infected
        self.under_attack = False

    def as_tuple(self):
        return self.id, self.is_infected

    def __repr__(self):
        return f'id: {self.id}, infected: {self.is_infected}'


class Edge:
    def __init__(self, node_a: Node, node_b: Node, weight: float):
        self.node_a = node_a
        self.node_b = node_b
        self.weight = weight

    def as_tuple(self):
        return self.node_a, self.node_b, {'weight': self.weight}

    def has_node(self, node: Node) -> bool:
        return self.node_a is node or self.node_b is node


class Graph:
    def __init__(self, use_weights=False):
        self.edges = []
        self.use_weights = use_weights
        self.rounds_survived = 0

    def add_edge(self, edge: Edge):
        self.edges.append(edge)

    def add_edges(self, edges: [Edge]):
        [self.edges.append(e) for e in edges]

    def get_nodes(self) -> [Node]:
        nodes = set()
        for edge in self.edges:
            nodes.add(edge.node_a)
            nodes.add(edge.node_b)
        return nodes

    def clear_attacked(self) -> None:
        for edge in self.edges:
            edge.node_a.under_attack = False
            edge.node_b.under_attack = False

    def get_adjacent_nodes(self, node: Node) -> [(Node, int)]:
        """
        :param node: Node to search for
        :return: An array of tuples (node, weight)
        """
        edges_with_node = filter(lambda ed: ed.has_node(node), self.edges)
        nodes = set()
        for e in edges_with_node:
            if e.node_a is node:
                nodes.add((e.node_b, e.weight))
            else:
                nodes.add((e.node_a, e.weight))

        return nodes

    def is_alive(self):
        nodes_alive = list(filter(lambda x: not x.is_infected, self.get_nodes()))
        return len(nodes_alive) > 0

    def update_survived(self):
        if not self.is_alive():
            return

        self.rounds_survived += 1

    def infect_step(self):
        infected_nodes = list(filter(lambda n: n.is_infected, self.get_nodes()))
        for node in infected_nodes:
            adjacent_nodes = self.get_adjacent_nodes(node)
            if self.use_weights:
                to_be_infected = random.choices([n[0] for n in adjacent_nodes], weights=[n[1] for n in adjacent_nodes])[0]
            else:
                to_be_infected = random.choice([n[0] for n in adjacent_nodes])
            to_be_infected.is_infected = True
            to_be_infected.under_attack = True
        self.update_survived()


def update(num, layout, g_repr, ax, our_graph: Graph):
    """
    This function is called every 'step', so if you wish to update the graph, do it here
    """
    if not our_graph.is_alive():
        return

    if num != 0:
        our_graph.infect_step()
    ax.clear()

    ax.set_title(f'Step: {num}', loc='right', fontsize=30)

    colors = ['red' if n.is_infected else 'blue' for n in g_repr]
    edgecolors = ['black' if n.under_attack else 'none' for n in g_repr]
    linewidths = [3 if c == 'black' else 0 for c in edgecolors]
    sizes = [300 if n.is_infected else 150 for n in g_repr]
    nx.draw(
        g_repr,
        ax=ax,
        pos=layout,
        node_color=colors,
        linewidths=linewidths,
        edgecolors=edgecolors,
        with_labels=False,
        node_size=sizes,
        alpha=0.7,
    )

    our_graph.clear_attacked()


def do_graph_animation(output_file_name: str, in_graph: Graph, frame_count: int, layout):
    g_repr = nx.Graph()
    # Convert our graph class into tuples understood by networkx
    g_repr.add_edges_from([e.as_tuple() for e in in_graph.edges])

    layout = layout(g_repr)
    fig, ax = plt.subplots()

    fig.set_figwidth(8)
    fig.set_figheight(8)

    anim = animation.FuncAnimation(
        fig, update, frames=frame_count, interval=500, fargs=(layout, g_repr, ax, in_graph)
    )
    anim.save(output_file_name)

    plt.style.use('seaborn')
    plt.show()


def rank_avg(edges, digits=2):
    ranks = {}
    for e in edges:
        ranks[e.node_a] = ranks.get(e.node_a, 0) + 1
        ranks[e.node_b] = ranks.get(e.node_b, 0) + 1
    return round(mean(ranks.values()), digits)


def bus_network(n=30, infected_idx=0) -> tuple[Graph, float]:
    network = Graph()
    nodes = [Node() for _ in range(n)]
    nodes[infected_idx].is_infected = True
    edges = [Edge(nodes[i], nodes[i + 1], 1.0) for i in range(n - 1)]

    network.add_edges(edges)
    return network, rank_avg(edges)


def star_network(cluster_count=5, starsize=6) -> tuple[Graph, float]:
    node_count = cluster_count + cluster_count * starsize + 1
    nodes = [Node() for _ in range(node_count)]

    edges = []
    for x in range(cluster_count):
        center_node = x * starsize + x
        edges += [Edge(nodes[center_node], nodes[i], 1.0) for i in range(center_node + 1, center_node + starsize + 1)]
        edges.append(Edge(nodes[-1], nodes[center_node], 1.0))

    network = Graph()
    network.add_edges(edges)

    return network, rank_avg(edges)


def ring_network(n=30) -> tuple[Graph, float]:
    network = Graph()
    nodes = [Node() for _ in range(n)]
    nodes[0].is_infected = True
    edges = [Edge(nodes[i], nodes[i + 1], 1.0) for i in range(n - 1)]
    end_edge = Edge(nodes[n - 1], nodes[0], 1.0)
    edges.append(end_edge)

    network.add_edges(edges)
    return network, rank_avg(edges)


def main():
    bus, star_avg_rank = bus_network(36)
    do_graph_animation('bus.gif', bus, 90, nx.spiral_layout)

    ring, star_avg_rank = ring_network(36)
    do_graph_animation('ring.gif', ring, 90, nx.circular_layout)

    star, star_avg_rank = star_network()
    list(star.get_nodes())[1].is_infected = True
    do_graph_animation('star.gif', star, 90, nx.kamada_kawai_layout)

    print(f'bus survived {bus.rounds_survived} rounds')
    print(f'ring survived {ring.rounds_survived} rounds')
    print(f'star survived {star.rounds_survived} rounds')


if __name__ == "__main__":
    main()