import argparse
from os import path, mkdir
from utils import SubCommandRunner
from core.data import read_features, read_lst, read_labels, write_line
import numpy as np
from sklearn.cluster import KMeans
import pickle
from clustering_modules.kmeans import kmeans
from clustering_modules.kmeans_mahalanobis import  kmeansMahalanobis
from clustering_modules.kmeans_multidistance import kmeansMultidistance

from sklearn.preprocessing import LabelEncoder
from sklearn.metrics import v_measure_score, homogeneity_score, completeness_score

import core.measures
import json


CLUSTERING_METHODS = {
    "k-means": kmeans(),
    "k-means-mahalanobis": kmeansMahalanobis(),
    "k-means-mahalanobis-constrained": kmeansMahalanobis(constrained=True),
    "k-means-basic-mahalanobis": kmeansMultidistance(distance="mahalanobis"),
    "k-means-basic-cosine": kmeansMultidistance(distance="cosine")
}

KMEANS_METHODS = [key for key in CLUSTERING_METHODS if key.startswith("k-means")]

EVALUATION_METHODS = {
    "entropy": core.measures.entropy_score,
    "purity": core.measures.purity_score,
    "v-measure": v_measure_score,
    "homogeneity": homogeneity_score,
    "completeness": completeness_score,
}


def disequilibrium_run():
    pass


def measure_run(measure: str, features: str, lst: str, truelabels: str, model: str, modeltype: str):
    """

    @param measure:
    @param features:
    @param lst:
    @param truelabels:
    @param model:
    @param modeltype:
    @return:
    """
    module = CLUSTERING_METHODS[modeltype]
    module.load(model)

    eval = {}
    for ms in measure:
        evaluation = EVALUATION_METHODS[ms]
        feats_dict = read_features(features)
        labels_dict = read_labels(truelabels)
        lst_dict = read_lst(lst)
        lst_keys = [key for key in lst_dict]
        feats = np.asarray([feats_dict[key] for key in lst_keys])
        Y_pred = module.predict(feats)
        Y_truth = [labels_dict[key][0] for key in lst_keys]

        le = LabelEncoder()
        le.fit(Y_truth)
        Y_truth = le.transform(Y_truth)

        eval[ms] = evaluation(Y_truth, Y_pred)

    print(json.dumps(eval))


def kmeans_run(features: str,
               lst: str,
               k:int,
               kmax: int,
               klist,
               maxiter: int,
               ninit: int,
               output: str,
               tol: float,
               modeltype: str,
               debug: bool = False):
    """

    @param features: output features
    @param lst: list file
    @param k: k (kmin if kmax specified)
    @param kmax: maximum k to compute
    @param klist: list of k values to compute, ignore k value
    @param output: output file if kmax not specified, else, output directory
    @param mahalanobis: distance option of k-means.
    """
    json_content = locals().copy()

    def fit_model(k: int, output_file):
        if debug:
            print(f"Computing clustering with k={k}")
        model = CLUSTERING_METHODS[modeltype]
        model.fit(X, k, tol, ninit, maxiter, debug)
        model.save(output_file)
        json_content["models"].append({
            "model_file": output_file,
            "k": k,
        })

    json_content["models"] = []

    # -- READ FILES --
    features_dict = read_features(features)
    lst_dict = read_lst(lst)
    X = np.asarray([features_dict[x] for x in lst_dict])

    # Exception cases
    if kmax is None and klist is None and path.isdir(output):
        raise Exception("The \"output\" is an existing directory while the system is waiting the path of a file.")

    if (kmax is not None or klist is not None) and path.isfile(output):
        raise Exception("The \"output\" is an existing file while the system is waiting the path of a directory.")

    # Mono value case
    if kmax is None and klist is None:
        fit_model(k, output)

    # Multi values case with kmax
    if kmax is not None:
        if not path.isdir(output):
            mkdir(output)
        Ks = range(k, kmax + 1)
        for i in Ks:
            fit_model(i, path.join(output, "clustering_" + str(i) + ".pkl"))

    # Second multi values case with klist
    if klist is not None:
        if not path.isdir(output):
            mkdir(output)
        for k in klist:
            k = int(k)
            fit_model(k, path.join(output, "clustering_" + str(k) + ".pkl"))

    print(json.dumps(json_content))


def extract_run(features, lst, model, modeltype, outfile):
    feats_dict = read_features(features)
    lst_dict = read_lst(lst)
    lst_keys = [key for key in lst_dict]
    feats = np.asarray([feats_dict[key] for key in lst_keys])

    module = CLUSTERING_METHODS[modeltype]
    module.load(model)
    Y_pred = module.predict(feats)
    with open(outfile, "w") as f:
        for i, key in enumerate(lst_keys):
            write_line(key, Y_pred[i], f)
    json_output = {
        "outfile": outfile
    }
    print(json.dumps(json_output))


if __name__ == "__main__":
    # Main parser
    parser = argparse.ArgumentParser(description="Clustering methods to apply")
    subparsers = parser.add_subparsers(title="action")

    # kmeans
    parser_kmeans = subparsers.add_parser(
        "kmeans", help="Compute clustering using k-means algorithm")

    parser_kmeans.add_argument("--features", required=True, type=str, help="Features file (works with list)")
    parser_kmeans.add_argument("--lst", required=True, type=str, help="List file (.lst)")
    parser_kmeans.add_argument("-k", default=2, type=int,
                               help="number of clusters to compute. It is kmin if kmax is specified.")
    parser_kmeans.add_argument("--kmax", default=None, type=int, help="if specified, k is kmin.")
    parser_kmeans.add_argument("--klist", nargs="+",
                               help="List of k values to test. As kmax, activate the multi values mod.")
    parser_kmeans.add_argument("--maxiter",
                               type=int,
                               default=300,
                               help="Max number of iteration before stoping if not converging")
    parser_kmeans.add_argument("--ninit",
                               type=int,
                               default=10,
                               help="Number of time the k-means algorithm will be run with different centroid seeds.")
    parser_kmeans.add_argument("--tol",
                               type=float,
                               default=0.0001,
                               help="Tolerance to finish of distance between centroids and their updates.")
    parser_kmeans.add_argument("--debug", action="store_true")
    parser_kmeans.add_argument("--output",
                               default=".kmeans",
                               help="output file if only k. Output directory if multiple kmax specified.")
    parser_kmeans.add_argument("--modeltype",
                                required=True,
                                choices=KMEANS_METHODS,
                                help="type of model for learning")
    parser_kmeans.set_defaults(which="kmeans")

    # measure
    parser_measure = subparsers.add_parser(
        "measure", help="compute the entropy")

    parser_measure.add_argument("--measure",
                                required=True,
                                nargs="+",
                                choices=[key for key in EVALUATION_METHODS],
                                help="...")
    parser_measure.add_argument("--features", required=True, type=str, help="...")
    parser_measure.add_argument("--lst", required=True, type=str, help="...")
    parser_measure.add_argument("--truelabels", required=True, type=str, help="...")
    parser_measure.add_argument("--model", required=True, type=str, help="...")
    parser_measure.add_argument("--modeltype",
                                required=True,
                                choices=[key for key in CLUSTERING_METHODS],
                                help="type of model for learning")
    parser_measure.set_defaults(which="measure")

    # disequilibrium
    parser_disequilibrium = subparsers.add_parser(
        "disequilibrium", help="...")

    parser_disequilibrium.add_argument("--features", required=True, type=str, help="...")
    parser_disequilibrium.add_argument("--lstrain", required=True, type=str, help="...")
    parser_disequilibrium.add_argument("--lstest", required=True, type=str, help="...")
    parser_disequilibrium.add_argument("--model", required=True, type=str, help="...")
    parser_disequilibrium.add_argument("--modeltype",
                                required=True,
                                choices=["kmeans", "2", "3"],
                                help="...")
    parser_disequilibrium.set_defaults(which="disequilibrium")

    # Extract
    parser_extract = subparsers.add_parser(
        "extract", help="extract cluster labels")

    parser_extract.add_argument("--features", required=True, type=str, help="...")
    parser_extract.add_argument("--lst", required=True, type=str, help="...")
    parser_extract.add_argument("--model", required=True, type=str, help="...")
    parser_extract.add_argument("--modeltype",
                                required=True,
                                choices=[key for key in CLUSTERING_METHODS],
                                help="type of model for learning")
    parser_extract.add_argument("--outfile", required=True, type=str, help="...")
    parser_extract.set_defaults(which="extract")

    # Parse
    args = parser.parse_args()

    # Run commands
    runner = SubCommandRunner({
        "kmeans": kmeans_run,
        "measure": measure_run,
        "disequilibrium": disequilibrium_run,
        "extract": extract_run
    })

    runner.run(args.which, args.__dict__, remove="which")