import os
import logging

import tensorflow as tf
from tensorflow.python.client import timeline
from tensorflow.python.training import training_util
from tensorflow.python.training import session_run_hook


class ProfileAtStepHook(session_run_hook.SessionRunHook):
  """Hook that requests stop at a specified step."""

  def __init__(self, at_step=None, checkpoint_dir=None, trace_level=tf.RunOptions.FULL_TRACE):
    self._at_step = at_step
    self._do_profile = False
    self._writer = tf.summary.FileWriter(checkpoint_dir)
    self._trace_level = trace_level

  def begin(self):
    self._global_step_tensor = tf.train.get_global_step()
    if self._global_step_tensor is None:
      raise RuntimeError("Global step should be created to use ProfileAtStepHook.")

  def before_run(self, run_context):  # pylint: disable=unused-argument
    if self._do_profile:
      options = tf.RunOptions(trace_level=self._trace_level)
    else:
      options = None
    return tf.train.SessionRunArgs(self._global_step_tensor, options=options)

  def after_run(self, run_context, run_values):
    global_step = run_values.results - 1
    if self._do_profile:
      self._do_profile = False
      self._writer.add_run_metadata(run_values.run_metadata,
                                    'trace_{}'.format(global_step), global_step)
      timeline_object = timeline.Timeline(run_values.run_metadata.step_stats)
      chrome_trace = timeline_object.generate_chrome_trace_format()
      chrome_trace_save_path = 'timeline_{}.json'.format(global_step)
      with open(chrome_trace_save_path, 'w') as f:
        f.write(chrome_trace)
      logging.info('Profile trace saved to {}'.format(chrome_trace_save_path))
    if global_step == self._at_step:
      self._do_profile = True
