import networkx as nx
import numpy as np
import scipy.sparse as sparse
from networkx.algorithms import bipartite, centrality


def generate_graph(hashtag_tweet_dict) -> dict:
    graph = build_hashtag_tweet_adjacency_matrix(hashtag_tweet_dict)
    if not graph:
        return

    graph["nx_graph"], graph["nx_projected"] = get_projected_graph(graph)
    # graph["edgelist"] = nx.to_edgelist(graph["nx_projected"])
    graph["projected_adjacency_matrix"] = nx.to_numpy_array(graph["nx_projected"])
    return graph


def get_projected_graph(_graph: dict, along: str = "hashtag") -> tuple:
    """Retrieve projected bipartite graph from csr matrix.

    Args:
        graph: Dict-like object as generated by the generate_adjacency_matrix function.

    Returns:
        NetworkX graphs (bipartite graph and projected representation of the bipartite).

    """

    nx_graph = nx.from_scipy_sparse_matrix(_graph["mtx"])

    if along == "hashtag":
        projection = list(map(int, list(_graph["idx_hashtag"])))
    elif along == "tweet":
        projection = list(map(int, list(_graph["idx_tweet"])))
    else:
        raise ValueError("Specify a valid projection class (user or hashtag)!")

    nx_projected_graph = bipartite.projected_graph(nx_graph, projection)

    return nx_graph, nx_projected_graph


def build_hashtag_tweet_adjacency_matrix(hashtag_tweet_dict: dict) -> dict:
    # building subsets of the hashtag_tweet_dict, tweet_hashtag_dict
    tweet_idx, idx_tweet, hashtag_idx, idx_hashtag = {}, {}, {}, {}
    node_idx = 0
    edge_list, edge_weights = [], {}
    for keyword in hashtag_tweet_dict:
        if keyword not in hashtag_idx:
            hashtag_idx[keyword] = node_idx
            idx_hashtag[node_idx] = keyword
            node_idx += 1

        # iterate over all tweets connected to each hashtag
        for tweet in hashtag_tweet_dict[keyword]:
            if tweet not in tweet_idx:
                tweet_idx[tweet] = node_idx
                idx_tweet[node_idx] = tweet
                node_idx += 1

            # build weighted adj matrix
            edge = (hashtag_idx[keyword], tweet_idx[tweet])
            if edge not in edge_weights:
                edge_list.append(edge)
                edge_weights[edge] = 1
            else:
                edge_weights[edge] += 1

    if node_idx == 0:
        return
    edge_array = np.array(edge_list)
    shape = (node_idx, node_idx)
    weights = np.array([edge_weights[edge] for edge in (edge_list)]).astype(np.int32)

    mtx = sparse.coo_matrix(
        (weights, (edge_array[:, 0], edge_array[:, 1])),
        dtype=edge_array.dtype,
        shape=shape,
    ).tocsr()

    return dict(
        mtx=mtx,
        idx_hashtag=idx_hashtag,
        hashtag_idx=hashtag_idx,
        idx_tweet=idx_tweet,
        tweet_idx=tweet_idx,
    )


def get_closeness_over_labelled_set(sp: dict, labelled_set: set, graph: dict) -> float:
    sub_sp = {}
    for node in labelled_set:
        # graph may not be connected
        if graph["hashtag_idx"][node] in sp:
            sub_sp[graph["hashtag_idx"][node]] = sp[graph["hashtag_idx"][node]]
            assert graph["hashtag_idx"][node] in graph["nx_projected"].nodes()

    totsp = sum(sub_sp.values())
    len_G = len(labelled_set)
    _closeness_centrality = 0.0
    if totsp > 0.0 and len_G > 1:
        _closeness_centrality = (len(sub_sp) - 1.0) / totsp
        s = (len(sub_sp) - 1.0) / (len_G - 1)
        _closeness_centrality *= s

    return _closeness_centrality


def compounded_closeness(
    graph: dict, positive_labels: set, negative_labels: set, graph_key: str
) -> dict:
    G = graph[graph_key]

    path_length = nx.single_source_shortest_path_length
    nodes = G.nodes

    closeness_centrality = {}
    for n in nodes:
        sp = path_length(G, n)
        tmp_pos_centrality = get_closeness_over_labelled_set(sp, positive_labels, graph)
        tmp_neg_centrality = get_closeness_over_labelled_set(sp, negative_labels, graph)
        closeness_centrality[n] = tmp_pos_centrality - tmp_neg_centrality

    return closeness_centrality


def get_centrality_measure(
    _graph: dict,
    centrality_measure: str = "degree",
    graph_key: str = "nx_projected",
    positive_labels: set = None,
    negative_labels: set = None,
) -> dict:
    """Use networkx to compute the centrality on each node of the graph.

    Perform centrality computation over the projected bipartite graph.

    Args:
        _graph: Dict-like object containing user-hashtag graph related information (
            adjacency matrix, edge list, networkx graph, networkx bipartite projected
            graph, etc).

    Returns:
        A dict of node ids to centrality measure in the graph.

    """

    assert graph_key in _graph, "Include the bipartite projection in _graph!"

    if centrality_measure == "degree":
        centrality_func = centrality.degree_centrality
    elif centrality_measure == "katz":
        centrality_func = centrality.katz_centrality
    elif centrality_measure == "eigenvector":
        centrality_func = centrality.eigenvector_centrality
    elif centrality_measure == "betweenness":
        centrality_func = centrality.betweenness_centrality
    elif centrality_measure == "closeness":
        centrality_func = centrality.closeness_centrality
    else:
        raise ValueError("Unrecognized centrality measure!")

    return {
        _graph["idx_hashtag"][idx]: centrality
        for idx, centrality in sorted(
            centrality_func(_graph[graph_key]).items(),
            key=lambda item: item[1],
            reverse=True,
        )
        if _graph["idx_hashtag"][idx] not in positive_labels
        and _graph["idx_hashtag"][idx] not in negative_labels
    }
