# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# ==============================================================================
"""Hooks for use with GTFlow Estimator."""

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

import os

from tensorflow.contrib.learn.python.learn import session_run_hook
from tensorflow.contrib.learn.python.learn.session_run_hook import SessionRunArgs
from tensorflow.core.framework.summary_pb2 import Summary
from tensorflow.python.framework import ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import state_ops
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.training import training_util
from tensorflow.python.training.summary_io import SummaryWriterCache


class FeatureImportanceSummarySaver(session_run_hook.SessionRunHook):
  """Hook to save feature importance summaries."""

  def __init__(self, model_dir, every_n_steps=1):
    """Create a FeatureImportanceSummarySaver Hook.

    This hook creates scalar summaries representing feature importance
    for each feature column during training.

    Args:
      model_dir: model base output directory.
      every_n_steps: frequency, in number of steps, for logging summaries.

    Raises:
      ValueError: If one of the arguments is invalid.
    """
    if model_dir is None:
      raise ValueError("model dir must be specified.")
    self._model_dir = model_dir
    self._every_n_steps = every_n_steps
    self._last_triggered_step = None

  def begin(self):
    self._global_step_tensor = training_util.get_global_step()
    if self._global_step_tensor is None:
      raise RuntimeError(
          "Global step should be created to use FeatureImportanceSummarySaver.")
    graph = ops.get_default_graph()
    self._feature_names_tensor = graph.get_tensor_by_name(
        "gbdt/feature_names:0")
    self._feature_usage_counts_tensor = graph.get_tensor_by_name(
        "gbdt/feature_usage_counts:0")
    self._feature_gains_tensor = graph.get_tensor_by_name(
        "gbdt/feature_gains:0")

  def before_run(self, run_context):
    del run_context  # Unused by feature importance summary saver hook.
    requests = {
        "global_step": self._global_step_tensor,
        "feature_names": self._feature_names_tensor,
        "feature_usage_counts": self._feature_usage_counts_tensor,
        "feature_gains": self._feature_gains_tensor
    }
    return SessionRunArgs(requests)

  def after_run(self, run_context, run_values):
    del run_context  # Unused by feature importance summary saver hook.

    # Read result tensors.
    global_step = run_values.results["global_step"]
    feature_names = run_values.results["feature_names"]
    feature_usage_counts = run_values.results["feature_usage_counts"]
    feature_gains = run_values.results["feature_gains"]

    # Ensure summaries are logged at desired frequency
    if (self._last_triggered_step is not None and
        global_step < self._last_triggered_step + self._every_n_steps):
      return

    # Validate tensors.
    if (len(feature_names) != len(feature_usage_counts) or
        len(feature_names) != len(feature_gains)):
      raise RuntimeError(
          "Feature names and importance measures have inconsistent lengths.")

    # Compute total usage.
    total_usage_count = 0.0
    for usage_count in feature_usage_counts:
      total_usage_count += usage_count
    usage_count_norm = 1.0 / total_usage_count if total_usage_count else 1.0

    # Compute total gain.
    total_gain = 0.0
    for gain in feature_gains:
      total_gain += gain
    gain_norm = 1.0 / total_gain if total_gain else 1.0

    # Output summary for each feature.
    self._last_triggered_step = global_step
    for (name, usage_count, gain) in zip(feature_names, feature_usage_counts,
                                         feature_gains):
      output_dir = os.path.join(self._model_dir, name.decode("utf-8"))
      summary_writer = SummaryWriterCache.get(output_dir)
      usage_count_summary = Summary(value=[
          Summary.Value(
              tag="feature_importance/usage_counts", simple_value=usage_count)
      ])
      usage_fraction_summary = Summary(value=[
          Summary.Value(
              tag="feature_importance/usage_fraction",
              simple_value=usage_count * usage_count_norm)
      ])
      summary_writer.add_summary(usage_count_summary, global_step)
      summary_writer.add_summary(usage_fraction_summary, global_step)
      gains_summary = Summary(value=[
          Summary.Value(tag="feature_importance/gains", simple_value=gain)
      ])
      gains_fraction_summary = Summary(value=[
          Summary.Value(
              tag="feature_importance/gains_fraction",
              simple_value=gain * gain_norm)
      ])
      summary_writer.add_summary(gains_summary, global_step)
      summary_writer.add_summary(gains_fraction_summary, global_step)


class FeedFnHook(session_run_hook.SessionRunHook):
  """Runs feed_fn and sets the feed_dict accordingly."""

  def __init__(self, feed_fn):
    self.feed_fn = feed_fn

  def before_run(self, run_context):
    del run_context  # unused by FeedFnHook.
    return session_run_hook.SessionRunArgs(fetches=None, feed_dict=self.feed_fn)


class StopAfterNTrees(session_run_hook.SessionRunHook):
  """Stop training after building N full trees."""

  def __init__(self, n, num_attempted_trees_tensor, num_finalized_trees_tensor,
               override_global_step_value=None):
    self._num_trees = n
    # num_attempted_trees_tensor and num_finalized_trees_tensor are both
    # tensors.
    self._num_attempted_trees_tensor = num_attempted_trees_tensor
    self._num_finalized_trees_tensor = num_finalized_trees_tensor
    self._override_global_step_value = override_global_step_value

  def begin(self):
    self._global_step_tensor = training_util.get_global_step()
    if self._global_step_tensor is None:
      raise RuntimeError("Global step should be created.")

    if self._override_global_step_value is not None:
      self._override_global_step_op = state_ops.assign(
          self._global_step_tensor, self._override_global_step_value)

  def before_run(self, run_context):
    del run_context  # unused by StopTrainingAfterNTrees.
    return session_run_hook.SessionRunArgs({
        "num_attempted_trees": self._num_attempted_trees_tensor,
        "num_finalized_trees": self._num_finalized_trees_tensor,
    })

  def after_run(self, run_context, run_values):
    num_attempted_trees = run_values.results["num_attempted_trees"]
    num_finalized_trees = run_values.results["num_finalized_trees"]
    assert num_attempted_trees is not None
    assert num_finalized_trees is not None
    # Stop when the required number of finalized trees is reached, or when we
    # try enough times to build a tree but keep failing.
    if (num_finalized_trees >= self._num_trees or
        num_attempted_trees > 2 * self._num_trees):
      logging.info("Requesting stop since we have reached %d trees.",
                   num_finalized_trees)
      if self._override_global_step_value is not None:
        logging.info("Overriding global steps value.")
        run_context.session.run(self._override_global_step_op)
      run_context.request_stop()


class SwitchTrainOp(session_run_hook.SessionRunHook):
  """Hook that switches the train op after specified number of steps.

  Hook that replaces the train op depending on the number of steps of training
  that have taken place. The first_train_op is used till train_steps steps
  are reached. Thereafter the second_train_op is used.
  """

  def __init__(self, first_train_op, train_steps, second_train_op):
    """Initializes a `SwitchTrainOp`."""
    self._first_train_op = first_train_op
    self._second_train_op = second_train_op
    self._train_steps = train_steps

  def _get_train_op_for_global_step(self, current_step):
    """Gets train_op for current global step."""
    if current_step < self._train_steps:
      return self._first_train_op
    return self._second_train_op

  def begin(self):
    self._global_step_tensor = training_util.get_global_step()
    self._current_train_op = control_flow_ops.no_op()
    if self._global_step_tensor is None:
      raise RuntimeError(
          "Global step should be created to use SwitchTrainOp.")

  def before_run(self, run_context):  # pylint: disable=unused-argument
    return session_run_hook.SessionRunArgs(
        {"global_step": self._global_step_tensor,
         "train_op": self._current_train_op})

  def after_run(self, run_context, run_values):
    self._current_train_op = self._get_train_op_for_global_step(
        run_values.results["global_step"])
