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 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()
  }
  
  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):
      """
  
      @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
      """

      # -- 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}")
          kmeans = KMeans(n_clusters=k, n_init=10, random_state=0).fit(X)
          preds = kmeans.predict(X)
          pickle.dump(kmeans, open(output, "wb"))
  
      # 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:
              print(f"Computing clustering with k={i}")
              kmeans = KMeans(n_clusters=i, n_init=10, random_state=0).fit(X)
              preds = kmeans.predict(X)
              pickle.dump(kmeans, open(path.join(output, "clustering_" + str(i) + ".pkl"), "wb"))
  
      # 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)
              print(f"Computing clustering with k={k}")
              kmeans = KMeans(n_clusters=k, n_init=10, random_state=0).fit(X)
              preds = kmeans.predict(X)
              pickle.dump(kmeans, open(path.join(output, "clustering_" + str(k) + ".pkl"), "wb"))
  
  
  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.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")