import argparse
  
  import os
  import core.data
  import math
  import numpy as np
  import scipy.stats
  import pickle
  import matplotlib.pyplot as plt
  import matplotlib.colors as mcolors

  from utils import SubCommandRunner

  from cycler import cycler

  def pred_distribution(predictions: str, labels: str, labelencoder: str, outdir: str):

      '''
      Distribution of the prediction.

      For each label, we plot the distribution of the class predicted.
      For example, for each character, we plot the distribution of the characters predicted.
      Another example, for each speaker, we plot the distribution of the characters predicted.
  
      '''

      predictions = core.data.read_id_values(args.predictions, float)
      labels = core.data.read_labels(args.labels)
  
      le = None
      with open(args.labelencoder, "rb") as f:
          le = pickle.load(f)
      stats = {}
  
      print("PREDICTIONS ---------------------------")
      for id_, predictions_ in predictions.items():
          label = labels[id_][0]
          if label not in stats:
              stats[label] = {
                  "nb_utt": 1,
                  "predictions": np.expand_dims(predictions_, axis=0)
              }
          else:
              stats[label]["nb_utt"] = stats[label]["nb_utt"] + 1
              stats[label]["predictions"] = np.append(stats[label]["predictions"], np.expand_dims(predictions_, axis=0), axis=0)

  
      colors = [
          "darkorange",
          "red",
          "blue"
      ]
      custom_cycler = (cycler(color=list(mcolors.TABLEAU_COLORS)) *
          cycler(linestyle=['-', '--', '-.']))

      print("CALCULATING ---------------------------")

      
      for label, stats_ in stats.items():
  
          plt.gca().set_prop_cycle(custom_cycler)
          stats_mean = np.mean(stats_["predictions"], axis=0)
          stats_std = np.std(stats_["predictions"], axis=0)

          #print(label)
          #print(stats_mean)
          #print(stats_std)
          kwargs = dict(alpha=0.5)
          
          for i in range(stats_["predictions"].shape[1]):
              label_str = le.inverse_transform([i])[0]
              #plt.hist(stats_["predictions"][:, i], bins=10, label=label_str, **kwargs)
              mu = stats_mean[i]
              variance = stats_std[i] * stats_std[i]
              sigma = stats_std[i]

              print(f"{i}: mu {mu}, var {variance}, sigma {sigma}")
  
              #x_values = np.arange(-1, 5, 0.1)
  
              #y_values = scipy.stats.norm(mu, variance)
              #y = scipy.stats.norm.pdf(x,mean,std)
  
              #plt.plot(x_values, y_values.pdf(x_values,))

              #x, step = np.linspace(mu - 3*sigma, mu + 3*sigma, 1000, retstep=True)
              x = np.linspace(0, 1, 1000)
              #x = np.linspace(mu - 3*sigma, mu + 3*sigma, 1000)
              #x, step = np.linspace(0, 1, 1000, retstep=True)
              
              P = scipy.stats.norm.cdf(x, mu, sigma)
              #print(step)
              plt.plot(x, P, label=label_str, **kwargs)
              #plt.savefig("simple_gaussian.pdf")
              
          plt.legend()
          plt.savefig(os.path.join(args.outdir, f"{label}_prediction_cdf.pdf"))
          plt.clf()
      

      print("Decisions")

  def pred_distribution_wt_sel(predictions: str, n: int, labels: str, labelencoder: str, outdir: str):
  
      '''
      Distribution of the predictions with selection process.
  
      1) For each dimension, select the n individus with the maximum values for the focused dimension.
      We name S_i the set of n selected individus for the dimension i.
      2) For each subset S_i, we plot the distribution of each dimension.
      '''
  
      le = None
      with open(args.labelencoder, "rb") as f:
          le = pickle.load(f)

  
      keys_preds, matrix_preds = core.data.read_features_with_matrix(predictions)

  
      colors = [
          "darkorange",
          "red",
          "blue"
      ]
      custom_cycler = (cycler(color=list(mcolors.TABLEAU_COLORS)) *
          cycler(linestyle=['-', '--', '-.']))
  
      kwargs = dict(alpha=0.5)
  
      stats_of = open(os.path.join(args.outdir, f"stats.txt"), "w")

      for j in range(matrix_preds.shape[1]):

  
          label_focused = le.inverse_transform([j])[0]

          indices = (-matrix_preds[:, j]).argsort()[:n]

  
          print(f"LABEL: {label_focused}", file=stats_of)
          print(f"INDICE: {j}", file=stats_of)
          print("indices", file=stats_of)
          print(indices, file=stats_of)
          print("Best values", file=stats_of)
          print(matrix_preds[indices, j], file=stats_of)
          print("All dimensions of best values", file=stats_of)
          print(matrix_preds[indices], file=stats_of)
  
          # Use it to build a plot.
          pred_ = matrix_preds[indices]
          stats_mean = np.mean(pred_, axis=0)
          stats_std = np.std(pred_, axis=0)
          for i in range(matrix_preds.shape[1]):
              label_str = le.inverse_transform([i])[0]
              mu = stats_mean[i]
              variance = stats_std[i] * stats_std[i]
              sigma = stats_std[i]
  
              print(f"{i}: mu {mu}, var {variance}, sigma {sigma}")
  
              x = np.linspace(0, 1, 1000)
  
              P = scipy.stats.norm.cdf(x, mu, sigma)
              plt.plot(x, P, label=label_str, **kwargs)
  
          plt.legend()
          plt.savefig(os.path.join(args.outdir, f"{label_focused}_prediction_cdf.pdf"))
          plt.clf()
      stats_of.close()

      pass

  def utt2dur(utt2dur: str, labels: str):
      if labels == None:
          pass
      else:
          pass
  
      durations = []
      with open(utt2dur, "r") as f:
          for line in f:
              splited = line.replace("
  ", "").split(" ")
              durations.append(float(splited[1]))
      
      durations = np.asarray(durations, dtype=float)
      print(durations.shape)
      mean = np.mean(durations)
      std = np.std(durations)

      print(f"mean: {mean}")
      print(f"std: {std}")
  
  
  if __name__ == "__main__":
  
      # Parser
      parser = argparse.ArgumentParser(description="Statistics")
      subparsers = parser.add_subparsers(title="actions")
  
      # pred-distribution
      parser_pred_dist = subparsers.add_parser("pred-distribution", help="plot distributions of prediction through labels")
      parser_pred_dist.add_argument("--predictions", type=str, help="prediction file", required=True)
      parser_pred_dist.add_argument("--labels", type=str, help="label file", required=True)
      parser_pred_dist.add_argument("--labelencoder", type=str, help="label encode pickle file", required=True)
      parser_pred_dist.add_argument("--outdir", type=str, help="output file", required=True)
      parser_pred_dist.set_defaults(which="pred_distribution")

      # pred-distribution-with-selection
      parser_pred_dist_wt_sel = subparsers.add_parser("pred-distribution-with-selection", help="plot distributions of prediction through labels with a selection of the n best records by column/class prediction.")
      parser_pred_dist_wt_sel.add_argument("--predictions", type=str, help="prediction file", required=True)

      parser_pred_dist_wt_sel.add_argument("-n", type=int, help="Number of maximum selected for each prediction y_i.", required=True)

      parser_pred_dist_wt_sel.add_argument("--labels", type=str, help="label file", required=True)
      parser_pred_dist_wt_sel.add_argument("--labelencoder", type=str, help="label encode pickle file", required=True)
      parser_pred_dist_wt_sel.add_argument("--outdir", type=str, help="output file", required=True)
      parser_pred_dist_wt_sel.set_defaults(which="pred_distribution_with_selection")

      # duration-stats
      parser_utt2dur = subparsers.add_parser("utt2dur", help="distribution of utt2dur")
      parser_utt2dur.add_argument("--utt2dur", type=str, help="utt2dur file", required=True)
      parser_utt2dur.add_argument("--labels", type=str, default=None, help="labels file")
      parser_utt2dur.set_defaults(which="utt2dur")
  
      # Parse
      args = parser.parse_args()
  
      # Run commands
      runner = SubCommandRunner({

          "pred_distribution": pred_distribution,
          "pred_distribution_with_selection": pred_distribution_wt_sel,

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