__all__ = [
    "get_structural_holes_HAM"
]
import sys
import numpy as np
import json, os
import scipy.sparse as sps
import scipy.linalg as spl
from sklearn import metrics
from scipy.cluster.vq import kmeans, vq, kmeans2
from collections import Counter
eps=2.220446049250313e-16
import scipy.stats as stat

def sym(w):
    '''
    Initialize a random orthogonal matrix F = w * (wT * w)^ (-1/2)
    Parameters
    ----------
    w : A random matrix.

    Returns
    -------
    F : a random orthogonal matrix.
    '''
    return w.dot(spl.inv(spl.sqrtm(w.T.dot(w))))

def avg_entropy(predicted_labels, actual_labels):
    '''
    Calculate the average entropy between predicted_labels and actual_labels.

    Parameters
    ----------
    predicted_labels : a Ndarray of predicted_labels.
    actual_labels : a Ndarray of actual_labels.

    Returns
    -------
    A float of average entropy.
    '''
    actual_labels_dict = {}
    predicted_labels_dict = {}
    for label in np.unique(actual_labels):
        actual_labels_dict[label] = np.nonzero(actual_labels==label)[0]
    for label in np.unique(predicted_labels):
        predicted_labels_dict[label] = np.nonzero(predicted_labels==label)[0]
    avg_value = 0
    N = len(predicted_labels)
    # store entropy for each community
    for label, items in predicted_labels_dict.items():
        N_i = float(len(items))
        p_i = []
        for label2, items2  in actual_labels_dict.items():
            common = set(items.tolist()).intersection(set(items2.tolist()))
            p_ij = float(len(common))/ N_i
            p_i.append(p_ij)
        entropy_i = stat.entropy(p_i)
        avg_value += entropy_i * (N_i / float(N))
    return avg_value

def load_adj_matrix(G):
    '''
    Transfer the graph into sparse matrix.
    Parameters
    ----------
    G : graph
        An undirected graph.

    Returns
    -------
    A : A sparse matrix A
    '''
    listE = []
    for edge in G.edges:
        listE.append(edge[0]-1)
        listE.append(edge[1]-1)
    adj_tuples = np.array(listE).reshape(-1,2)
    n = len(np.unique(adj_tuples))
    vals = np.array([1] * len(G.edges))
    max_id = max(max(adj_tuples[:, 0]), max(adj_tuples[:, 1])) + 1
    A = sps.csr_matrix((vals, (adj_tuples[:, 0], adj_tuples[:, 1])), shape=(max_id, max_id))
    A = A + A.T
    return sps.csr_matrix(A)

def majority_voting(votes):
    '''
    majority voting.

    Parameters
    ----------
    votes : a Ndarray of votes

    Returns
    -------
    the most common label.
    '''
    C = Counter(votes)
    pairs = C.most_common(2)
    if len(pairs)==0:
        return 0
    if pairs[0][0] > 0:
        return pairs[0][0]
    elif len(pairs)>1:
        return pairs[1][0]
    else:
        return 0

def label_by_neighbors(AdjMat,labels):
    '''
    classifify SHS using majority voting.

    Parameters
    ----------
    AdjMat : adjacency matrix
    labels : a Ndarray of labeled communities of the nodes.

    Returns
    -------
    labels : a Ndarray of labeled communities of the nodes.
    '''
    assert (AdjMat.shape[0] == len(labels)), "dimensions are not equal"
    unlabeled_idx = (labels==0)
    num_unlabeled = sum(unlabeled_idx)
    count = 0
    while num_unlabeled > 0:
        idxs = np.array(np.nonzero(unlabeled_idx)[0])
        next_labels = np.zeros(len(labels))
        for idx in idxs:
            neighbors = np.nonzero(AdjMat[idx,:] > 0)[1]
            if len(neighbors)==0:
                next_labels[idx] = majority_voting(labels)
            else :
                neighbor_labels = labels[neighbors]
                next_labels[idx] = majority_voting(neighbor_labels)
        labels[idxs] = next_labels[idxs]
        unlabeled_idx = (labels==0)
        num_unlabeled = sum(unlabeled_idx)
    return labels


def get_structural_holes_HAM(G, k, c,ground_truth_labels):
    '''
    using HAM to jointly detect SHS and communities.
    https://dl.acm.org/doi/10.1145/2939672.2939807

    Parameters
    ----------
    G : graph
        An undirected graph.
    k : int
        top - k structural hole spanners
    c : the number of communities

    ground_truth_labels : ground truth labels of nodes.

    Returns
    -------
    a Ndarray of top k nodes as structural hole spanners, and a Ndarray of labeled communities of the nodes.
    '''
    A_mat = load_adj_matrix(G)
    A = A_mat  # adjacency matrix
    n = A.shape[0]  # the number of nodes

    epsilon = 1e-4  # smoothing value: epsilon
    max_iter = 50  # maximum iteration value
    seeeed = 5433
    np.random.seed(seeeed)
    topk = k

    invD = sps.diags((np.array(A.sum(axis=0))[0, :]+eps) ** (-1.0), 0) # Inv of degree matrix D^-1
    L = (sps.identity(n) - invD.dot(A)).tocsr()  # Laplacian matrix L = I - D^-1 * A
    F = sym(np.random.random((n, c))) # Initialize a random orthogonal matrix F

    # Algorithm 1
    for step in range(max_iter):
        Q = sps.identity(n).tocsr()
        P = L.dot(F)
        for i in range(n):
            Q[i, i] = 0.5 / (spl.norm(P[i, :]) + epsilon)

        R = L.T.dot(Q).dot(L)

        W, V = np.linalg.eigh(R.todense())
        Wsort = np.argsort(W)  # sort from smallest to largest
        F = V[:, Wsort[0:c]]  # select the smallest eigenvectors

    # find SH spanner
    SH = np.zeros((n,))
    for i in range(n):
        SH[i] = np.linalg.norm(F[i, :])
    SHrank = np.argsort(SH)  # index of SH



    # METRICS BEGIN

    to_keep_index = np.sort(SHrank[topk:])
    A_temp = A[to_keep_index, :]
    A_temp = A_temp[:, to_keep_index]
    HAM_labels_keep = np.asarray(ground_truth_labels)[to_keep_index]
    allLabels = np.asarray(ground_truth_labels)

    cluster_matrix = F
    labelbook, distortion = kmeans(cluster_matrix[to_keep_index, :], c)
    HAM_labels, dist = vq(cluster_matrix[to_keep_index, :], labelbook)

    print("AMI")
    print('HAM: ' + str(metrics.adjusted_mutual_info_score(HAM_labels, HAM_labels_keep.T[0])))

    # classifify SHS using majority voting
    predLabels = np.zeros(len(ground_truth_labels))
    predLabels[to_keep_index] = HAM_labels + 1

    HAM_predLabels = label_by_neighbors(A, predLabels)
    print('HAM_all: ' + str(metrics.adjusted_mutual_info_score(HAM_predLabels, allLabels.T[0])))

    print("NMI")
    print('HAM: ' + str(metrics.normalized_mutual_info_score(HAM_labels, HAM_labels_keep.T[0])))
    print('HAM_all: ' + str(metrics.normalized_mutual_info_score(HAM_predLabels, allLabels.T[0])))

    print("Entropy")
    print('HAM: ' + str(avg_entropy(HAM_labels, HAM_labels_keep.T[0])))
    print('HAM_all: ' + str(avg_entropy(HAM_predLabels, allLabels.T[0])))

    # METRICS END
