'''
This script allows the user to evaluate a classification system on new labels using clustering methods.
The algorithms are applied on the given latent space (embedding).
'''
import argparse
import numpy as np
import pandas as pd
import os
import time
import pickle
import csv
import json

from sklearn.preprocessing import LabelEncoder
from sklearn.metrics.pairwise import pairwise_distances
from sklearn.cluster import KMeans
from sklearn.manifold import TSNE
from sklearn.metrics import f1_score, homogeneity_score, completeness_score, v_measure_score
import matplotlib.pyplot as plt

from volia.data_io import read_features,read_lst
from volia.measures import entropy_score, purity_score

'''
TODO: 
- Add an option allowing the user to choose the number of 
clustering to train in order to compute the average and the
'''


def train_clustering(label_encoder, feats, classes, outdir):
    num_classes = len(label_encoder.classes_)
    estimator = None
    kmeans_filepath = os.path.join(outdir, f"{args.prefix}kmeans.pkl") 
    if args.onlymeasures:
        print(f"Loading model: {kmeans_filepath}")
        with open(kmeans_filepath, "rb") as f:
            estimator = pickle.load(f)
    else:
        # Compute KMEANS clustering on data
        print("Saving parameters")
        kmeans_parameters = {
            "n_clusters": num_classes,
            "n_init": 100,
            "tol": 10-6,
            "algorithm": "elkan"
        }
        with open(os.path.join(outdir, f"{args.prefix}kmeans_parameters.json"), "w") as f:
            json.dump(kmeans_parameters, f)

        # Fit the model and Save parameters
        print(f"Fit the model: {kmeans_filepath}")
        estimator = KMeans(
            **kmeans_parameters
        )
        estimator.fit(feats)
        print(f"Kmeans: processed {estimator.n_iter_} iterations - intertia={estimator.inertia_}")

        with open(kmeans_filepath, "wb") as f:
            pickle.dump(estimator, f)
    
    # contains distance to each cluster for each sample
    dist_space = estimator.transform(feats)
    predictions = np.argmin(dist_space, axis=1)

    # gives each cluster a name (considering most represented character)
    dataframe = pd.DataFrame({
        "label": pd.Series(list(map(lambda x: le.classes_[x], labels))),
        "prediction": pd.Series(predictions)
    })

    def find_cluster_name_fn(c):
        mask = dataframe["prediction"] == c
        return dataframe[mask]["label"].value_counts(sort=False).idxmax()
    
    cluster_names = list(map(find_cluster_name_fn, range(num_classes)))
    predicted_labels = le.transform(
        [cluster_names[pred] for pred in predictions])
    
    # F-measure
    fscores = f1_score(labels, predicted_labels, average=None)
    fscores_str = "\n".join(map(lambda i: "{0:25s}: {1:.4f}".format(le.classes_[i], fscores[i]), range(len(fscores))))
    
    # Entropy
    _, _, entropy = entropy_score(labels, predicted_labels)

    # Homogenity
    homogeneity = homogeneity_score(labels, predicted_labels)

    # Completeness
    completeness = completeness_score(labels, predicted_labels)

    # V-Measure
    v_measure = v_measure_score(labels, predicted_labels)

    # Purity
    purity_scores = purity_score(labels, predicted_labels)
    purity_class_score = purity_scores["purity_class_score"]
    purity_cluster_score = purity_scores["purity_cluster_score"]
    K = purity_scores["K"]

    # Write results
    with open(os.path.join(outdir, args.prefix + "eval_clustering.log"), "w") as fd:
        print(f"F1-scores for each classes:\n{fscores_str}", file=fd)
        print(f"Entropy: {entropy}", file=fd)
        print(f"Global score : {np.mean(fscores)}", file=fd)
        print(f"Homogeneity: {homogeneity}", file=fd)
        print(f"completeness: {completeness}", file=fd)
        print(f"v-measure: {v_measure}", file=fd)
        print(f"purity class score: {purity_class_score}", file=fd)
        print(f"purity cluster score: {purity_cluster_score}", file=fd)
        print(f"purity overall evaluation criterion (K): {K}", file=fd)

    # Process t-SNE and plot
    tsne_estimator = TSNE()
    embeddings = tsne_estimator.fit_transform(feats)
    print("t-SNE: processed {0} iterations - KL_divergence={1:.4f}".format(
        tsne_estimator.n_iter_, tsne_estimator.kl_divergence_))

    fig, [axe1, axe2] = plt.subplots(1, 2, figsize=(10, 5))
    for c, name in enumerate(le.classes_):
        c_mask = np.where(labels == c)
        axe1.scatter(embeddings[c_mask][:, 0], embeddings[c_mask][:, 1], label=name, alpha=0.2, edgecolors=None)

        try:
            id_cluster = cluster_names.index(name)
        except ValueError:
            print("WARNING: no cluster found for {}".format(name))
            continue
        c_mask = np.where(predictions == id_cluster)
        axe2.scatter(embeddings[c_mask][:, 0], embeddings[c_mask][:, 1], label=name, alpha=0.2, edgecolors=None)
    
    axe1.legend(loc="lower center", bbox_to_anchor=(0.5, -0.35))
    axe1.set_title("true labels")
    axe2.legend(loc="lower center", bbox_to_anchor=(0.5, -0.35))
    axe2.set_title("predicted cluster label")

    plt.suptitle("Kmeans Clustering")

    loc = os.path.join(
        outdir,
        args.prefix + "kmeans.pdf"
    )
    plt.savefig(loc, bbox_inches="tight")
    plt.close()

    print("INFO: figure saved at {}".format(loc))

    end = time.time()
    print("program ended in {0:.2f} seconds".format(end-start))
    return {
        "f1": np.mean(fscores),
        "entropy": entropy,
        "homogeneity": homogeneity,
        "completeness": completeness,
        "v-measure": v_measure,
        "purity_class_score": purity_class_score,
        "purity_cluster score": purity_cluster_score,
        "K": K
    }


if __name__ == "__main__":
    # Argparse
    parser = argparse.ArgumentParser("Compute clustering on a latent space")
    parser.add_argument("features")
    parser.add_argument("utt2",
                        type=str,
                        help="file with [utt] [value]")
    parser.add_argument("--idsfrom",
                        type=str,
                        default="utt2",
                        choices=[
                            "features",
                            "utt2"
                        ],
                        help="from features or from utt2?")
    parser.add_argument("--prefix",
                        default="",
                        type=str,
                        help="prefix of saved files")
    parser.add_argument("--outdir",
                        default=None,
                        type=str,
                        help="Output directory")
    parser.add_argument("--nmodels",
                        type=int,
                        default=1,
                        help="specifies the number of models to train")
    parser.add_argument("--onlymeasures",
                        action='store_true',
                        help="Don't compute the clustering, compute only the measures")
    args = parser.parse_args()

    assert args.outdir

    start = time.time()

    # Load features and utt2
    features = read_features(args.features)
    utt2 = read_lst(args.utt2)

    # Take id list
    if args.idsfrom == "features":
        ids = list(features.keys())
    elif args.idsfrom == "utt2":
        ids = list(utt2.keys())
    else:
        print(f"idsfrom is not good: {args.idsfrom}")
        exit(1)
    
    feats = np.vstack([ features[id_] for id_ in ids ])
    classes = [ utt2[id_] for id_ in ids ]

    # Encode labels
    le = LabelEncoder()
    labels = le.fit_transform(classes)
    
    measures = {}
    for i in range(1, args.nmodels+1):
        subdir = os.path.join(args.outdir, str(i))
        if not os.path.exists(subdir):
            os.mkdir(subdir)
        print(f"[{i}/{args.nmodels}] => {subdir}")
        results = train_clustering(le, feats, classes, subdir)

        for key, value in results.items():
            if key not in measures:
                measures[key] = []
            measures[key].append(results[key])
    

    # File with results
    file_results = os.path.join(args.outdir, args.prefix + "clustering_measures.txt")

    with open(file_results, "w") as f:
        f.write(f"[nmodels: {args.nmodels}]\n")
        for key in measures.keys():
            values = np.asarray(measures[key], dtype=float)
            mean = np.mean(values)
            std = np.std(values)
            f.write(f"[{key} => mean: {mean}, std: {std}] \n")
            
    # CSV File with all the values
    file_csv_measures = os.path.join(args.outdir, args.prefix + "clustering_measures.csv")

    with open(file_csv_measures, "w", newline="") as f:
        writer = csv.writer(f, delimiter=",")
        writer.writerow(["measure"] + list(range(1, args.nmodels+1)) + ["mean"] + ["std"])
        for key in measures.keys():
            values = np.asarray(measures[key], dtype=float)
            mean = np.mean(values)
            std = np.std(values)
            writer.writerow([key] + list(values) + [mean] + [std])