# 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.
# ==============================================================================
"""TensorFlow monitoring APIs."""

import collections
import functools
import time

from tensorflow.core.framework import summary_pb2
from tensorflow.python import pywrap_tfe
from tensorflow.python.client import pywrap_tf_session
from tensorflow.python.framework import c_api_util
from tensorflow.python.util import compat
from tensorflow.python.util.tf_export import tf_export

_MetricMethod = collections.namedtuple('MetricMethod', 'create delete get_cell')
_counter_methods = [
    _MetricMethod(
        create=pywrap_tfe.TFE_MonitoringNewCounter0,
        delete=pywrap_tfe.TFE_MonitoringDeleteCounter0,
        get_cell=pywrap_tfe.TFE_MonitoringGetCellCounter0),
    _MetricMethod(
        create=pywrap_tfe.TFE_MonitoringNewCounter1,
        delete=pywrap_tfe.TFE_MonitoringDeleteCounter1,
        get_cell=pywrap_tfe.TFE_MonitoringGetCellCounter1),
    _MetricMethod(
        create=pywrap_tfe.TFE_MonitoringNewCounter2,
        delete=pywrap_tfe.TFE_MonitoringDeleteCounter2,
        get_cell=pywrap_tfe.TFE_MonitoringGetCellCounter2),
]
_int_gauge_methods = [
    _MetricMethod(
        create=pywrap_tfe.TFE_MonitoringNewIntGauge0,
        delete=pywrap_tfe.TFE_MonitoringDeleteIntGauge0,
        get_cell=pywrap_tfe.TFE_MonitoringGetCellIntGauge0),
    _MetricMethod(
        create=pywrap_tfe.TFE_MonitoringNewIntGauge1,
        delete=pywrap_tfe.TFE_MonitoringDeleteIntGauge1,
        get_cell=pywrap_tfe.TFE_MonitoringGetCellIntGauge1),
    _MetricMethod(
        create=pywrap_tfe.TFE_MonitoringNewIntGauge2,
        delete=pywrap_tfe.TFE_MonitoringDeleteIntGauge2,
        get_cell=pywrap_tfe.TFE_MonitoringGetCellIntGauge2),
]
_string_gauge_methods = [
    _MetricMethod(
        create=pywrap_tfe.TFE_MonitoringNewStringGauge0,
        delete=pywrap_tfe.TFE_MonitoringDeleteStringGauge0,
        get_cell=pywrap_tfe.TFE_MonitoringGetCellStringGauge0),
    _MetricMethod(
        create=pywrap_tfe.TFE_MonitoringNewStringGauge1,
        delete=pywrap_tfe.TFE_MonitoringDeleteStringGauge1,
        get_cell=pywrap_tfe.TFE_MonitoringGetCellStringGauge1),
    _MetricMethod(
        create=pywrap_tfe.TFE_MonitoringNewStringGauge2,
        delete=pywrap_tfe.TFE_MonitoringDeleteStringGauge2,
        get_cell=pywrap_tfe.TFE_MonitoringGetCellStringGauge2),
    _MetricMethod(
        create=pywrap_tfe.TFE_MonitoringNewStringGauge3,
        delete=pywrap_tfe.TFE_MonitoringDeleteStringGauge3,
        get_cell=pywrap_tfe.TFE_MonitoringGetCellStringGauge3),
    _MetricMethod(
        create=pywrap_tfe.TFE_MonitoringNewStringGauge4,
        delete=pywrap_tfe.TFE_MonitoringDeleteStringGauge4,
        get_cell=pywrap_tfe.TFE_MonitoringGetCellStringGauge4),
]
_bool_gauge_methods = [
    _MetricMethod(
        create=pywrap_tfe.TFE_MonitoringNewBoolGauge0,
        delete=pywrap_tfe.TFE_MonitoringDeleteBoolGauge0,
        get_cell=pywrap_tfe.TFE_MonitoringGetCellBoolGauge0),
    _MetricMethod(
        create=pywrap_tfe.TFE_MonitoringNewBoolGauge1,
        delete=pywrap_tfe.TFE_MonitoringDeleteBoolGauge1,
        get_cell=pywrap_tfe.TFE_MonitoringGetCellBoolGauge1),
    _MetricMethod(
        create=pywrap_tfe.TFE_MonitoringNewBoolGauge2,
        delete=pywrap_tfe.TFE_MonitoringDeleteBoolGauge2,
        get_cell=pywrap_tfe.TFE_MonitoringGetCellBoolGauge2),
]
_sampler_methods = [
    _MetricMethod(
        create=pywrap_tfe.TFE_MonitoringNewSampler0,
        delete=pywrap_tfe.TFE_MonitoringDeleteSampler0,
        get_cell=pywrap_tfe.TFE_MonitoringGetCellSampler0),
    _MetricMethod(
        create=pywrap_tfe.TFE_MonitoringNewSampler1,
        delete=pywrap_tfe.TFE_MonitoringDeleteSampler1,
        get_cell=pywrap_tfe.TFE_MonitoringGetCellSampler1),
    _MetricMethod(
        create=pywrap_tfe.TFE_MonitoringNewSampler2,
        delete=pywrap_tfe.TFE_MonitoringDeleteSampler2,
        get_cell=pywrap_tfe.TFE_MonitoringGetCellSampler2),
]


class Metric(object):
  """The base class of metric."""

  __slots__ = ["_metric", "_metric_name", "_metric_methods", "_label_length"]

  def __init__(self, metric_name, metric_methods, label_length, *args):
    """Creates a new metric.

    Args:
      metric_name: name of the metric class.
      metric_methods: list of swig metric methods.
      label_length: length of label args.
      *args: the arguments to call create method.
    """
    self._metric_name = metric_name
    self._metric_methods = metric_methods
    self._label_length = label_length

    if label_length >= len(self._metric_methods):
      raise ValueError('Cannot create {} metric with label >= {}'.format(
          self._metric_name, len(self._metric_methods)))

    self._metric = self._metric_methods[self._label_length].create(*args)

  def __del__(self):
    try:
      deleter = self._metric_methods[self._label_length].delete
      metric = self._metric
    except AttributeError:
      return

    if deleter is not None:
      deleter(metric)

  def get_cell(self, *labels):
    """Retrieves the cell."""
    if len(labels) != self._label_length:
      raise ValueError('The {} expects taking {} labels'.format(
          self._metric_name, self._label_length))
    return self._metric_methods[self._label_length].get_cell(
        self._metric, *labels)


class CounterCell(object):
  """CounterCell stores each value of a Counter."""

  __slots__ = ["_cell"]

  def __init__(self, cell):
    """Creates a new CounterCell.

    Args:
      cell: A c pointer of TFE_MonitoringCounterCell.
    """
    self._cell = cell

  def increase_by(self, value):
    """Atomically increments the value.

    Args:
      value: non-negative value.
    """
    pywrap_tfe.TFE_MonitoringCounterCellIncrementBy(self._cell, value)

  def value(self):
    """Retrieves the current value."""
    return pywrap_tfe.TFE_MonitoringCounterCellValue(self._cell)


class Counter(Metric):
  """A stateful class for updating a cumulative integer metric.

  This class encapsulates a set of values (or a single value for a label-less
  metric). Each value is identified by a tuple of labels. The class allows the
  user to increment each value.
  """

  __slots__ = []

  def __init__(self, name, description, *labels):
    """Creates a new Counter.

    Args:
      name: name of the new metric.
      description: description of the new metric.
      *labels: The label list of the new metric.
    """
    super(Counter, self).__init__('Counter', _counter_methods, len(labels),
                                  name, description, *labels)

  def get_cell(self, *labels):
    """Retrieves the cell."""
    return CounterCell(super(Counter, self).get_cell(*labels))


class IntGaugeCell(object):
  """A single integer value stored in an `IntGauge`."""

  __slots__ = ["_cell"]

  def __init__(self, cell):
    """Creates a new IntGaugeCell.

    Args:
      cell: A c pointer of TFE_MonitoringIntGaugeCell.
    """
    self._cell = cell

  def set(self, value):
    """Atomically set the value.

    Args:
      value: integer value.
    """
    pywrap_tfe.TFE_MonitoringIntGaugeCellSet(self._cell, value)

  def value(self):
    """Retrieves the current value."""
    return pywrap_tfe.TFE_MonitoringIntGaugeCellValue(self._cell)


class IntGauge(Metric):
  """A stateful class for updating a gauge-like integer metric.

  This class encapsulates a set of integer values (or a single value for a
  label-less metric). Each value is identified by a tuple of labels. The class
  allows the user to set each value.
  """

  __slots__ = []

  def __init__(self, name, description, *labels):
    """Creates a new IntGauge.

    Args:
      name: name of the new metric.
      description: description of the new metric.
      *labels: The label list of the new metric.
    """
    super(IntGauge, self).__init__('IntGauge', _int_gauge_methods, len(labels),
                                   name, description, *labels)

  def get_cell(self, *labels):
    """Retrieves the cell."""
    return IntGaugeCell(super(IntGauge, self).get_cell(*labels))


class StringGaugeCell(object):
  """A single string value stored in an `StringGauge`."""

  __slots__ = ["_cell"]

  def __init__(self, cell):
    """Creates a new StringGaugeCell.

    Args:
      cell: A c pointer of TFE_MonitoringStringGaugeCell.
    """
    self._cell = cell

  def set(self, value):
    """Atomically set the value.

    Args:
      value: string value.
    """
    pywrap_tfe.TFE_MonitoringStringGaugeCellSet(self._cell, value)

  def value(self):
    """Retrieves the current value."""
    with c_api_util.tf_buffer() as buffer_:
      pywrap_tfe.TFE_MonitoringStringGaugeCellValue(self._cell, buffer_)
      value = pywrap_tf_session.TF_GetBuffer(buffer_).decode('utf-8')
    return value


class StringGauge(Metric):
  """A stateful class for updating a gauge-like string metric.

  This class encapsulates a set of string values (or a single value for a
  label-less metric). Each value is identified by a tuple of labels. The class
  allows the user to set each value.
  """

  __slots__ = []

  def __init__(self, name, description, *labels):
    """Creates a new StringGauge.

    Args:
      name: name of the new metric.
      description: description of the new metric.
      *labels: The label list of the new metric.
    """
    super(StringGauge, self).__init__('StringGauge', _string_gauge_methods,
                                      len(labels), name, description, *labels)

  def get_cell(self, *labels):
    """Retrieves the cell."""
    return StringGaugeCell(super(StringGauge, self).get_cell(*labels))


class BoolGaugeCell(object):
  """A single boolean value stored in an `BoolGauge`."""

  __slots__ = ["_cell"]

  def __init__(self, cell):
    """Creates a new BoolGaugeCell.

    Args:
      cell: A c pointer of TFE_MonitoringBoolGaugeCell.
    """
    self._cell = cell

  def set(self, value):
    """Atomically set the value.

    Args:
      value: bool value.
    """
    pywrap_tfe.TFE_MonitoringBoolGaugeCellSet(self._cell, value)

  def value(self):
    """Retrieves the current value."""
    return pywrap_tfe.TFE_MonitoringBoolGaugeCellValue(self._cell)


@tf_export("__internal__.monitoring.BoolGauge", v1=[])
class BoolGauge(Metric):
  """A stateful class for updating a gauge-like bool metric.

  This class encapsulates a set of boolean values (or a single value for a
  label-less metric). Each value is identified by a tuple of labels. The class
  allows the user to set each value.
  """

  __slots__ = []

  def __init__(self, name, description, *labels):
    """Creates a new BoolGauge.

    Args:
      name: name of the new metric.
      description: description of the new metric.
      *labels: The label list of the new metric.
    """
    super(BoolGauge, self).__init__('BoolGauge', _bool_gauge_methods,
                                    len(labels), name, description, *labels)

  def get_cell(self, *labels):
    """Retrieves the cell."""
    return BoolGaugeCell(super(BoolGauge, self).get_cell(*labels))


class SamplerCell(object):
  """SamplerCell stores each value of a Sampler."""

  __slots__ = ["_cell"]

  def __init__(self, cell):
    """Creates a new SamplerCell.

    Args:
      cell: A c pointer of TFE_MonitoringSamplerCell.
    """
    self._cell = cell

  def add(self, value):
    """Atomically add a sample.

    Args:
      value: float value.
    """
    pywrap_tfe.TFE_MonitoringSamplerCellAdd(self._cell, value)

  def value(self):
    """Retrieves the current distribution of samples.

    Returns:
      A HistogramProto describing the distribution of samples.
    """
    with c_api_util.tf_buffer() as buffer_:
      pywrap_tfe.TFE_MonitoringSamplerCellValue(self._cell, buffer_)
      proto_data = pywrap_tf_session.TF_GetBuffer(buffer_)
    histogram_proto = summary_pb2.HistogramProto()
    histogram_proto.ParseFromString(compat.as_bytes(proto_data))
    return histogram_proto


class Buckets(object):
  """Bucketing strategies for the samplers."""

  __slots__ = ["buckets"]

  def __init__(self, buckets):
    """Creates a new Buckets.

    Args:
      buckets: A c pointer of TFE_MonitoringBuckets.
    """
    self.buckets = buckets

  def __del__(self):
    pywrap_tfe.TFE_MonitoringDeleteBuckets(self.buckets)


class ExponentialBuckets(Buckets):
  """Exponential bucketing strategy.

  Sets up buckets of the form:
      [-DBL_MAX, ..., scale * growth^i,
       scale * growth_factor^(i + 1), ..., DBL_MAX].
  """

  __slots__ = []

  def __init__(self, scale, growth_factor, bucket_count):
    """Creates a new exponential Buckets.

    Args:
      scale: float
      growth_factor: float
      bucket_count: integer
    """
    super(ExponentialBuckets, self).__init__(
        pywrap_tfe.TFE_MonitoringNewExponentialBuckets(scale, growth_factor,
                                                       bucket_count))


class Sampler(Metric):
  """A stateful class for updating a cumulative histogram metric.

  This class encapsulates a set of histograms (or a single histogram for a
  label-less metric) configured with a list of increasing bucket boundaries.
  Each histogram is identified by a tuple of labels. The class allows the
  user to add a sample to each histogram value.
  """

  __slots__ = []

  def __init__(self, name, buckets, description, *labels):
    """Creates a new Sampler.

    Args:
      name: name of the new metric.
      buckets: bucketing strategy of the new metric.
      description: description of the new metric.
      *labels: The label list of the new metric.
    """
    super(Sampler, self).__init__('Sampler', _sampler_methods, len(labels),
                                  name, buckets.buckets, description, *labels)

  def get_cell(self, *labels):
    """Retrieves the cell."""
    return SamplerCell(super(Sampler, self).get_cell(*labels))


# Keeping track of current MonitoredTimer sections to prevent repetitive
# counting.
MonitoredTimerSections = []


class MonitoredTimer(object):
  """A context manager to measure the walltime and increment a Counter cell."""

  __slots__ = [
      "cell",
      "t",
      "monitored_section_name",
      "_counting",
      "_avoid_repetitive_counting",
  ]

  def __init__(
      self, cell, monitored_section_name=None, avoid_repetitive_counting=False
  ):
    """Creates a new MonitoredTimer.

    Args:
      cell: the cell associated with the time metric that will be inremented.
      monitored_section_name: name of action being monitored here.
      avoid_repetitive_counting: when set to True, if already in a monitored
        timer section with the same monitored_section_name, skip counting.
    """
    self.cell = cell
    self.monitored_section_name = monitored_section_name
    self._avoid_repetitive_counting = avoid_repetitive_counting
    self._counting = True

  def __enter__(self):
    if (
        self._avoid_repetitive_counting
        and self.monitored_section_name
        and self.monitored_section_name in MonitoredTimerSections
    ):
      self._counting = False
      return self

    self.t = time.time()
    if self.monitored_section_name:
      MonitoredTimerSections.append(self.monitored_section_name)

    return self

  def __exit__(self, exception_type, exception_value, traceback):
    del exception_type, exception_value, traceback
    if self._counting:
      micro_seconds = (time.time() - self.t) * 1000000
      self.cell.increase_by(int(micro_seconds))
      if self.monitored_section_name:
        MonitoredTimerSections.remove(self.monitored_section_name)


def monitored_timer(cell):
  """A function decorator for adding MonitoredTimer support.

  Args:
    cell: the cell associated with the time metric that will be inremented.
  Returns:
    A decorator that measure the function runtime and increment the specified
    counter cell.
  """

  def actual_decorator(func):

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
      with MonitoredTimer(cell):
        return func(*args, **kwargs)

    return wrapper

  return actual_decorator