# 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.
# ==============================================================================
"""Data structures and algorithms for numerics alert.

The alerts are generated when a Tensor's elements contain bad values,
including nan, -inf and +inf.
"""

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

import collections
import json
import re

import tensorflow as tf

from tensorflow.python import debug as tf_debug
from tensorboard.plugins.debugger import constants

# The following two namedtuples are the same except that
# 1) `timestamp` in `NumericsAlert` is the timestamp of a single alerting event,
#    while `first_timestamp` in `NumericsAlertReportRow` is the first (earliest)
#    timestamp of a se tof aggregated `NumericsAlert`s.
# 2) The counts in `NumericsAlert` are the counts of elements in the tensor,
#    while the event counts in `NumericsAlertReportRow` are counts of previous
#    `NumericsAlert` events of the corresponding categories.
NumericsAlert = collections.namedtuple(
    "NumericsAlert",
    [
        "device_name",
        "tensor_name",
        "timestamp",
        "nan_count",
        "neg_inf_count",
        "pos_inf_count",
    ],
)

NumericsAlertReportRow = collections.namedtuple(
    "NumericsAlertReportRow",
    [
        "device_name",
        "tensor_name",
        "first_timestamp",
        "nan_event_count",
        "neg_inf_event_count",
        "pos_inf_event_count",
    ],
)

# Used to reconstruct an _EventTracker from data read from disk. When updating
# this named tuple, make sure to keep the properties of _EventTracker in sync.
EventTrackerDescription = collections.namedtuple(
    "EventTrackerDescription",
    ["event_count", "first_timestamp", "last_timestamp"],
)

# Used to reconstruct NumericsAlertHistory.
HistoryTriplet = collections.namedtuple(
    "HistoryTriplet", ["device", "tensor", "jsonable_history"]
)


class _EventTracker(object):
    """Track events for a single category of values (NaN, -Inf, or +Inf)."""

    def __init__(self, event_count=0, first_timestamp=-1, last_timestamp=-1):
        """Tracks events for a single category of values.

        Args:
          event_count: The initial event count to use.
          first_timestamp: The timestamp of the first event with this value.
          last_timestamp: The timestamp of the last event with this category of
              values.
        """

        # When updating the properties of this class, make sure to keep
        # EventTrackerDescription in sync so that data can be written to and from
        # disk correctly.
        self.event_count = event_count
        self.first_timestamp = first_timestamp
        self.last_timestamp = last_timestamp

    def add(self, timestamp):
        if self.event_count == 0:
            self.first_timestamp = timestamp
            self.last_timestamp = timestamp
        else:
            if timestamp < self.first_timestamp:
                self.first_timestamp = timestamp
            if timestamp > self.last_timestamp:
                self.last_timestamp = timestamp
        self.event_count += 1

    def get_description(self):
        return EventTrackerDescription(
            self.event_count, self.first_timestamp, self.last_timestamp
        )


class NumericsAlertHistory(object):
    """History of numerics alerts."""

    def __init__(self, initialization_list=None):
        """Stores alert history for a single device, tensor pair.

        Args:
          initialization_list: (`list`) An optional list parsed from JSON read
            from disk. That entity is used to initialize this NumericsAlertHistory.
            Use the create_jsonable_object method of this class to create such an
            object.
        """
        if initialization_list:
            # Use data to initialize this NumericsAlertHistory.
            self._trackers = {}
            for (
                value_category_key,
                description_list,
            ) in initialization_list.items():
                description = EventTrackerDescription._make(description_list)
                self._trackers[value_category_key] = _EventTracker(
                    event_count=description.event_count,
                    first_timestamp=description.first_timestamp,
                    last_timestamp=description.last_timestamp,
                )
        else:
            # Start cleanly. With no prior data.
            self._trackers = {
                constants.NAN_KEY: _EventTracker(),
                constants.NEG_INF_KEY: _EventTracker(),
                constants.POS_INF_KEY: _EventTracker(),
            }

    def add(self, numerics_alert):
        if numerics_alert.nan_count:
            self._trackers[constants.NAN_KEY].add(numerics_alert.timestamp)
        if numerics_alert.neg_inf_count:
            self._trackers[constants.NEG_INF_KEY].add(numerics_alert.timestamp)
        if numerics_alert.pos_inf_count:
            self._trackers[constants.POS_INF_KEY].add(numerics_alert.timestamp)

    def first_timestamp(self, event_key=None):
        """Obtain the first timestamp.

        Args:
          event_key: the type key of the sought events (e.g., constants.NAN_KEY).
          If None, includes all event type keys.

        Returns:
          First (earliest) timestamp of all the events of the given type (or all
            event types if event_key is None).
        """
        if event_key is None:
            timestamps = [
                self._trackers[key].first_timestamp for key in self._trackers
            ]
            return min(timestamp for timestamp in timestamps if timestamp >= 0)
        else:
            return self._trackers[event_key].first_timestamp

    def last_timestamp(self, event_key=None):
        """Obtain the last timestamp.

        Args:
          event_key: the type key of the sought events (e.g., constants.NAN_KEY). If
          None, includes all event type keys.

        Returns:
          Last (latest) timestamp of all the events of the given type (or all
            event types if event_key is None).
        """
        if event_key is None:
            timestamps = [
                self._trackers[key].first_timestamp for key in self._trackers
            ]
            return max(timestamp for timestamp in timestamps if timestamp >= 0)
        else:
            return self._trackers[event_key].last_timestamp

    def event_count(self, event_key):
        """Obtain event count.

        Args:
          event_key: the type key of the sought events (e.g., constants.NAN_KEY). If
          None, includes all event type keys.

        Returns:
          If event_key is None, return the sum of the event_count of all event
          types. Otherwise, return the event_count of the specified event type.
        """
        return self._trackers[event_key].event_count

    def create_jsonable_history(self):
        """Creates a JSON-able representation of this object.

        Returns:
          A dictionary mapping key to EventTrackerDescription (which can be used to
          create event trackers).
        """
        return {
            value_category_key: tracker.get_description()
            for (value_category_key, tracker) in self._trackers.items()
        }


class NumericsAlertRegistry(object):
    """A registry for alerts on numerics (e.g., due to NaNs and infinities)."""

    def __init__(self, capacity=100, initialization_list=None):
        """Constructor.

        Args:
          capacity: (`int`) maximum number of device-tensor keys to store.
          initialization_list: (`list`) An optional list (parsed from JSON) that
            is used to initialize the data within this registry. Use the
            create_jsonable_registry method of NumericsAlertRegistry to create such
            a list.
        """
        self._capacity = capacity

        # A map from device-tensor key to a the TensorAlertRecord namedtuple.
        # The device-tensor key is a 2-tuple of the format (device_name, node_name).
        # E.g., ("/job:worker/replica:0/task:1/gpu:0", "cross_entropy/Log:0").
        self._data = dict()

        if initialization_list:
            # Initialize the alert registry using the data passed in. This might be
            # backup data used to restore the registry after say a borg pre-emption.
            for entry in initialization_list:
                triplet = HistoryTriplet._make(entry)
                self._data[
                    (triplet.device, triplet.tensor)
                ] = NumericsAlertHistory(
                    initialization_list=triplet.jsonable_history
                )

    def register(self, numerics_alert):
        """Register an alerting numeric event.

        Args:
          numerics_alert: An instance of `NumericsAlert`.
        """
        key = (numerics_alert.device_name, numerics_alert.tensor_name)
        if key in self._data:
            self._data[key].add(numerics_alert)
        else:
            if len(self._data) < self._capacity:
                history = NumericsAlertHistory()
                history.add(numerics_alert)
                self._data[key] = history

    def report(self, device_name_filter=None, tensor_name_filter=None):
        """Get a report of offending device/tensor names.

        The report includes information about the device name, tensor name, first
        (earliest) timestamp of the alerting events from the tensor, in addition to
        counts of nan, positive inf and negative inf events.

        Args:
          device_name_filter: regex filter for device name, or None (not filtered).
          tensor_name_filter: regex filter for tensor name, or None (not filtered).

        Returns:
          A list of NumericsAlertReportRow, sorted by the first_timestamp in
          asecnding order.
        """
        report = []
        for key in self._data:
            device_name, tensor_name = key
            history = self._data[key]
            report.append(
                NumericsAlertReportRow(
                    device_name=device_name,
                    tensor_name=tensor_name,
                    first_timestamp=history.first_timestamp(),
                    nan_event_count=history.event_count(constants.NAN_KEY),
                    neg_inf_event_count=history.event_count(
                        constants.NEG_INF_KEY
                    ),
                    pos_inf_event_count=history.event_count(
                        constants.POS_INF_KEY
                    ),
                )
            )

        if device_name_filter:
            device_name_pattern = re.compile(device_name_filter)
            report = [
                item
                for item in report
                if device_name_pattern.match(item.device_name)
            ]
        if tensor_name_filter:
            tensor_name_pattern = re.compile(tensor_name_filter)
            report = [
                item
                for item in report
                if tensor_name_pattern.match(item.tensor_name)
            ]
        # Sort results chronologically.
        return sorted(report, key=lambda x: x.first_timestamp)

    def create_jsonable_registry(self):
        """Creates a JSON-able representation of this object.

        Returns:
          A dictionary mapping (device, tensor name) to JSON-able object
          representations of NumericsAlertHistory.
        """
        # JSON does not support tuples as keys. Only strings. Therefore, we store
        # the device name, tensor name, and dictionary data within a 3-item list.
        return [
            HistoryTriplet(pair[0], pair[1], history.create_jsonable_history())
            for (pair, history) in self._data.items()
        ]


def extract_numerics_alert(event):
    """Determines whether a health pill event contains bad values.

    A bad value is one of NaN, -Inf, or +Inf.

    Args:
      event: (`Event`) A `tensorflow.Event` proto from `DebugNumericSummary`
        ops.

    Returns:
      An instance of `NumericsAlert`, if bad values are found.
      `None`, if no bad values are found.

    Raises:
      ValueError: if the event does not have the expected tag prefix or the
        debug op name is not the expected debug op name suffix.
    """
    value = event.summary.value[0]
    debugger_plugin_metadata_content = None
    if value.HasField("metadata"):
        plugin_data = value.metadata.plugin_data
        if plugin_data.plugin_name == constants.DEBUGGER_PLUGIN_NAME:
            debugger_plugin_metadata_content = plugin_data.content

    if not debugger_plugin_metadata_content:
        raise ValueError(
            "Event proto input lacks debugger plugin SummaryMetadata."
        )

    debugger_plugin_metadata_content = tf.compat.as_text(
        debugger_plugin_metadata_content
    )
    try:
        content_object = json.loads(debugger_plugin_metadata_content)
        device_name = content_object["device"]
    except (KeyError, ValueError) as e:
        raise ValueError(
            "Could not determine device from JSON string %r, %r"
            % (debugger_plugin_metadata_content, e)
        )

    debug_op_suffix = ":DebugNumericSummary"
    if not value.node_name.endswith(debug_op_suffix):
        raise ValueError(
            "Event proto input does not have the expected debug op suffix %s"
            % debug_op_suffix
        )
    tensor_name = value.node_name[: -len(debug_op_suffix)]

    elements = tf_debug.load_tensor_from_event(event)
    nan_count = elements[constants.NAN_NUMERIC_SUMMARY_OP_INDEX]
    neg_inf_count = elements[constants.NEG_INF_NUMERIC_SUMMARY_OP_INDEX]
    pos_inf_count = elements[constants.POS_INF_NUMERIC_SUMMARY_OP_INDEX]
    if nan_count > 0 or neg_inf_count > 0 or pos_inf_count > 0:
        return NumericsAlert(
            device_name,
            tensor_name,
            event.wall_time,
            nan_count,
            neg_inf_count,
            pos_inf_count,
        )
    return None
