# coding=utf-8
# Copyright 2022 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# python3
# pylint: disable=line-too-long
r"""Trains a model on an eviction trace.

Example Usage:

  python3 -m cache_replacement.policy_learning.cache_model.main \
    --experiment_base_dir=/tmp \
    --experiment_name=sample_model_llc \
    --cache_configs=cache_replacement/policy_learning/cache/configs/default.json \
    --model_bindings="loss=[\"ndcg\", \"reuse_dist\"]" \
    --train_memtrace=cache_replacement/policy_learning/cache/traces/sample_trace.csv \
    --valid_memtrace=cache_replacement/policy_learning/cache/traces/sample_trace.csv

Trains full model with all additions on trained and validated on the sample trace.
"""
# pylint: enable=line-too-long

import io
import os
from absl import app
from absl import flags
from absl import logging
from baselines.common import schedules
import numpy as np
import prettytable
import tensorflow.compat.v1 as tf
import torch
from torch import optim
import tqdm
from cache_replacement.policy_learning.cache import cache as cache_mod
from cache_replacement.policy_learning.cache import evict_trace
from cache_replacement.policy_learning.cache import eviction_policy
from cache_replacement.policy_learning.cache import memtrace
from cache_replacement.policy_learning.cache_model import eviction_policy as model_eviction_policy
from cache_replacement.policy_learning.cache_model import metric
from cache_replacement.policy_learning.cache_model import model
from cache_replacement.policy_learning.cache_model import utils
from cache_replacement.policy_learning.common import config as cfg
from cache_replacement.policy_learning.common import utils as common_utils


FLAGS = flags.FLAGS
flags.DEFINE_string(
    "train_memtrace",
    "cache_replacement/policy_learning/cache/traces/omnetpp_train.csv",
    "Path to the training memory trace.")
flags.DEFINE_string(
    "valid_memtrace",
    "cache_replacement/policy_learning/cache/traces/omnetpp_valid.csv",
    "Path to the validation memory trace.")
flags.DEFINE_string(
    "experiment_base_dir", "/tmp/experiments",
    "Base directory to store all experiments in. Should not frequently change.")
flags.DEFINE_string(
    "experiment_name", "unnamed",
    "All data related to this experiment is written to"
    " experiment_base_dir/experiment_name.")
flags.DEFINE_integer("batch_size", 32, "Size of model input batches.")
flags.DEFINE_integer(
    "total_steps", int(1e6), "Number of training steps to take.")
flags.DEFINE_integer("tb_freq", 100, "Steps between logging to tensorboard.")
flags.DEFINE_integer(
    "small_eval_size", 30000,
    "Number of examples to evaluate on in small evaluations.")
flags.DEFINE_integer(
    "small_eval_freq", 4000,
    "Steps between evaluating on small_eval_size examples of validation data.")
flags.DEFINE_integer(
    "full_eval_freq", 48000,
    "Steps between evaluating on ALL of validation data.")
flags.DEFINE_integer(
    "save_freq", 20000, "Steps between saving model checkpoints.")
flags.DEFINE_integer(
    "collection_multiplier", 5,
    ("This times more training data is collected than updated on to avoid"
     " temporal correlations."))
flags.DEFINE_bool(
    "force_overwrite", False,
    ("If true, overwrites directory at "
     " experiment_base_dir/experiment_name if it exists."))
flags.DEFINE_multi_string(
    "model_configs",
    ["cache_replacement/policy_learning/cache_model/configs/default.json"],  # pylint: disable=line-too-long
    ("List of config filenames merged front to back for the model. "
     "See config.Config.merge for merging behavior."))
flags.DEFINE_multi_string(
    "model_bindings", [], "bindings to override model config.")
flags.DEFINE_multi_string(
    "cache_configs",
    ["cache_replacement/policy_learning/cache/configs/default.json"],  # pylint: disable=line-too-long
    "List of config filenames merged front to back for the cache.")
flags.DEFINE_multi_string(
    "cache_bindings", [], "bindings to override cache config.")
flags.DEFINE_multi_string(
    "dagger_schedule_configs",
    ["cache_replacement/policy_learning/cache_model/configs/schedule/linear.json"],  # pylint: disable=line-too-long
    "List of config filenames for the DAgger schedule.")
flags.DEFINE_multi_string(
    "dagger_schedule_bindings", [], "bindings to override DAgger config.")
flags.DEFINE_enum(
    "oracle_eviction_policy", "belady", ["lru", "belady"],
    "Which eviction policy to learn.")
flags.DEFINE_integer(
    "seed", 0, "Determines which random seed to use.")


def schedule_from_config(config):
  if config.get("type") == "linear":
    return schedules.LinearSchedule(
        config.get("num_steps"), config.get("final"), config.get("initial"))
  elif config.get("type") == "constant":
    return schedules.ConstantSchedule(config.get("value"))
  else:
    raise ValueError("U
nsupported schedule type: {}".format(config.get("type")))


def evaluate(policy_model, data, step, descriptor, tb_writer, log_dir, k=5):
  """Computes metrics about the model on the data and logs them to tensorboard.

  Args:
    policy_model (EvictionPolicyModel): model to evaluate with.
    data (list[EvictionEntry]): consecutive eviction entries.
    step (int): the current step in training. Affects how things are written to
      tensorboard and where predictions are logged.
    descriptor (str): des
cription of the data. Predictions are written to
      log_dir/descriptor-step.txt and metrics are logged to tensorboard with
      descriptor as the tb prefix.
    tb_writer (FileWriter): all tensorboard metrics are written here.
    log_dir (str): the directory for writing txt file predictions.
    k (int): computes top-1 to top-k success rates.

  Returns:
    metrics (list[CacheEvictionMetric]): metrics tracked during evaluation.
  """
  def pretty_print(entry, probs, attention, pred_reuse_distances):
    """Returns a human-readable string of the entry, probs, and attention.

    Args:
      entry (EvictionEntry): entry to print
      probs (torch.FloatTensor): probs of evicting entry's eviction candidates.
      attention (list[torch.FloatTensor, CacheAccess]): pairs of attention
        weight on past EvictionEntry sorted from most distant to most recent.
      pred_reuse_distances (torch.FloatTensor): the predicted reuse distance of
        each cache line in the entry of shape (num_cache_lines).

    Returns:
      string: error analysis string for probs and attention on the entry.
    """
    cache_access = entry.cache_access
    eviction_decision = entry.eviction_decision

    s = ["PC: {}\nAddress: {}\nEvict: {}\nCache lines:\n"
         .format(hex(cache_access.pc), hex(cache_access.address),
                 eviction_decision.evict)]

    _, true_ranks = probs.sort(descending=True)
    true_rank_to_pred_rank = dict(
        zip(true_ranks.cpu().data.numpy(), range(len(true_ranks))))
    headers = ["true rank", "pc", "address", "pred rank", "prob",
               "oracle score", "pred reuse distance", "rank correct?",
               "in history?"]
    cache_lines_table = prettytable.PrettyTable(headers)
    for i, (line, prob, pred_reuse) in enumerate(
        zip(cache_access.cache_lines, probs, pred_reuse_distances)):
      cand, pc = line
      pred_rank = true_rank_to_pred_rank[i]
      success = "SUCCESS" if pred_rank == i else "FAILURE"

      present_in_history = any(
          cand == prev_access.address for _, prev_access in attention)
      present = "PRESENT" if present_in_history else "ABSENT"
      cache_lines_table.add_row(
          [i, hex(pc), hex(cand), pred_rank, "{:.2f}".format(prob),
           "{:.2f}".for