# Code in this file is copied and adapted from
# https://github.com/openai/evolution-strategies-starter.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from collections import OrderedDict
import os
import sys
import time

import tensorflow as tf
from tensorflow.core.util import event_pb2
from tensorflow.python import pywrap_tensorflow
from tensorflow.python.util import compat

DEBUG = 10
INFO = 20
WARN = 30
ERROR = 40

DISABLED = 50


class TbWriter(object):
  """Based on SummaryWriter, but changed to allow for a different prefix."""
  def __init__(self, dir, prefix):
    self.dir = dir
    # Start at 1, because EvWriter automatically generates an object with
    # step = 0.
    self.step = 1
    self.evwriter = pywrap_tensorflow.EventsWriter(
        compat.as_bytes(os.path.join(dir, prefix)))

  def write_values(self, key2val):
    summary = tf.Summary(value=[tf.Summary.Value(tag=k, simple_value=float(v))
                                for (k, v) in key2val.items()])
    event = event_pb2.Event(wall_time=time.time(), summary=summary)
    event.step = self.step
    self.evwriter.WriteEvent(event)
    self.evwriter.Flush()
    self.step += 1

  def close(self):
    self.evwriter.Close()

# API


def start(dir):
  if _Logger.CURRENT is not _Logger.DEFAULT:
    sys.stderr.write("WARNING: You asked to start logging (dir=%s), but you "
                     "never stopped the previous logger (dir=%s)."
                     "\n" % (dir, _Logger.CURRENT.dir))
  _Logger.CURRENT = _Logger(dir=dir)


def stop():
  if _Logger.CURRENT is _Logger.DEFAULT:
    sys.stderr.write("WARNING: You asked to stop logging, but you never "
                     "started any previous logger."
                     "\n" % (dir, _Logger.CURRENT.dir))
    return
  _Logger.CURRENT.close()
  _Logger.CURRENT = _Logger.DEFAULT


def record_tabular(key, val):
  """Log a value of some diagnostic.

  Call this once for each diagnostic quantity, each iteration.
  """
  _Logger.CURRENT.record_tabular(key, val)


def dump_tabular():
  """Write all of the diagnostics from the current iteration."""
  _Logger.CURRENT.dump_tabular()


def log(*args, **kwargs):
  """Write the sequence of args, with no separators.

  This is written to the console and output files (if you've configured an
  output file).
  """
  level = kwargs['level'] if 'level' in kwargs else INFO
  _Logger.CURRENT.log(*args, level=level)


def debug(*args):
  log(*args, level=DEBUG)


def info(*args):
  log(*args, level=INFO)


def warn(*args):
  log(*args, level=WARN)


def error(*args):
  log(*args, level=ERROR)


def set_level(level):
  """
  Set logging threshold on current logger.
  """
  _Logger.CURRENT.set_level(level)


def get_dir():
  """
  Get directory that log files are being written to.
  will be None if there is no output directory (i.e., if you didn't call start)
  """
  return _Logger.CURRENT.get_dir()


def get_expt_dir():
  sys.stderr.write("get_expt_dir() is Deprecated. Switch to get_dir()\n")
  return get_dir()

# Backend


class _Logger(object):
  # A logger with no output files. (See right below class definition) so that
  # you can still log to the terminal without setting up any output files.
  DEFAULT = None
  # Current logger being used by the free functions above.
  CURRENT = None

  def __init__(self, dir=None):
    self.name2val = OrderedDict()  # Values this iteration.
    self.level = INFO
    self.dir = dir
    self.text_outputs = [sys.stdout]
    if dir is not None:
      os.makedirs(dir, exist_ok=True)
      self.text_outputs.append(open(os.path.join(dir, "log.txt"), "w"))
      self.tbwriter = TbWriter(dir=dir, prefix="events")
    else:
      self.tbwriter = None

  # Logging API, forwarded

  def record_tabular(self, key, val):
    self.name2val[key] = val

  def dump_tabular(self):
    # Create strings for printing.
    key2str = OrderedDict()
    for (key, val) in self.name2val.items():
      if hasattr(val, "__float__"):
        valstr = "%-8.3g" % val
      else:
        valstr = val
      key2str[self._truncate(key)] = self._truncate(valstr)
    keywidth = max(map(len, key2str.keys()))
    valwidth = max(map(len, key2str.values()))
    # Write to all text outputs
    self._write_text("-" * (keywidth + valwidth + 7), "\n")
    for (key, val) in key2str.items():
      self._write_text("| ", key, " " * (keywidth - len(key)), " | ", val,
                       " " * (valwidth - len(val)), " |\n")
    self._write_text("-" * (keywidth + valwidth + 7), "\n")
    for f in self.text_outputs:
      try:
        f.flush()
      except OSError:
        sys.stderr.write('Warning! OSError when flushing.\n')
    # Write to tensorboard
    if self.tbwriter is not None:
      self.tbwriter.write_values(self.name2val)
      self.name2val.clear()

  def log(self, *args, **kwargs):
    level = kwargs['level'] if 'level' in kwargs else INFO
    if self.level <= level:
      self._do_log(*args)

  # Configuration

  def set_level(self, level):
    self.level = level

  def get_dir(self):
    return self.dir

  def close(self):
    for f in self.text_outputs[1:]:
      f.close()
    if self.tbwriter:
      self.tbwriter.close()

  # Misc

  def _do_log(self, *args):
    self._write_text(*args + ('\n',))
    for f in self.text_outputs:
      try:
        f.flush()
      except OSError:
        print('Warning! OSError when flushing.')

  def _write_text(self, *strings):
    for f in self.text_outputs:
      for string in strings:
        f.write(string)

  def _truncate(self, s):
    if len(s) > 33:
      return s[:30] + "..."
    else:
      return s


_Logger.DEFAULT = _Logger()
_Logger.CURRENT = _Logger.DEFAULT