import argparse
  from os import path, mkdir
  from utils import SubCommandRunner

  from core.data import read_features, read_lst, read_labels

  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 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()

  }
  
  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, output: str, mahalanobis: str = 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.

      """

      # -- 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:
          print(f"Computing clustering with k={k}")

          model = CLUSTERING_METHODS["k-means"]
          if mahalanobis:
              print("Computing with mahalanobis distance")
              model = CLUSTERING_METHODS["k-means-mahalanobis"]
          model.fit(X, k)
          model.save(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:

              model = CLUSTERING_METHODS["k-means"]
              if mahalanobis:
                  model = CLUSTERING_METHODS["k-means-mahalanobis"]
              model.fit(X, i)
              model.save(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)

              model = CLUSTERING_METHODS["k-means"]
              if mahalanobis:
                  print("Computing with mahalanobis distance")
                  model = CLUSTERING_METHODS["k-means-mahalanobis"]
              model.fit(X, k)
              model.save(path.join(output, "clustering_" + str(k) + ".pkl"))

  
  
  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("--output",
                                 default=".kmeans",
                                 help="output file if only k. Output directory if multiple kmax specified.")
      parser_kmeans.add_argument("--mahalanobis", action="store_true")

      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("--model-type",
                                  required=True,
                                  choices=["kmeans", "2", "3"],
                                  help="...")

      parser_disequilibrium.set_defaults(which="disequilibrium")

      # Parse
      args = parser.parse_args()
  
      # Run commands
      runner = SubCommandRunner({

          "kmeans": kmeans_run,
          "measure": measure_run,
          "disequilibrium": disequilibrium_run

      })

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