# Copyright 2016 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.
# ==============================================================================
"""Stores debugging information regarding TensorFlow model."""

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

import os
import parser
import re
import token

from google.protobuf import json_format

from tensorflow.contrib.tensorboard.plugins.trace.trace_info_pb2 import TraceInfo
from tensorflow.python.framework import ops
from tensorflow.python.platform import gfile

# List of regex patterns that match files in the core tensorflow library.
TF_LIB_REGEX_FPATHS = [os.sep + os.path.join('tensorflow', 'python')]

LEFT_TOKENS = [token.LPAR, token.LSQB, token.LBRACE]
RIGHT_TOKENS = [token.RPAR, token.RSQB, token.RBRACE]
TOKENS = LEFT_TOKENS + RIGHT_TOKENS


def store_trace_info(output_file_path,
                     graph=None,
                     ignore_regex_fpaths=None):
  """Collects and stores trace information for a TensorFlow model.

  The output proto is stored in json format.

  Args:
    output_file_path: The path where to store the output proto.
    graph: Optional. The data flow graph. Defaults to `tf.get_default_graph()`.
    ignore_regex_fpaths: Optional. Files whose path matches any of the regexes
        in this list will be ignored. Defaults to patterns that match the core
        tensorflow python library.
  """
  graph = graph or ops.get_default_graph()

  if not ignore_regex_fpaths:
    ignore_regex_fpaths = TF_LIB_REGEX_FPATHS

  trace_info = TraceInfo()
  # Extract trace information for every op in the graph.
  source_fpaths = set()
  for op in graph.get_operations():
    op_info = trace_info.ops.add()
    op_info.name = op.name
    op_info.op_type = op.type
    op_info.device = op.device
    for trace in op.traceback:
      fname, lineno, _, _ = trace
      # Ignore traces in specified file paths.
      if os.path.isabs(fname) and not _ignore_file_path(fname,
                                                        ignore_regex_fpaths):
        line_trace = op_info.traceback.add()
        line_trace.file_path = fname
        line_trace.line_number = lineno
        source_fpaths.add(fname)
    _add_data_from_tensors(op.inputs, op_info.inputs)
    _add_data_from_tensors(op.outputs, op_info.outputs)

  # Read the source files involved in the graph construction.
  for fpath in source_fpaths:
    file_info = trace_info.files.add()

    with gfile.Open(fpath, 'r') as f:
      source = f.read()

    file_info.file_path = fpath
    file_info.source_code = source

    line2start = find_multiline_statements(source)

    for key, value in line2start.items():
      file_info.multiline_statements[key] = value

  # Make sure the directory for the output file exists.
  output_file_path = os.path.expanduser(output_file_path)
  output_dir = os.path.dirname(output_file_path)
  if not gfile.Exists(output_dir):
    gfile.MakeDirs(output_dir)

  # Store the debug information.
  with gfile.Open(output_file_path, 'w') as f:
    f.write(json_format.MessageToJson(trace_info))


def find_multiline_statements(source):
  """Parses the python source and finds multiline statements.

  Based on counting the number of open and closed parenthesis on each line.

  Args:
    source: The source code string.

  Returns:
    A dict that maps a line index A to a line index B, where A is the end of a
    multiline statement and B is the start. Line indexing is 0-based.
  """
  # Get the AST.
  tree = parser.suite(source)
  line2paren_count = [0] * (source.count('\n') + 1)
  _count_brackets_braces_parenthesis(tree.totuple(True), line2paren_count)

  line2start = {}
  for end in range(len(line2paren_count)):
    if line2paren_count[end] >= 0:
      # This is not the end of a multiline statement.
      continue
    cumulative_paren_count = 0
    for start in range(end, -1, -1):
      cumulative_paren_count += line2paren_count[start]
      if cumulative_paren_count == 0:
        line2start[end] = start
        break
  return line2start


def _add_data_from_tensors(tensors, info):
  for t in tensors:
    tensor_info = info.add()

    shape = t.get_shape()
    if shape.ndims:
      shape = [(-1 if s is None else s) for s in shape.as_list()]
      tensor_info.shape.extend(shape)
    tensor_info.dtype = t.dtype.name
    tensor_info.num_bytes_per_elem = t.dtype.size

    for c in t.consumers():
      tensor_info.consumers.append(c.name)


def _ignore_file_path(fname, ignore_regex_fpaths):
  for regex_pattern in ignore_regex_fpaths:
    if re.search(regex_pattern, fname):
      return True
  return False


def _count_brackets_braces_parenthesis(node, line2par):
  if isinstance(node[1], tuple):
    for child in node[1:]:
      _count_brackets_braces_parenthesis(child, line2par)
  else:
    tok = node[0]
    if tok in TOKENS:
      lineno = node[2]
      line2par[lineno - 1] += (1 if tok in LEFT_TOKENS else -1)
  return line2par
