<reponame>felixnext/tf-argonaut
'''
Defines generally used items in the dataset.

author: <NAME>
'''

import os
import shutil
import re
from time import time
from datetime import datetime
import numpy as np
import pandas as pd
import tensorflow as tf
from tensorflow.keras.utils import plot_model
import matplotlib.pyplot as plt

class Entry():
  def __init__(self, name, desc, params):
    self.name = name
    self.description = desc
    self.params = params
    self.time = time()

  def __str__(self):
    res = []
    # handle main data
    res.append('\033[1m{} - {}\033[0m'.format(datetime.fromtimestamp(self.time), self.name))
    if self.description is not None:
      res.append('> {}'.format(self.description.replace("\n", "\n> ")))

    # handle params
    if self.params is not None:
      res.append('> PARAMS:')
      # check types of params
      if isinstance(self.params, dict):
        for key in self.params.keys():
          res.append('> - {:5}: {}'.format(key, self.params[key]))
      elif isinstance(self.params, list) or isinstance(self.params, np.ndarray):
        for x in self.params:
          res.append('> - {}'.format(x))
      elif isinstance(self.params, pd.DataFrame):
        # NOTE: adjust
        res.append('> {}'.format(self.params))
      elif isinstance(self.params, str) or isinstance(self.params, int):
        res.append('> > {}'.format(self.params))
      elif isinstance(self.params, float):
        res.append('> > {:.4f}'.format(self.params))
      else:
        res.append('> > \033[93munkown datatype\033[0m - {}'.format(self.params))

    return "\n".join(res)

class Summary():
  '''Defines a summary class that stored relevant information to all changes done in the experiment.'''

  def __init__(self, name=None):
    self.entries = []
    # define the name
    if name is None:
      self.name = "Model_{}".format(int(time()))
    else:
      self.name = name

  def add(self, name, description, params=None, default_name=None):
    '''Adds a new entry to the summary.

    Args:
      title (str): Name of the entry
      description (str): description of the entry (type of operations performed)
      params (list): Object (either list, bunch or dict) with relevant parameters for the operation
    '''
    # simply add
    self.entries.append(Entry(name if name is not None else default_name, description, params))

  def add_str(self, str):
    '''Adds a single string message to the summary.'''
    self.add("message", str, None)

  def __str__(self):
    '''Shows the entire summary log.'''
    res = []
    res.append("\n--- Summary ---")

    for entry in self.entries:
      res.append(str(entry))

    res.append("---------------\n")

    return "\n".join(res)

  def generate_summary(self, folder, model=None, histories=None, horizontal=False, overwrite=True):
    '''Generates a summary and stores it at the given location.

    Args:
      folder (str): Folder where to store the summary
      model (tf.keras.Model): Model that is trained
      histories (tf.keras.History): dict of histories that contain information about individual training processes
      validation (tf.Results): Results from the model evaluation for every task at the end of the training process
      horizontal (bool): Defines if the model should be plotted horizontal
      overwrite (bool): Defines if old data should be cleared
    '''
    # generate folder structure
    #target = os.path.join(folder, self.name)
    target = os.path.join(folder, "results")
    if os.path.exists(target) and overwrite is True:
      shutil.rmtree(target)
    if not os.path.exists(target):
      os.makedirs(target)

    if model is not None:
      # store the image
      # FIXME: Appears not work (due to get_config of the actual models) - For now display in tensorboard...
      #plot_model(model, os.path.join(target, 'model.png'), show_shapes=True, expand_nested=True, dpi=128, rankdir='LR' if horizontal else 'TB')

      # generate summary
      with open(os.path.join(target, 'model_summary.txt'), 'w') as f:
        model.summary(print_fn=lambda s: f.write("{}\n".format(s)))

    # generate log
    with open(os.path.join(target, 'experiment.log'), 'w') as f:
      f.write(re.sub("\\033\\[[0-9]+m", "", str(self)))

    hist_data = []
    # generate training data and plots
    if histories is not None:
      for task_id in histories:
        # Plot training & validation accuracy values
        history = histories[task_id]
        hist = history.history
        used = {}
        # iterate through all keys in histogram
        for key in hist:
          # check if used
          if key in used:
            continue
          # plot data
          plt.plot(hist[key])
          hist_data.append(np.array([task_id, key] + list(hist[key])))
          used[key] = True
          has_val = False
          # check for validation data
          vkey = "val_{}".format(key)
          if vkey in hist:
            plt.plot(hist[vkey])
            hist_data.append(np.array([task_id, vkey] + list(hist[vkey])))
            has_val = True
            used[vkey] = True

          # configure plot
          plt.title('Model {}'.format(key))
          plt.ylabel(key)
          plt.xlabel('Epoch')
          if has_val is True:
            plt.legend(['Train', 'Test'], loc='upper left')
          plt.savefig(os.path.join(target, '{}_{}.png'.format(task_id, key)))
          # clear for next picture
          plt.clf()

        # generate pandas elements and store csv
        df = pd.DataFrame(np.array(hist_data))
        df.to_csv(os.path.join(target, "train_histories.csv"))

      # FEAT: generate combined graphs and csv (to integrate in HTML file)

    # generate index-html that can load and display all these information and copy it to target
    # FEAT: update the graphical output
    #shutil.copyfile(os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html"), os.path.join(target, "index.html"))

class TempCallback(tf.keras.callbacks.Callback):
  '''Callback that stores temporary results from the training process.
  Useful in case the system crashes - Can be used additionally to TensorBoard Callback.
  '''
  def __init__(self):
    pass
