from sklearn.cluster import KMeans
from sklearn.metrics.cluster import normalized_mutual_info_score, adjusted_rand_score 

def cluster_test(embedding, n_clusters, y, random_state, test_num = 10):
    np_embedding = embedding.detach().numpy() 
    np_y = y.detach().numpy() 
    NMI = 0
    ARI = 0
    for i in range(test_num):
        NMI_, ARI_ = cluster_one_test(np_embedding, n_clusters, np_y, random_state)
        NMI+=NMI_
        ARI+=ARI_

    return NMI/test_num, ARI/test_num

def cluster_one_test(embedding, n_clusters, y,random_state):

    y_pred = KMeans(n_clusters = n_clusters, random_state=random_state, max_iter=500, n_init= 'auto').fit_predict(embedding)
    
    NMI = normalized_mutual_info_score(y, y_pred)
    ARI = adjusted_rand_score(y, y_pred)
    return NMI, ARI
