################################################################################
# © Copyright 2021-2022 Zapata Computing Inc.
################################################################################
# TODO : Remove everythin in this file
import itertools
import random
import warnings
from typing import Iterator, Optional, Sequence, TypeVar

import networkx as nx


def compare_graphs(graph1: nx.Graph, graph2: nx.Graph) -> bool:
    """Compares two NetworkX graph objects to see if they are identical.
    NOTE: this is *not* solving isomorphism problem.
    """

    for n1, n2 in zip(graph1.nodes, graph2.nodes):
        if n1 != n2:
            return False
    for e1, e2 in zip(graph1.edges, graph2.edges):
        if e1 != e2:
            return False
    return True


def generate_graph_node_dict(graph: nx.Graph) -> dict:
    """Generates a dictionary containing key:value pairs in the form of
    nx.Graph node : integer index of the node
    """
    nodes_int_map = []
    for node_index, node in enumerate(graph.nodes):
        nodes_int_map.append((node, node_index))
    nodes_dict = dict(nodes_int_map)
    return nodes_dict


T = TypeVar("T")


Sampler = Iterator[T]


def uniform_sampler(min_value: float = 0.0, max_value: float = 1.0) -> Sampler[float]:
    while True:
        yield random.uniform(min_value, max_value)


def constant_sampler(value: T) -> Sampler[T]:
    return itertools.repeat(value)


def choice_sampler(choices: Sequence[T]) -> Sampler[T]:
    while True:
        yield random.choice(choices)


def normal_sampler(mu: float = 0.0, sigma: float = 1.0) -> Sampler[float]:
    while True:
        yield random.normalvariate(mu, sigma)


def generate_random_graph_erdos_renyi(
    num_nodes: int,
    edge_probability: float,
    weight_sampler: Sampler = constant_sampler(1.0),
    seed: Optional[int] = None,
) -> nx.Graph:
    """Randomly generate a graph from Erdos-Renyi ensemble.  A graph is constructed by
    connecting nodes randomly.  Each edge is included in the graph with probability p
    independent from every other edge. Equivalently, all graphs with n nodes and M edges
    have equal probability.
    Args:
        num_nodes: Number of nodes in the result graph.
        edge_probability: Probability of connecting two nodes.
        weight_sampler: Used to sample edge weights. Defaults `static_sampler`,
            i.e. all edge weights are set to 1.0.
        seed: if provided, sets the global seed
    """
    output_graph = nx.erdos_renyi_graph(n=num_nodes, p=edge_probability, seed=seed)
    _weight_graph_edges(output_graph, weight_sampler, seed)

    return output_graph


def generate_random_regular_graph(
    num_nodes: int,
    degree: int,
    weight_sampler: Sampler = constant_sampler(1.0),
    seed: Optional[int] = None,
) -> nx.Graph:
    """Randomly generate a d-regular graph.
    A graph is generated by picking uniformly a graph among the set of graphs
    with the desired number of nodes and degree.
    Args:
        num_nodes: Number of nodes in the generated graph.
        degree: Degree of each edge.
        weight_sampler: Used to sample edge weights. Defaults `static_sampler`,
            i.e. all edge weights are set to 1.0.
        seed: if provided, sets the global seed
    """
    output_graph = nx.random_regular_graph(d=degree, n=num_nodes, seed=seed)
    _weight_graph_edges(output_graph, weight_sampler, seed)

    return output_graph


def generate_caveman_graph(
    number_of_cliques: int,
    size_of_cliques: int,
    weight_sampler: Sampler = constant_sampler(1.0),
    seed: Optional[int] = None,
) -> nx.Graph:
    output_graph = nx.caveman_graph(number_of_cliques, size_of_cliques)
    _weight_graph_edges(output_graph, weight_sampler, seed)
    return output_graph


def generate_ladder_graph(
    length_of_ladder: int,
    weight_sampler: Sampler = constant_sampler(1.0),
    seed: Optional[int] = None,
) -> nx.Graph:
    graph = nx.ladder_graph(length_of_ladder)
    _weight_graph_edges(graph, weight_sampler, seed)
    return graph


def generate_barbell_graph(
    number_of_vertices_complete_graph: int,
    weight_sampler: Sampler = constant_sampler(1.0),
    seed: Optional[int] = None,
) -> nx.Graph:
    graph = nx.barbell_graph(number_of_vertices_complete_graph, 0)
    _weight_graph_edges(graph, weight_sampler, seed)

    return graph


def _weight_graph_edges(
    graph: nx.Graph,
    sampler: Sampler,
    seed: Optional[int] = None,
) -> nx.Graph:
    """Update the weights of all the edges of a graph in place.
    Args:
        graph: The graph to mutate.
        sampler:
    """
    if graph.is_multigraph():
        raise ValueError("Cannot deal with multigraphs")

    if seed is not None:
        random.seed(seed)

    weighted_edges = [(e[0], e[1], next(sampler)) for e in graph.edges]

    # If edges already present, it will effectively update them (except for multigraph)
    graph.add_weighted_edges_from(weighted_edges)
