#!/usr/bin/python3
"""Read intermediate tensors generated by DumpAllTensors activity

Tools for reading/ parsing intermediate tensors.
"""

import argparse
import datetime
import numpy as np
import os
import pandas as pd
import tensorflow as tf
import json
import seaborn as sns
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import multiprocessing

from matplotlib.pylab import *
from tqdm import tqdm
# Enable large animation size
matplotlib.rcParams['animation.embed_limit'] = 2**128
# Enable tensor.numpy()
tf.compat.v1.enable_eager_execution()


############################ Helper Functions ############################
def reshape_to_matrix(array):
  """Reshape an array to a square matrix padded with np.nan at the end."""
  array = array.astype(float)
  width = math.ceil(len(array)**0.5)
  height = math.ceil(len(array)/ width)
  padded = np.pad(array=array,
                  pad_width=(0, width * height - len(array)),
                  mode='constant',
                  constant_values=np.nan)
  padded = padded.reshape(width, -1)
  return padded

def save_ani_to_video(ani, save_video_path, video_fps=5):
  Writer = animation.writers['ffmpeg']
  writer = Writer(fps=video_fps)
  #Save the movie
  ani.save(save_video_path, writer=writer, dpi=250)

def save_ani_to_html(ani, save_html_path):
  with open(save_html_path, 'w') as f:
    f.write(ani.to_jshtml())

############################ ModelMetaDataManager ############################
class ModelMetaDataManager(object):
  """Maps model name in nnapi to its graph architecture with lazy initialization.

  # Arguments
    android_build_top: the root directory of android source tree dump_dir:
    directory containing intermediate tensors pulled from device
    tflite_model_json_path: directory containing intermediate json output of
    model visualization tool (third_party/tensorflow/lite/tools:visualize) The
    json output path from the tool is always /tmp.
  """
  ############################ ModelMetaData ############################
  class ModelMetaData(object):
    """Store graph information of a model."""

    def __init__(self, tflite_model_json_path='/tmp'):
      with open(tflite_model_json_path, 'rb') as f:
        model_json = json.load(f)
      self.operators = model_json['subgraphs'][0]['operators']
      self.operator_codes = [item['builtin_code']\
                            for item in model_json['operator_codes']]
      self.output_meta_data = []
      self.load_output_meta_data()

    def load_output_meta_data(self):
      for operator in self.operators:
        data = {}
        # Each operator can only have one output
        assert(len(operator['outputs']) == 1)
        data['output_tensor_index'] = operator['outputs'][0]
        data['fused_activation_function'] = operator\
          .get('builtin_options', {})\
          .get('fused_activation_function', '')
        data['operator_code'] = self.operator_codes[operator['opcode_index']]
        self.output_meta_data.append(data)

  def __init__(self, android_build_top, dump_dir, tflite_model_json_dir='/tmp'):
    # key: nnapi model name, value: ModelMetaData
    self.models = dict()
    self.ANDROID_BUILD_TOP = android_build_top + "/"
    self.TFLITE_MODEL_JSON_DIR = tflite_model_json_dir + "/"
    self.DUMP_DIR = dump_dir + "/"
    self.nnapi_to_tflite_name = dict()
    self.tflite_to_nnapi_name = dict()
    self.__load_mobilenet_topk_aosp()
    self.model_names = sorted(os.listdir(dump_dir))

  def __load_mobilenet_topk_aosp(self):
    """Load information about tflite and nnapi model names."""
    json_path = '{}/{}'.format(
        self.ANDROID_BUILD_TOP,
        'test/mlts/models/assets/models_list/mobilenet_topk_aosp.json')
    with open(json_path, 'rb') as f:
      topk_aosp = json.load(f)
    for model in topk_aosp['models']:
      self.nnapi_to_tflite_name[model['name']] = model['modelFile']
      self.tflite_to_nnapi_name[model['modelFile']] = model['name']

  def __get_model_json_path(self, tflite_model_name):
    """Return tflite model jason path."""
    json_path = '{}/{}.json'.format(self.TFLITE_MODEL_JSON_DIR,
                                    tflite_model_name)
    return json_path

  def __load_model(self, tflite_model_name):
    """Initialize a ModelMetaData for this model."""
    model = self.ModelMetaData(self.__get_model_json_path(tflite_model_name))
    nnapi_model_name = self.model_name_tflite_to_nnapi(tflite_model_name)
    self.models[nnapi_model_name] = model

  def model_name_nnapi_to_tflite(self, nnapi_model_name):
    return self.nnapi_to_tflite_name.get(nnapi_model_name, nnapi_model_name)

  def model_name_tflite_to_nnapi(self, tflite_model_name):
    return self.tflite_to_nnapi_name.get(tflite_model_name, tflite_model_name)

  def get_model_meta_data(self, nnapi_model_name):
    """Retrieve the ModelMetaData with lazy initialization."""
    tflite_model_name = self.model_name_nnapi_to_tflite(nnapi_model_name)
    if nnapi_model_name not in self.models:
      self.__load_model(tflite_model_name)
    return self.models[nnapi_model_name]

  def generate_animation_html(self, output_file_path, model_names=None, heatmap=True):
    """Generate a html file containing the hist and heatmap animation of all models"""
    model_names = self.model_names if model_names is None else model_names
    html_data = ''
    for model_name in tqdm(model_names):
      print(datetime.datetime.now(), 'Processing', model_name)
      html_data += '<h3>{}</h3>'.format(model_name)
      model_data = ModelData(nnapi_model_name=model_name, manager=self)
      ani = model_data.gen_error_hist_animation()
      html_data += ani.to_jshtml()
      if heatmap:
        ani = model_data.gen_heatmap_animation()
        html_data += ani.to_jshtml()
    with open(output_file_path, 'w') as f:
      f.write(html_data)

  def generate_hist_animation_html(self, model_name):
    """Generate a html hist animation for a model, used for multiprocessing"""
    html_data = '<h3>{}</h3>'.format(model_name)
    model_data = ModelData(nnapi_model_name=model_name, manager=self)
    ani = model_data.gen_error_hist_animation()
    html_data += ani.to_jshtml()
    print(datetime.datetime.now(), "Done histogram for", model_name)
    self.return_dict[model_name + "-hist"] = html_data

  def generate_heatmap_animation_html(self, model_name):
    """Generate a html hist animation for a model, used for multiprocessing"""
    model_data = ModelData(nnapi_model_name=model_name, manager=self)
    ani = model_data.gen_heatmap_animation()
    html_data = ani.to_jshtml()
    print(datetime.datetime.now(), "Done heatmap for", model_name)
    self.return_dict[model_name + "-heatmap"] = html_data

  def multiprocessing_generate_animation_html(self, output_file_path,
                                       model_names=None, heatmap=True):
    """
    Generate a html file containing the hist and heatmap animation of all models
    with multiple process.
    """
    model_names = self.model_names if model_names is None else model_names
    manager = multiprocessing.Manager()
    self.return_dict = manager.dict()
    jobs = []
    for model_name in model_names:
      for target_func in [self.generate_hist_animation_html, self.generate_heatmap_animation_html]:
        p = multiprocessing.Process(target=target_func, args=(model_name,))
        jobs.append(p)
        p.start()
    # wait for completion
    for job in jobs:
      job.join()

    with open(output_file_path, 'w') as f:
      for model_name in model_names:
        f.write(self.return_dict[model_name + "-hist"])
        f.write(self.return_dict[model_name + "-heatmap"])


############################ TensorDict ############################
class TensorDict(dict):
  """A class to store cpu and nnapi tensors.

  # Arguments
    model_dir: directory containing intermediate tensors pulled from device
  """
  def __init__(self, model_dir):
    super().__init__()
    for useNNAPIDir in ['cpu', 'nnapi']:
      dir_path = model_dir + useNNAPIDir + "/"
      self[useNNAPIDir] = self.read_tensors_from_dir(dir_path)
    self.tensor_sanity_check()
    self.max_absolute_diff, self.min_absolute_diff = 0.0, 0.0
    self.max_relative_diff, self.min_relative_diff = 0.0, 0.0
    self.layers = sorted(self['cpu'].keys())
    self.calc_range()

  def bytes_to_numpy_tensor(self, file_path):
    """Load bytes outputed from DumpIntermediateTensor into numpy tensor."""
    if 'quant' in file_path or '8bit' in file_path:
      tensor_type = tf.int8
    else:
      tensor_type = tf.float32
    with open(file_path, mode='rb') as f:
      tensor_bytes = f.read()
      tensor = tf.decode_raw(input_bytes=tensor_bytes, out_type=tensor_type)
    if np.isnan(np.sum(tensor)):
      print('WARNING: tensor contains inf or nan')
    return tensor.numpy()

  def read_tensors_from_dir(self, dir_path):
    tensor_dict = dict()
    for tensor_file in os.listdir(dir_path):
      tensor = self.bytes_to_numpy_tensor(dir_path + tensor_file)
      tensor_dict[tensor_file] = tensor
    return tensor_dict

  def tensor_sanity_check(self):
    # Make sure the cpu tensors and nnapi tensors have the same outputs
    assert(set(self['cpu'].keys()) == set(self['nnapi'].keys()))
    print('Tensor sanity check passed')

  def calc_range(self):
    for layer in self.layers:
      diff = self.calc_diff(layer, relative_error=False)
      # update absolute max, min
      self.max_absolute_diff = max(self.max_absolute_diff, np.max(diff))
      self.min_absolute_diff = min(self.min_absolute_diff, np.min(diff))
      self.absolute_range = max(abs(self.min_absolute_diff),
                                abs(self.max_absolute_diff))

  def calc_diff(self, layer, relative_error=True):
    cpu_tensor = self['cpu'][layer]
    nnapi_tensor = self['nnapi'][layer]
    assert(cpu_tensor.shape == nnapi_tensor.shape)
    diff = cpu_tensor - nnapi_tensor
    if not relative_error:
      return diff
    diff = diff.astype(float)
    cpu_tensor = cpu_tensor.astype(float)
    # Devide by max so the relative error range is conveniently [-1, 1]
    max_cpu_nnapi_tensor = np.maximum(np.abs(cpu_tensor), np.abs(nnapi_tensor))
    relative_diff = np.divide(diff, max_cpu_nnapi_tensor, out=np.zeros_like(diff),
                              where=max_cpu_nnapi_tensor>0)
    return relative_diff

  def gen_tensor_diff_stats(self, relative_error=True, return_df=True, plot_diff=False):
    stats = []
    for layer in self.layers:
      diff = self.calc_diff(layer, relative_error)
      if plot_diff:
        self.plot_tensor_diff(diff)
      if return_df:
        stats.append({
          'layer': layer,
          'min': np.min(diff),
          'max': np.max(diff),
          'mean': np.mean(diff),
          'median': np.median(diff)
        })
    if return_df:
      return pd.DataFrame(stats)

  def plot_tensor_diff(diff):
    plt.figure()
    plt.hist(diff, bins=50, log=True)
    plt.plot()


############################ Model Data ############################
class ModelData(object):
  """A class to store all relevant inormation of a model.

  # Arguments
    nnapi_model_name: the name of the model
    manager: ModelMetaDataManager
  """
  def __init__(self, nnapi_model_name, manager, seq_limit=10):
    self.nnapi_model_name = nnapi_model_name
    self.manager = manager
    self.model_dir = self.get_target_model_dir(manager.DUMP_DIR,
                                               nnapi_model_name)
    self.tensor_dict = TensorDict(self.model_dir)
    self.mmd = manager.get_model_meta_data(nnapi_model_name)
    self.stats = self.tensor_dict.gen_tensor_diff_stats(relative_error=True,
                                                        return_df=True)
    self.layers = sorted(self.tensor_dict['cpu'].keys())
    self.cmap = sns.diverging_palette(220, 20, sep=20, as_cmap=True)
    self.seq_limit = seq_limit

  def get_target_model_dir(self, dump_dir, target_model_name):
    # Get the model directory path
    target_model_dir = dump_dir + target_model_name + "/"
    return target_model_dir

  def __sns_distplot(self, layer, bins, ax, range, relative_error):
    sns.distplot(self.tensor_dict.calc_diff(layer, relative_error=relative_error), bins=bins,
             hist_kws={"range":range, "log":True}, ax=ax, kde=False)

  def __plt_hist(self, layer, bins, ax, range, relative_error):
    ax.hist(self.tensor_dict.calc_diff(layer, relative_error=relative_error), bins=bins,
             range=range, log=True)

  def __get_layer_num(self):
    if self.seq_limit:
      return min(len(self.layers), len(self.mmd.output_meta_data) * self.seq_limit)
    return len(self.layers)

  def update_hist_data(self, i, fig, ax1, ax2, bins=50, plot_library='sns'):
    # Use % because there may be multiple testing samples
    operation = self.mmd.output_meta_data[i % len(self.mmd.output_meta_data)]['operator_code']
    layer = self.layers[i]
    subtitle = fig.suptitle('{} | {}\n{}'
                      .format(self.nnapi_model_name, layer, operation),
                      fontsize='x-large')
    for ax in (ax1, ax2):
      ax.clear()
    ax1.set_title('Relative Error')
    ax2.set_title('Absolute Error')
    absolute_range = self.tensor_dict.absolute_range

    # Determine underlying plotting library
    hist_func = self.__plt_hist if plot_library == 'matplotlib' else self.__sns_distplot
    hist_func(layer=layer, bins=bins, ax=ax1,
              range=(-1, 1), relative_error=True)
    hist_func(layer=layer, bins=bins, ax=ax2,
              range=(-absolute_range, absolute_range), relative_error=False)

  def gen_error_hist_animation(self, save_video_path=None, video_fps=10):
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12,9))
    ani = animation.FuncAnimation(fig, self.update_hist_data, self.__get_layer_num(),
                                  fargs=(fig, ax1, ax2),
                                  interval=200, repeat=False)
    # close before return to avoid dangling plot
    if save_video_path:
      save_ani_to_video(ani, save_video_path, video_fps)
    plt.close()
    return ani

  def __sns_heatmap(self, data, ax, cbar_ax, **kwargs):
    return sns.heatmap(data, cmap=self.cmap, cbar=True, ax=ax, cbar_ax=cbar_ax,
                     cbar_kws={"orientation": "horizontal"}, center=0, **kwargs)

  def update_heatmap_data(self, i, fig, axs):
    # Use % because there may be multiple testing samples
    operation = self.mmd.output_meta_data[i % len(self.mmd.output_meta_data)]['operator_code']
    layer = self.layers[i]
    subtitle = fig.suptitle('{} | {}\n{}\n'
                      .format(self.nnapi_model_name, layer, operation),
                      fontsize='x-large')
    # Clear all the axs and redraw
    # It's important to clear the colorbars as well to avoid duplicate colorbars
    for ax_tuple in axs:
      for ax in ax_tuple:
        ax.clear()
    axs[0][0].set_title('Diff')
    axs[0][1].set_title('CPU Tensor')
    axs[0][2].set_title('NNAPI Tensor')

    reshaped_diff = reshape_to_matrix(self.tensor_dict.calc_diff(layer, relative_error=False))
    reshaped_cpu = reshape_to_matrix(self.tensor_dict['cpu'][layer])
    reshaped_nnapi = reshape_to_matrix(self.tensor_dict['nnapi'][layer])
    absolute_range = self.tensor_dict.absolute_range
    g1 = self.__sns_heatmap(data=reshaped_diff, ax=axs[0][0], cbar_ax=axs[1][0],
                            vmin=-absolute_range, vmax=absolute_range)
    g2 = self.__sns_heatmap(data=reshaped_cpu, ax=axs[0][1], cbar_ax=axs[1][1])
    g3 = self.__sns_heatmap(data=reshaped_nnapi, ax=axs[0][2], cbar_ax=axs[1][2])

  def gen_heatmap_animation(self, save_video_path=None, video_fps=10, figsize=(13,6)):
    fig = plt.figure(constrained_layout=True, figsize=figsize)
    widths = [1, 1, 1]
    heights = [7, 1]
    spec = fig.add_gridspec(ncols=3, nrows=2, width_ratios=widths,
                            height_ratios=heights)
    axs = []
    for row in range(2):
      axs.append([])
      for col in range(3):
          axs[-1].append(fig.add_subplot(spec[row, col]))

    ani = animation.FuncAnimation(fig, self.update_heatmap_data, self.__get_layer_num(),
                                  fargs=(fig, axs),
                                  interval=200, repeat=False)
    if save_video_path:
      save_ani_to_video(ani, save_video_path, video_fps)
    # close before return to avoid dangling plot
    plt.close()
    return ani

  def plot_error_heatmap(self, target_layer, vmin=None, vmax=None):
    # Plot the diff heatmap for a given layer
    target_diff = self.tensor_dict['cpu'][target_layer] - \
                  self.tensor_dict['nnapi'][target_layer]
    reshaped_target_diff = reshape_to_matrix(target_diff)
    fig, ax = subplots(figsize=(9, 9))
    plt.title('Heat Map of Error between CPU and NNAPI')
    sns.heatmap(reshaped_target_diff,
                cmap=self.cmap,
                mask=np.isnan(reshaped_target_diff),
                center=0)
    plt.show()


############################ ModelDataComparison ############################
class ModelDataComparison:
  """A class to store and compare multiple ModelData.

  # Arguments
    model_data_list: a list of ModelData to be compared. Can be modified through
    the class variable.
  """
  def __init__(self, dump_dir_list, android_build_top, tflite_model_json_dir, model_name):
    self.dump_dir_list = dump_dir_list
    self.android_build_top = android_build_top
    self.tflite_model_json_dir = tflite_model_json_dir
    self.set_model_name(model_name)

  def set_model_name(self, model_name):
    # Set model to be compared and load/ reload all model data
    self.model_name = model_name
    self.__load_data()

  def __load_data(self):
    # Load all model data
    self.manager_list = []
    self.model_data_list = []
    for i, dump_dir in enumerate(self.dump_dir_list):
      manager = ModelMetaDataManager(self.android_build_top,
                                     dump_dir,
                                     tflite_model_json_dir=self.tflite_model_json_dir)
      model_data = ModelData(nnapi_model_name=self.model_name, manager=manager)
      self.manager_list.append(manager)
      self.model_data_list.append(model_data)
    self.sanity_check()

  def sanity_check(self):
    # Check
    # 1) if there are more than one model to be compared
    # 2) The data has the same intermediate layers
    assert(len(self.model_data_list) >= 1)
    sample_model_data = self.model_data_list[0]
    for i in range(1, len(self.model_data_list)):
      assert(set(sample_model_data.tensor_dict['cpu'].keys()) ==
             set(self.model_data_list[i].tensor_dict['nnapi'].keys()))
    print('Sanity Check Passed')
    self.layers = sample_model_data.layers
    self.mmd = sample_model_data.mmd

  def update_hist_comparison_data(self, i, fig, axs, bins=50):
    # Use % because there may be multiple testing samples
    sample_model_data = self.model_data_list[0]
    operation = self.mmd.output_meta_data[i % len(self.mmd.output_meta_data)]['operator_code']
    layer = self.layers[i]
    subtitle = fig.suptitle('{} | {}\n{}'
                      .format(sample_model_data.nnapi_model_name, layer, operation),
                      fontsize='x-large')
    for row in axs:
      for ax in row:
        ax.clear()

    hist_ax = axs[0][0]
    hist_ax.set_title('Diff Histogram')
    labels = [dump_dir.split('/')[-2] for dump_dir in self.dump_dir_list]
    cmap = sns.diverging_palette(220, 20, sep=20, as_cmap=True)
    for i, ax in enumerate(axs[1]):
      model_data = self.model_data_list[i]
      axs[1][i].set_title(labels[i])
      reshaped_diff = reshape_to_matrix(
        self.model_data_list[i].tensor_dict.calc_diff(layer, relative_error=False))
      sns.heatmap(reshaped_diff, cmap=cmap, cbar=True, ax=axs[1][i], cbar_ax=axs[2][i],
                     cbar_kws={"orientation": "horizontal"}, center=0)
      sns.distplot(model_data.tensor_dict.calc_diff(layer, relative_error=False), bins=bins,
             hist_kws={"log":True}, ax=hist_ax, kde=False)
    hist_ax.legend(labels)

  def gen_error_hist_comparison_animation(self, save_video_path=None, video_fps=10):
    layers = self.layers
    N = len(self.model_data_list)
    widths = [1] * N
    heights = [N * 0.7, 1, 0.2]
    fig = plt.figure(figsize=(5 * N, 4 * N))
    gs = fig.add_gridspec(3, N, width_ratios=widths, height_ratios=heights)
    axs = [[], [], []]
    axs[0].append(fig.add_subplot(gs[0, :]))
    for i in range(N):
      # heatmap
      axs[1].append(fig.add_subplot(gs[1, i]))
      # colorbar
      axs[2].append(fig.add_subplot(gs[2, i]))
    ani = animation.FuncAnimation(fig, self.update_hist_comparison_data, len(layers),
                                  fargs=(fig, axs),
                                  interval=200, repeat=False)
    if save_video_path:
      save_ani_to_video(ani, save_video_path, video_fps)
    # close before return to avoid dangling plot
    plt.close()
    return ani


############################ NumpyEncoder ############################
class NumpyEncoder(json.JSONEncoder):
  """Enable numpy array serilization in a dictionary.

  Usage:
    a = np.array([[1, 2, 3], [4, 5, 6]])
    json.dumps({'a': a, 'aa': [2, (2, 3, 4), a], 'bb': [2]}, cls=NumpyEncoder)
  """
  def default(self, obj):
      if isinstance(obj, np.ndarray):
          return obj.tolist()
      return json.JSONEncoder.default(self, obj)

def main(args):
  output_file_path = args.output_file_path if args.output_file_path else '/tmp/intermediate.html'

  manager = ModelMetaDataManager(
    args.android_build_top,
    args.dump_dir,
    tflite_model_json_dir='/tmp')

  if args.no_parallel or args.model_name:
    generation_func = manager.generate_animation_html
  else:
    generation_func = manager.multiprocessing_generate_animation_html

  if args.model_name:
    model_data = ModelData(nnapi_model_name=model_name, manager=manager)
    print(model_data.tensor_dict)
    generation_func(output_file_path=output_file_path, model_names=[args.model_name])
  else:
    generation_func(output_file_path=output_file_path)


if __name__ == '__main__':
  # Example usage
  # python tensor_utils.py ~/android/master/ ~/android/master/intermediate/ tts_float
  parser = argparse.ArgumentParser(description='Utilities for parsing intermediate tensors.')
  parser.add_argument('android_build_top', help='Your Android build top path.')
  parser.add_argument('dump_dir', help='The dump dir pulled from the device.')
  parser.add_argument('--model_name', help='NNAPI model name. Run all models if not specified.')
  parser.add_argument('--output_file_path', help='Animation HTML path.')
  parser.add_argument('--no_parallel', help='Run on a single process instead of multiple processes.')
  args = parser.parse_args()
  main(args)