## label propagation routines

## label propagation algorithms:
import networkx as nx
import pandas as pd
from sklearn.model_selection import ShuffleSplit
from ..general.benchmark_classification import *
import numpy as np
import scipy.sparse as sp
import time
import multiprocessing as mp ## initialize the MP part
import mkl

def label_propagation_normalization(matrix):
    matrix = matrix.tocsr()
    try:
        matrix.setdiag(0)
    except TypeError:
        matrix.setdiag(np.zeros(matrix.shape[0]))
    d = matrix.sum(axis=1).getA1()
    nzs = np.where(d > 0)
    d[nzs] = 1 / np.sqrt(d[nzs])
    dm = sp.diags(d, 0).tocsc()
    return dm.dot(matrix).dot(dm)

    
def label_propagation(graph_matrix, class_matrix, alpha, epsilon=1e-12, max_steps=100000):
    # This method assumes the label-propagation normalization and a symmetric matrix with no rank sinks.
    diff = np.inf
    steps = 0
    current_labels = class_matrix
    while diff > epsilon and steps < max_steps:
        steps += 1
        new_labels = alpha * graph_matrix.dot(current_labels) + (1 - alpha) * class_matrix
        diff = np.linalg.norm(new_labels - current_labels) / np.linalg.norm(new_labels)
        current_labels = new_labels
    return current_labels


def validate_label_propagation(core_network,labels,dataset_name="test",repetitions=5):

    try:
        labels= labels.todense()
    except:
        pass
    
    matrix = label_propagation_normalization(core_network)
    print("Propagation..")
    results = []
    for k in range(repetitions):
        for j in np.arange(0.1,1,0.1):
            print("Train size:{}".format(j))
            rs = ShuffleSplit(n_splits=10, test_size=j, random_state=5123)
            micros = []
            macros = []
            times = []
            for X_train, X_test in rs.split(labels):
                start = time.time()
                tmp_labels = labels.copy()
                true_labels = tmp_labels[X_test].copy()
                tmp_labels[X_test] = 0
                probs = label_propagation(matrix,tmp_labels,0.001)
                micro,macro = evaluate_oracle_F1(probs,None,tmp_labels,input_proba=True)

                end = time.time()
                elapsed = end - start
                micros.append(micro)
                macros.append(macro)
                times.append(elapsed)

            outarray = [str(np.round(1-j,1)),str(np.mean(micros)),str(np.mean(macros)) ,"LP_basic" ,dataset_name,str(np.mean(times))]
            results.append(outarray)

    ## create pd object for easier manipulation
    cnames = ["percent_train","micro_F","macro_F","setting","dataset","time"]
    df = pd.DataFrame(results)
    df.columns = cnames
    return df

def label_propagation_tf():
    ## todo..
    pass
