# 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.
# ==============================================================================
"""AST conversion templates.

Adapted from Tangent.
"""

import ast
import textwrap

import gast

from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import ast_util
from tensorflow.python.autograph.pyct import parser
from tensorflow.python.autograph.pyct import qual_names


class ContextAdjuster(gast.NodeTransformer):
  """Adjusts the ctx field of nodes to ensure consistency.

  This transformer can change the ctx fields of a variable, tuple and other
  AST elements that allow one, based on whether the element is being read or
  written.
  """

  def __init__(self, override_value):
    self._ctx_override = override_value

  def visit(self, node):
    original_override = self._ctx_override
    node = super(ContextAdjuster, self).visit(node)
    if hasattr(node, 'ctx'):
      assert node.ctx is not None, 'node {} has ctx unset'.format(node)
    self._ctx_override = original_override
    return node

  def _apply_override(self, node):
    if self._ctx_override is not None:
      node.ctx = self._ctx_override()

  def visit_Attribute(self, node):
    self._apply_override(node)
    self._ctx_override = gast.Load
    node = self.generic_visit(node)
    return node

  def visit_Tuple(self, node):
    self._apply_override(node)
    return self.generic_visit(node)

  def visit_List(self, node):
    self._apply_override(node)
    return self.generic_visit(node)

  def visit_Name(self, node):
    self._apply_override(node)
    return self.generic_visit(node)

  def visit_Call(self, node):
    self._apply_override(node)
    # We may be able to override these to Load(), but for now it's simpler
    # to just assert that they're set.
    self._ctx_override = None
    return self.generic_visit(node)

  def visit_Dict(self, node):
    # We may be able to override these to Load(), but for now it's simpler
    # to just assert that they're set.
    self._ctx_override = None
    return self.generic_visit(node)

  def visit_Subscript(self, node):
    self._apply_override(node)
    self._ctx_override = gast.Load
    node.value = self.visit(node.value)
    return self.generic_visit(node)

  def visit_comprehension(self, node):
    # We may be able to override some of these, but for now it's simpler
    # to just assert that they're set.
    self._ctx_override = None
    return self.generic_visit(node)

  def visit_Lambda(self, node):
    # We may be able to override some of these, but for now it's simpler
    # to just assert that they're set.
    self._ctx_override = None
    return self.generic_visit(node)


class ReplaceTransformer(gast.NodeTransformer):
  """Replace AST nodes."""

  def __init__(self, replacements):
    """Create a new ReplaceTransformer.

    Args:
      replacements: A mapping from placeholder names to (lists of) AST nodes
          that these placeholders will be replaced by.
    """
    self.replacements = replacements
    self.in_replacements = False
    self.preserved_annos = {
        anno.Basic.DIRECTIVES,
        anno.Basic.EXTRA_LOOP_TEST,
        anno.Basic.ORIGIN,
        anno.Basic.SKIP_PROCESSING,
        anno.Static.ORIG_DEFINITIONS,
        'function_context_name',
    }

  def _prepare_replacement(self, replaced, key):
    """Prepares a replacement AST that's safe to swap in for a node.

    Args:
      replaced: ast.AST, the node being replaced
      key: Hashable, the key of the replacement AST
    Returns:
      ast.AST, the replacement AST
    """
    repl = self.replacements[key]

    new_nodes = ast_util.copy_clean(repl, preserve_annos=self.preserved_annos)
    if isinstance(new_nodes, gast.AST):
      new_nodes = [new_nodes]

    return new_nodes

  def visit_Expr(self, node):
    # When replacing a placeholder with an entire statement, the replacement
    # must stand on its own and not be wrapped in an Expr.
    new_value = self.visit(node.value)
    if new_value is node.value:
      return node
    return new_value

  def visit_keyword(self, node):
    if node.arg not in self.replacements:
      return self.generic_visit(node)

    repl = self._prepare_replacement(node, node.arg)
    if isinstance(repl, gast.keyword):
      return repl
    elif (repl and isinstance(repl, (list, tuple)) and
          all(isinstance(r, gast.keyword) for r in repl)):
      return repl
    # TODO(mdan): We may allow replacing with a string as well.
    # For example, if one wanted to replace foo with bar in foo=baz, then
    # we could allow changing just node arg, so that we end up with bar=baz.
    raise ValueError(
        'a keyword argument may only be replaced by another keyword or a '
        'non-empty list of keywords. Found: {} for keyword {}'.format(
            repl, node.arg))

  def visit_FunctionDef(self, node):
    node = self.generic_visit(node)
    if node.name not in self.replacements:
      return node

    repl = self.replacements[node.name]
    if not isinstance(repl, (gast.Name, ast.Name)):
      raise ValueError(
          'a function name can only be replaced by a Name node. Found: %s' %
          repl)
    node.name = repl.id
    return node

  def visit_Attribute(self, node):
    node = self.generic_visit(node)
    if node.attr not in self.replacements:
      return node

    repl = self.replacements[node.attr]
    if not isinstance(repl, gast.Name):
      raise ValueError(
          'An attribute can only be replaced by a Name node. Found: %s' % repl)
    node.attr = repl.id
    return node

  def visit_Name(self, node):
    if node.id not in self.replacements:
      return node

    new_nodes = self._prepare_replacement(node, node.id)

    if not new_nodes:
      return new_nodes

    # Preserve the target context.
    adjuster = ContextAdjuster(type(node.ctx))
    for n in new_nodes:
      if hasattr(n, 'ctx'):
        adjuster.visit(n)

    if len(new_nodes) == 1:
      new_nodes, = new_nodes

    return new_nodes


def _convert_to_ast(n):
  """Converts from a known data type to AST."""
  # Note: When generating AST nodes from strings/QNs in isolation, ctx is
  # unknown. ctx must be filled in according to the template being used.
  # See ReplaceTransformer.visit_Name.
  if isinstance(n, str):
    return gast.Name(id=n, ctx=None, annotation=None, type_comment=None)
  if isinstance(n, qual_names.QN):
    return n.ast()
  if isinstance(n, list):
    return [_convert_to_ast(e) for e in n]
  if isinstance(n, tuple):
    return tuple(_convert_to_ast(e) for e in n)
  return n


def replace(template, **replacements):
  """Replaces placeholders in a Python template.

  AST Name and Tuple nodes always receive the context that inferred from
  the template. However, when replacing more complex nodes (that can potentially
  contain Name children), then the caller is responsible for setting the
  appropriate context.

  Args:
    template: A string representing Python code. Any symbol name can be used
        that appears in the template code can be used as placeholder.
    **replacements: A mapping from placeholder names to (lists of) AST nodes
        that these placeholders will be replaced by. String values are also
        supported as a shorthand for AST Name nodes with the respective ID.

  Returns:
    An AST node or list of AST nodes with the replacements made. If the
    template was a function, a list will be returned. If the template was a
    node, the same node will be returned. If the template was a string, an
    AST node will be returned (a `Module` node in the case of a multi-line
    string, an `Expr` node otherwise).

  Raises:
    ValueError: if the arguments are incorrect.
  """
  if not isinstance(template, str):
    raise ValueError('Expected string template, got %s' % type(template))
  for k in replacements:
    replacements[k] = _convert_to_ast(replacements[k])
  template_str = parser.STANDARD_PREAMBLE + textwrap.dedent(template)
  nodes = parser.parse(
      template_str,
      preamble_len=parser.STANDARD_PREAMBLE_LEN,
      single_node=False)
  results = []
  for node in nodes:
    node = ReplaceTransformer(replacements).visit(node)
    if isinstance(node, (list, tuple)):
      results.extend(node)
    else:
      results.append(node)
  results = [qual_names.resolve(r) for r in results]
  return results


def replace_as_expression(template, **replacements):
  """Variant of replace that generates expressions, instead of code blocks."""
  replacement = replace(template, **replacements)
  if len(replacement) != 1:
    raise ValueError(
        'single expression expected; for more general templates use replace')
  node, = replacement

  if isinstance(node, gast.Expr):
    return node.value
  elif isinstance(node, gast.Name):
    return node

  raise ValueError(
      'the template is expected to generate an expression or a name node;'
      ' instead found %s' % node)