"""Analyze results generated by train.py.

From qhbm-research/test_qaia run:
poetry run python qaia_chain/analyze.py
  --timestamp=<name of directory in results that you want to explore>
"""

import json
import os

from absl import app
from absl import flags
from absl import logging
#import altair
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import pandas
import tensorflow as tf
from tensorflow.python.summary.summary_iterator import summary_iterator

# Configuration settings
flags.DEFINE_string("root_dir", os.path.dirname(__file__),
                    "Where the training file is located.")
flags.DEFINE_string(
    "timestamp", None, "Log path for the results data.", required=True)

FLAGS = flags.FLAGS

P_SLICE = 25


def main(argv):
  print("\n\n\n\n\n\nBEGIN MAIN")
  del argv
  results_dir = os.path.join(FLAGS.root_dir, "results", FLAGS.timestamp)
  with open(os.path.join(results_dir, "config.json")) as f:
    loaded_config = json.load(f)
  logging.info(f"Parameter configuration:\n{loaded_config}")

  num_sites = loaded_config["dataset"]["num_sites"]
  bias_min = loaded_config["dataset"]["bias_min"]
  bias_max = loaded_config["dataset"]["bias_max"]
  bias_steps = loaded_config["dataset"]["bias_steps"]
  beta_min = loaded_config["dataset"]["beta_min"]
  beta_max = loaded_config["dataset"]["beta_max"]
  beta_steps = loaded_config["dataset"]["beta_steps"]
  digits = loaded_config["dataset"]["digits"]
  p_min = loaded_config["hparams"]["p_min"]
  p_max = loaded_config["hparams"]["p_max"]
  p_step = loaded_config["hparams"]["p_step"]

  # ========================================================================== #
  # Load training data.
  # ========================================================================== #

  metrics_dict = {
      "per_data_point": {
          "true_correlation": {},
          "true_long_correlation": {},
          "true_entropy": {},
      },
      "per_run": {
          "relative_entropy": {},
          "model_correlation": {},
          "model_long_correlation": {},
      },
      "per_epoch": {
          "loss": {},
          "fidelity": {},
      },
  }
  metrics_dir = os.path.join(results_dir, "metrics")
  print("Collecting all event files under {}".format(metrics_dir))
  for float_bias in tf.linspace(bias_min, bias_max,
                                bias_steps).numpy().tolist():
    bias = round(float_bias, digits)
    for m_type in metrics_dict.values():
      for m in m_type.values():
        m[bias] = {}
    for float_beta in tf.linspace(beta_min, beta_max,
                                  beta_steps).numpy().tolist():
      beta = round(float_beta, digits)
      print(f"Ingesting bias {bias}, beta {beta}")
      for m_type in metrics_dict.values():
        for m in m_type.values():
          m[bias][beta] = {}
      data_point_label = (
          f"bias_{str(bias).replace('.','p')}_beta_{str(beta).replace('.','p')}"
      )
      data_point_dir = os.path.join(metrics_dir, data_point_label)
      runs = os.listdir(data_point_dir)
      for payload in runs:
        model_dir = os.path.join(data_point_dir, payload)
        event_name = os.listdir(model_dir)[0]
        this_event_file = os.path.join(model_dir, event_name)
        if "data_point" in payload:
          for event in summary_iterator(this_event_file):
            for v in event.summary.value:
              for key in metrics_dict["per_data_point"]:
                if v.tag == key:
                  metrics_dict["per_data_point"][key][bias][
                      beta] = tf.make_ndarray(v.tensor)
        else:
          mystr = payload[:payload.find("_iteration")]
          p = int(mystr[2:])
          i = int(payload[-1])

          if p not in metrics_dict["per_epoch"]["loss"][bias][beta]:
            for key in ["per_run", "per_epoch"]:
              for metric in metrics_dict[key].values():
                metric[bias][beta][p] = {}
          if i not in metrics_dict["per_epoch"]["loss"][bias][beta][p]:
            for metric in metrics_dict["per_epoch"].values():
              metric[bias][beta][p][i] = {}

          for event in summary_iterator(this_event_file):
            for v in event.summary.value:
              for key in metrics_dict["per_run"]:
                if v.tag == key:
                  metrics_dict["per_run"][key][bias][beta][p][
                      i] = tf.make_ndarray(v.tensor)
              for key in metrics_dict["per_epoch"]:
                if v.tag == key:
                  metrics_dict["per_epoch"][key][bias][beta][p][i][
                      event.step] = tf.make_ndarray(v.tensor)

  # ========================================================================== #
  # Plot true correlations and entropy.
  # ========================================================================== #

  def plot_array(array, n_color_bins, color_bin_min, color_bin_max,
                 bias_downsample, plot_title, plot_name):
    bias_list = list(
        np.around(np.linspace(bias_min, bias_max, bias_steps), decimals=digits))
    beta_list = list(
        np.around(np.linspace(beta_min, beta_max, beta_steps), decimals=digits))
    fig, ax = plt.subplots()
    this_f_array = np.asarray(array).T
    cmap = matplotlib.cm.coolwarm
    levels = matplotlib.ticker.MaxNLocator(nbins=n_color_bins).tick_values(
        color_bin_min, color_bin_max)
    norm = matplotlib.colors.BoundaryNorm(levels, ncolors=cmap.N, clip=True)

    this_plot = ax.pcolor(this_f_array, cmap=cmap, norm=norm)
    this_colorbar = fig.colorbar(this_plot, ax=ax)
    ax.set_xlabel("bias")
    ax.set_ylabel("beta")
    ax.set_title(plot_title)
    plt.gcf().text(-0.5, 0.5, "test_metrics", fontsize=14)
    plt.xticks(
        np.arange(0, len(bias_list), bias_downsample) + 0.5,
        bias_list[::bias_downsample])
    plt.yticks(np.arange(len(beta_list)) + 0.5, beta_list)
    figpath = os.path.join(results_dir, plot_name)
    print("savefig at {}.png".format(figpath))
    plt.savefig(figpath)

  ######
  # Print values at this point
  #####
  print_this_bias = round(1.75, digits)
  print_this_beta = round(0.25, digits)

  # Extract metrics into plottable arrays
  metrics_array_dict = {}
  for outer_key in metrics_dict:
    metrics_array_dict[outer_key] = {}
    for inner_key in metrics_dict[outer_key]:
      metrics_array_dict[outer_key][inner_key] = np.zeros(
          (bias_steps, beta_steps))
  max_p_array = np.zeros((bias_steps, beta_steps))
  max_epochs_array = np.zeros((bias_steps, beta_steps))
  model_correlation_relative_difference_array = np.zeros(
      (bias_steps, beta_steps))
  model_long_correlation_relative_difference_array = np.zeros(
      (bias_steps, beta_steps))
  for r, float_bias in enumerate(
      tf.linspace(bias_min, bias_max, bias_steps).numpy().tolist()):
    bias = round(float_bias, digits)
    for c, float_beta in enumerate(
        tf.linspace(beta_min, beta_max, beta_steps).numpy().tolist()):
      beta = round(float_beta, digits)

      # Get the last p value
      last_p = 0
      for p in metrics_dict["per_epoch"]["loss"][bias][beta]:
        last_p = max(last_p, p)

      # Get the best i at the last p
      best_i = 0
      best_loss = 1000000000
      best_i_step = 0
      for i, i_dict in metrics_dict["per_epoch"]["loss"][bias][beta][
          last_p].items():
        max_step = 0
        for step in i_dict:
          max_step = max(max_step, step)
        last_loss = i_dict[max_step]
        if last_loss < best_loss:
          best_loss = last_loss
          best_i = i
          best_i_step = max_step

      max_epochs_array[r][c] = best_i_step
      max_p_array[r][c] = last_p

      for key in metrics_dict["per_data_point"]:
        if key == "true_correlation" and bias == print_this_bias and beta == print_this_beta:
          print(
              f"true_correlation: {metrics_dict['per_data_point'][key][bias][beta]}"
          )
        metrics_array_dict["per_data_point"][key][r][c] = metrics_dict[
            "per_data_point"][key][bias][beta]

      for key in metrics_dict["per_run"]:
        if key == "model_correlation" and bias == print_this_bias and beta == print_this_beta:
          print(
              f"model_correlation: {metrics_dict['per_run'][key][bias][beta][last_p][best_i]}"
          )
        metrics_array_dict["per_run"][key][r][c] = metrics_dict["per_run"][key][
            bias][beta][last_p][best_i]

      for key in metrics_dict["per_epoch"]:
        metrics_array_dict["per_epoch"][key][r][c] = metrics_dict["per_epoch"][
            key][bias][beta][last_p][best_i][best_i_step]

      model_correlation_relative_difference_array[r][c] = (
          metrics_array_dict["per_run"]["model_correlation"][r][c] -
          metrics_array_dict["per_data_point"]["true_correlation"][r][c]
      ) / metrics_array_dict["per_data_point"]["true_correlation"][r][c]
      model_long_correlation_relative_difference_array[r][c] = (
          metrics_array_dict["per_run"]["model_long_correlation"][r][c] -
          metrics_array_dict["per_data_point"]["true_long_correlation"][r][c]
      ) / metrics_array_dict["per_data_point"]["true_long_correlation"][r][c]

  plot_array(metrics_array_dict["per_data_point"]["true_correlation"], 20, 0.0,
             1.0, 1,
             "Average nearest-neighbor correlation for true thermal state",
             "true_correlation_plot")

  plot_array(metrics_array_dict["per_data_point"]["true_long_correlation"], 20,
             0.0, 1.0, 1, "Total correlation for true thermal state",
             "true_long_correlation_plot")

  plot_array(metrics_array_dict["per_data_point"]["true_entropy"], 20, 0.0,
             np.amax(metrics_array_dict["per_data_point"]["true_entropy"]), 1,
             "Entropy of true thermal state", "true_entropy_plot")

  plot_array(metrics_array_dict["per_run"]["relative_entropy"], 20, 0.0,
             np.amax(metrics_array_dict["per_run"]["relative_entropy"]), 1,
             "Relative entropy between final model and true thermal state",
             "relative_entropy_plot")

  plot_array(metrics_array_dict["per_run"]["model_correlation"], 20, 0.0, 1.0,
             1, "Average nearest-neighbor correlation for final model",
             "model_correlation_plot")

  plot_array(metrics_array_dict["per_run"]["model_long_correlation"], 20, 0.0,
             1.0, 1, "Total correlation for final model",
             "model_long_correlation_plot")

  plot_array(metrics_array_dict["per_epoch"]["loss"], 20,
             np.amin(metrics_array_dict["per_epoch"]["loss"]), 0.0, 1,
             "Final loss value for final model", "loss_plot")

  plot_array(metrics_array_dict["per_epoch"]["fidelity"], 20, 0.8, 1.0, 1,
             "Fidelity between final model and true thermal state",
             "fidelity_0p8_plot")

  plot_array(metrics_array_dict["per_epoch"]["fidelity"], 20, 0.9, 1.0, 1,
             "Fidelity between final model and true thermal state",
             "fidelity_0p9_plot")

  plot_array(metrics_array_dict["per_epoch"]["fidelity"], 20, 0.95, 1.0, 1,
             "Fidelity between final model and true thermal state",
             "fidelity_0p95_plot")

  plot_array(
      model_correlation_relative_difference_array, 20, -0.2, 0.2, 1,
      "Relative difference in correlations between final model and true state",
      "model_correlation_relative_difference_plot")

  plot_array(
      model_long_correlation_relative_difference_array, 20, -0.2, 0.2, 1,
      "Relative difference in total correlations between final model and true state",
      "model_long_correlation_relative_difference_plot")

  plot_array(max_p_array, len(range(p_min, p_max + 1, p_step)), p_min, p_max, 1,
             f"Layers required to converge", "max_p_plot")

  plot_array(max_epochs_array, 20, 0, np.amax(max_epochs_array), 1,
             f"Number of epochs to train final model", "max_epochs_plot")


if __name__ == "__main__":
  app.run(main)
