# Copyright 2015 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.
# ==============================================================================

"""Logic for TensorBoard inspector to help humans investigate event files.

Example usages:
tensorboard --inspect --event_file=myevents.out
tensorboard --inspect --event_file=myevents.out --tag=loss
tensorboard --inspect --logdir=mylogdir
tensorboard --inspect --logdir=mylogdir --tag=loss


This script runs over a logdir and creates an InspectionUnit for every
subdirectory with event files. If running over an event file, it creates only
one InspectionUnit. One block of output is printed to console for each
InspectionUnit.

The primary content of an InspectionUnit is the dict field_to_obs that maps
fields (e.g. "scalar", "histogram", "session_log:start", etc.) to a list of
Observations for the field. Observations correspond one-to-one with Events in an
event file but contain less information because they only store what is
necessary to generate the final console output.

The final output is rendered to console by applying some aggregating function
to the lists of Observations. Different functions are applied depending on the
type of field. For instance, for "scalar" fields, the inspector shows aggregate
statistics. For other fields like "session_log:start", all observed steps are
printed in order to aid debugging.


[1] Query a logdir or an event file for its logged tags and summary statistics
using --logdir or --event_file.

[[event_file]] contains these tags:
histograms
   binary/Sign/Activations
   binary/nn_tanh/act/Activations
   binary/nn_tanh/biases
   binary/nn_tanh/biases:gradient
   binary/nn_tanh/weights
   binary/nn_tanh/weights:gradient
images
   input_images/image/0
   input_images/image/1
   input_images/image/2
scalars
   Learning Rate
   Total Cost
   Total Cost (raw)

Debug output aggregated over all tags:
graph
   first_step           0
   last_step            0
   max_step             0
   min_step             0
   num_steps            1
   outoforder_steps     []
histograms
   first_step           491
   last_step            659823
   max_step             659823
   min_step             491
   num_steps            993
   outoforder_steps     []
images -
scalars
   first_step           0
   last_step            659823
   max_step             659823
   min_step             0
   num_steps            1985
   outoforder_steps     []
sessionlog:checkpoint
   first_step           7129
   last_step            657167
   max_step             657167
   min_step             7129
   num_steps            99
   outoforder_steps     []
sessionlog:start
   outoforder_steps     []
   steps                [0L]
sessionlog:stop -


[2] Drill down into a particular tag using --tag.

Debug output for binary/Sign/Activations:
histograms
   first_step           491
   last_step            659823
   max_step             659823
   min_step             491
   num_steps            993
   outoforder_steps     []
"""

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

import collections
import itertools
import os

from mxconsole.backend.event_processing import event_accumulator
from mxconsole.backend.event_processing import event_file_loader
from mxconsole.backend.event_processing import event_multiplexer
from mxconsole.platform import app, flags, gfile
from mxconsole.protobuf.event_pb2 import SessionLog

FLAGS = flags.FLAGS


# Map of field names within summary.proto to the user-facing names that this
# script outputs.
SUMMARY_TYPE_TO_FIELD = {'simple_value': 'scalars',
                         'histo': 'histograms',
                         'image': 'images',
                         'audio': 'audio'}
for summary_type in event_accumulator.SUMMARY_TYPES:
  if summary_type not in SUMMARY_TYPE_TO_FIELD:
    SUMMARY_TYPE_TO_FIELD[summary_type] = summary_type

# Types of summaries that we may want to query for by tag.
TAG_FIELDS = list(SUMMARY_TYPE_TO_FIELD.values())

# Summaries that we want to see every instance of.
LONG_FIELDS = ['sessionlog:start', 'sessionlog:stop']

# Summaries that we only want an abridged digest of, since they would
# take too much screen real estate otherwise.
SHORT_FIELDS = ['graph', 'sessionlog:checkpoint'] + TAG_FIELDS

# All summary types that we can inspect.
TRACKED_FIELDS = SHORT_FIELDS + LONG_FIELDS

# An `Observation` contains the data within each Event file that the inspector
# cares about. The inspector accumulates Observations as it processes events.
Observation = collections.namedtuple('Observation', ['step', 'wall_time',
                                                     'tag'])

# An InspectionUnit is created for each organizational structure in the event
# files visible in the final terminal output. For instance, one InspectionUnit
# is created for each subdirectory in logdir. When asked to inspect a single
# event file, there may only be one InspectionUnit.

# The InspectionUnit contains the `name` of the organizational unit that will be
# printed to console, a `generator` that yields `Event` protos, and a mapping
# from string fields to `Observations` that the inspector creates.
InspectionUnit = collections.namedtuple('InspectionUnit', ['name', 'generator',
                                                           'field_to_obs'])

PRINT_SEPARATOR = '=' * 70 + '\n'


def get_field_to_observations_map(generator, query_for_tag=''):
  """Return a field to `Observations` dict for the event generator.

  Args:
    generator: A generator over event protos.
    query_for_tag: A string that if specified, only create observations for
      events with this tag name.

  Returns:
    A dict mapping keys in `TRACKED_FIELDS` to an `Observation` list.
  """

  def increment(stat, event, tag=''):
    assert stat in TRACKED_FIELDS
    field_to_obs[stat].append(Observation(step=event.step,
                                          wall_time=event.wall_time,
                                          tag=tag)._asdict())

  field_to_obs = dict([(t, []) for t in TRACKED_FIELDS])

  for event in generator:
    ## Process the event
    if event.HasField('graph_def') and (not query_for_tag):
      increment('graph', event)
    if event.HasField('session_log') and (not query_for_tag):
      status = event.session_log.status
      if status == SessionLog.START:
        increment('sessionlog:start', event)
      elif status == SessionLog.STOP:
        increment('sessionlog:stop', event)
      elif status == SessionLog.CHECKPOINT:
        increment('sessionlog:checkpoint', event)
    elif event.HasField('summary'):
      for value in event.summary.value:
        if query_for_tag and value.tag != query_for_tag:
          continue

        for proto_name, display_name in SUMMARY_TYPE_TO_FIELD.items():
          if value.HasField(proto_name):
            increment(display_name, event, value.tag)
  return field_to_obs


def get_unique_tags(field_to_obs):
  """Returns a dictionary of tags that a user could query over.

  Args:
    field_to_obs: Dict that maps string field to `Observation` list.

  Returns:
    A dict that maps keys in `TAG_FIELDS` to a list of string tags present in
    the event files. If the dict does not have any observations of the type,
    maps to an empty list so that we can render this to console.
  """
  return {field: sorted(set([x.get('tag', '') for x in observations]))
          for field, observations in field_to_obs.items()
          if field in TAG_FIELDS}


def print_dict(d, show_missing=True):
  """Prints a shallow dict to console.