import networkx as nx
import numpy as np
import random
import matplotlib.pyplot as plt
from robustness import *

RANDOM = 0  # Add a random edge
PREFERENTIAL_MIN_MIN = 1    # Connect 2 nodes with the smallest degrees
PREFERENTIAL_MIN_MAX = 2    # Connect nodes with the smallest and largest degrees
PREFERENTIAL_MAX_MAX = 3    # Connect 2 nodes with the largest degrees


def random_addition(graph, fr):
    if fr > 1 or fr < 0:
        raise Exception("Fraction of edges to add has to be between 0 and 1.")
    graph = graph.copy()
    # size of graph is the number of edges
    size = graph.size()
    # number of edges to be added
    n = size * fr
    for i in range(int(n)):
        graph = add_random_edge(graph)
    return graph


def preferential_addition(graph, fr):
    if fr > 1 or fr < 0:
        raise Exception("Fraction of edges to add has to be between 0 and 1.")
    graph = graph.copy()
    # size of graph is the number of edges
    size = graph.size()
    # number of edges to be added
    n = size * fr
    for i in range(int(n)):
        graph = add_preferential_edge_min_min(graph)
    return graph


def add_random_edge(graph):
    """
    Create a new random edge. 
    :param graph: networkx graph
    :return: networkx graph
    """
    edges = list(graph.edges)
    nonedges = list(nx.non_edges(graph))

    # random edge choice
    chosen_edge = random.choice(edges)
    chosen_nonedge = random.choice([x for x in nonedges if chosen_edge[0] == x[0]])

    # add new edge
    graph.add_edge(chosen_nonedge[0], chosen_nonedge[1])

    return graph


def add_preferential_edge_min_min(graph):
    """
    Add an edge between 2 vertices with the least degree.
    :param graph: networkx graph
    :return: networkx graph
    """
    graph = graph.copy()
    nodes = sorted(graph.degree, key=lambda x: x[1], reverse=False)
    graph.add_edge(nodes[0][0], nodes[1][0])
    return graph


def add_preferential_edge_min_max(graph):
    """
    Add an edge between 2 vertices with the smallest and largest degree.
    :param graph: networkx graph
    :return: networkx graph
    """
    graph = graph.copy()
    nodes = sorted(graph.degree, key=lambda x: x[1], reverse=False)
    n_small = nodes[0][0]
    n_large = nodes[-1][0]
    i = 1
    while graph.has_edge(n_small, n_large):
        n_large = nodes[-1 - i][0]
        i += 1
    graph.add_edge(n_small, n_large)
    return graph


def add_preferential_edge_max_max(graph):
    """
    Add an edge between 2 vertices with the largest degree.
    :param graph: networkx graph
    :return: networkx graph
    """
    graph = graph.copy()
    nodes = sorted(graph.degree, key=lambda x: x[1], reverse=True)
    i = 1
    n_small = nodes[0][0]
    n_large = nodes[i][0]
    while graph.has_edge(n_small, n_large):
        i = i + 1
        if i >= len(nodes):
            break
        else:
            n_large = nodes[i][0]
    graph.add_edge(n_small, n_large)
    return graph


def recover_to_initial_diameter(initial_diameter, initial_lcc, attacked_graph, recovery_option=0):
    d, lcc, av_cc = get_robustness(attacked_graph)
    new_diameter = d
    num_edges = 0

    recovered_graph = attacked_graph.copy()

    # Fist we have to make sure there is only 1 LCC
    while lcc < initial_lcc:
        recovered_graph = add_edge(recovery_option, recovered_graph)
        num_edges += 1
        lcc = get_LCC_size(recovered_graph)

    # Then recover the diameter
    while new_diameter > initial_diameter:
        recovered_graph = add_edge(recovery_option, recovered_graph)
        num_edges += 1
        new_diameter = get_diameter(recovered_graph)

    # print_robustness(recovered_graph)
    # print("\nNumber of edges needed to recover:", num_edges)

    return recovered_graph, num_edges


def add_edge(recovery_option, recovered_graph):
    # Select respective recovery option
    if recovery_option == RANDOM:
        recovered_graph = add_random_edge(recovered_graph)
    elif recovery_option == PREFERENTIAL_MIN_MIN:
        recovered_graph = add_preferential_edge_min_min(recovered_graph)
    elif recovery_option == PREFERENTIAL_MIN_MAX:
        recovered_graph = add_preferential_edge_min_max(recovered_graph)
    elif recovery_option == PREFERENTIAL_MAX_MAX:
        recovered_graph = add_preferential_edge_max_max(recovered_graph)
    else:
        raise Exception("Incorrect recovery option specified")

    return recovered_graph