# 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.
# ==============================================================================
"""A node transformer that includes utilities for SCT."""

import collections
import enum

import gast

from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import parser
from tensorflow.python.autograph.pyct import pretty_printer
from tensorflow.python.autograph.pyct import templates


class AnalysisLevel(enum.IntEnum):

  NONE = 0
  ACTIVITY = 1
  DEFINEDNESS = 2
  LIVENESS = 3


# TODO(znado): Use namedtuple.
class Context(object):
  """Contains information about a source code transformation.

  This object is mutable, and is updated during conversion. Not thread safe.

  Attributes:
    info: EntityInfo, immutable.
    namer: naming.Namer.
    current_origin: origin_info.OriginInfo, holds the OriginInfo of the last
      AST node to be processed successfully. Useful for error handling.
    user: An user-supplied context object. The object is opaque to the
      infrastructure, but will pe passed through to all custom transformations.
  """

  def __init__(self, info, namer, user_context):
    self.info = info
    self.namer = namer
    self.current_origin = None
    self.user = user_context


# TODO(mdan): Move to a standalone file.
class EntityInfo(
    collections.namedtuple(
        'EntityInfo',
        ('name', 'source_code', 'source_file', 'future_features', 'namespace'))
):
  """Contains information about a Python entity.

  Immutable.

  Examples of entities include functions and classes.

  Attributes:
    name: The name that identifies this entity.
    source_code: The entity's source code.
    source_file: The entity's source file.
    future_features: Tuple[Text], the future features that this entity was
      compiled with. See
      https://docs.python.org/2/reference/simple_stmts.html#future.
    namespace: Dict[str, ], containing symbols visible to the entity (excluding
      parameters).
  """
  pass


class _StateStack(object):
  """Templated context manager.

  This class provides syntactic sugar for a stack of objects of known
  type. It allows accessing attributes of the object at the top of the stack
  directly against this object, which allows for very terse syntax.

  For example, this code:

    stack = _StateStack(Foo)
    stack.enter()
    stack.bar

  Is equivalent to:

    stack = []
    stack.append(Foo())
    foo = stack[-1]
    foo.bar

  See _State for more on how this is used.

  Attributes:
    type: Any, the type of objects that this stack holds
    level: int, the current stack depth
    stack: List[Any], the actual stack
    value: Any, the instance of the object at the top of the stack
  """

  def __init__(self, type_):
    # Because we override __setattr__, we need to attach these attributes using
    # the superclass' setattr.
    object.__setattr__(self, 'type', type_)
    object.__setattr__(self, '_stack', [])
    if not hasattr(type_, 'no_root'):
      self.enter()

  def __enter__(self):
    self.enter()
    return self

  def __exit__(self, exc_type, exc_value, traceback):
    self.exit()

  def enter(self):
    self._stack.append(self.type())

  def exit(self):
    self._stack.pop()

  @property
  def stack(self):
    return self._stack

  @property
  def level(self):
    return len(self._stack)

  @property
  def value(self):
    return self._stack[-1]

  def __iter__(self):
    return iter(self._stack)

  def __getattr__(self, key):
    return getattr(self._stack[-1], key)

  def __setattr__(self, key, value):
    setattr(self._stack[-1], key, value)


class _State(object):
  """Syntactic sugar for accessing an instance of a StateStack context manager.

  This structure offers syntactic sugar over a dict of stacks of objects
  of known type. These structures are useful to keep state during AST walks.
  Multiple different scopes can be tracked in parallel. For example:

    s = _State()

    s[foo].enter()
    s[bar].enter()  # this will not affect s[foo]

  Element access has special semantics:
    * keys are a data type
    * element values are _StateStack(type=key) objects
    * missing elements are automatically added, similarly to defaultdict

  For example, the following block :

    _State s
    s[Foo]

  Is equivalent to:

    s = {}
    if Foo not in s:
      s[Foo] = Foo()
    s[Foo]

  See Base for how it's used.
  """

  def __init__(self):
    self._value = {}

  def __getitem__(self, key):
    if key not in self._value:
      self._value[key] = _StateStack(key)
    return self._value[key]


class NodeStateTracker(object):
  """Base class for general-purpose Python code transformation.

  This abstract class provides helpful functions, like state tracking within
  the scope of arbitrary node, helpers for processing code blocks, debugging,
  mapping of transformed code to original code, and others.

  Scope-local state tracking: to keep state across nodes, at the level of
  (possibly nested) scopes, use enter/exit_local_scope and set/get_local.
  You must call enter/exit_local_scope manually, but the transformer detects
  when they are not properly paired.

  The transformer allows keeping state across calls that is local
  to arbitrary nodes and their descendants, using the self.state attribute.
  Multiple independent scopes are allowed and automatically constructed.

  For example, to keep track of the `If` node that encloses any `Name` node,
  one can write:

  ```
    class FooType(object):

      def __init__(self):
        self.foo_property = None

    class DummyTransformer(NodeStateTracker, ast.NodeTransformer):

      def visit_If(self, node):
        self.state[FooType].enter()
        self.state[FooType].foo_property = node
        node = self.veneric_visit(node)
        self.state[FooType].exit()
        return node

      def visit_Name(self, node):
        self.state[FooType].foo_property  # will hold the innermost enclosing if
  ```

  Alternatively, the `enter()`/`exit()` calls can be managed by a `with`
  statement:

  ```
      def visit_If(self, node):
        with self.state[FooType] as foo:
          foo.foo_property = node
          return self.generic_visit(node)
  ```
  """

  # TODO(mdan): Document all extra features.

  def __init__(self, ctx):
    """Initialize the transformer.

    Subclasses should call this.

    Args:
      ctx: A Context object.
    """
    self._lineno = 0
    self._col_offset = 0
    self.ctx = ctx

    # Allows scoping of local variables to keep state across calls to visit_*
    # methods. Multiple scope hierarchies may exist and are keyed by tag. A
    # scope is valid at one or more nodes and all its children. Scopes created
    # in child nodes supersede their parent. Scopes are isolated from one
    # another.
    self.state = _State()

  def debug_print(self, node):
    """Helper method useful for debugging. Prints the AST."""
    if __debug__:
      print(pretty_printer.fmt(node))
    return node

  def debug_print_src(self, node):
    """Helper method useful for debugging. Prints the AST as code."""
    if __debug__:
      print(parser.unparse(node))
    return node

  def visit_block(self, nodes, before_visit=None, after_visit=None):
    """A more powerful version of generic_visit for statement blocks.

    An example of a block is the body of an if statement.

    This function allows specifying a postprocessing callback (the
    after_visit argument) argument which can be used to move nodes to a new
    destination. This is done by after_visit by returning a non-null
    second return value, e.g. return new_node, new_destination.

    For example, a transformer could perform the following move:

        foo()
        bar()
        baz()

        foo()
        if cond:
          bar()
          baz()

    The above could be done with a postprocessor of this kind:

        def after_visit(node):
          if node_is_function_call(bar):
            new_container_node = build_cond()
            new_container_node.body.append(node)
            return new_container_node, new_container_node.body
          else:
            # Once we set a new destination, all subsequent items will be
            # moved to it, so we don't need to explicitly handle baz.
            return node, None

    Args:
      nodes: enumerable of AST node objects. If None, the function returns None.
      before_visit: optional callable that is called before visiting each item
        in nodes
      after_visit: optional callable that takes in an AST node and returns a
        tuple (new_node, new_destination). It is called after visiting each item
        in nodes. Is used in the same was as the
          visit_* methods: new_node will replace the node; if not None,
            new_destination must be a list, and subsequent nodes will be placed
            in this list instead of the list returned by visit_block.

    Returns:
      A list of AST node objects containing the transformed items fron nodes,
      except those nodes that have been relocated using after_visit.
    """
    if nodes is None:
      return None

    results = []
    node_destination = results
    for node in nodes:
      if before_visit:
        # TODO(mdan): We can modify node here too, if ever needed.
        before_visit()

      replacement = self.visit(node)

      if after_visit and replacement:
        replacement, new_destination = after_visit(replacement)
      else:
        new_destination = None

      if replacement:
        if isinstance(replacement, (list, tuple)):
          node_destination.extend(replacement)
        else:
          node_destination.append(replacement)

      # Allow the postprocessor to reroute the remaining nodes to a new list.
      if new_destination is not None:
        node_destination = new_destination
    return results


# TODO(mdan): Rename to PythonCodeTransformer.
class Base(NodeStateTracker, gast.NodeTransformer):
  """Base class for general-purpose Python-to-Python code transformation.

  This is an extension of ast.NodeTransformer that provides the additional
  functions offered by NodeStateTracker.
  """

  def create_assignment(self, target, expression):
    template = """
      target = expression
    """
    return templates.replace(template, target=target, expression=expression)

  # TODO(mdan): Remove.
  def apply_to_single_assignments(self, targets, values, apply_fn):
    """Applies a function to each individual assignment.

    This function can process a possibly-unpacked (e.g. a, b = c, d) assignment.
    It tries to break down the unpacking if possible. In effect, it has the same
    effect as passing the assigned values in SSA form to apply_fn.

    Examples:

    The following will result in apply_fn(a, c), apply_fn(b, d):

        a, b = c, d

    The following will result in apply_fn(a, c[0]), apply_fn(b, c[1]):

        a, b = c

    The following will result in apply_fn(a, (b, c)):

        a = b, c

    It uses the visitor pattern to allow subclasses to process single
    assignments individually.

    Args:
      targets: list, tuple of or individual AST node. Should be used with the
        targets field of an ast.Assign node.
      values: an AST node.
      apply_fn: a function of a single argument, which will be called with the
        respective nodes of each single assignment. The signature is
        apply_fn(target, value), no return value.
    """
    if not isinstance(targets, (list, tuple)):
      targets = (targets,)
    for target in targets:
      if isinstance(target, (gast.Tuple, gast.List)):
        for i in range(len(target.elts)):
          target_el = target.elts[i]
          if isinstance(values, (gast.Tuple, gast.List)):
            value_el = values.elts[i]
          else:
            value_el = gast.Subscript(values, i, ctx=gast.Store())
          self.apply_to_single_assignments(target_el, value_el, apply_fn)
      else:
        # TODO(mdan): Look into allowing to rewrite the AST here.
        apply_fn(target, values)

  def visit(self, node):
    if not isinstance(node, gast.AST):
      # This is not that uncommon a mistake: various node bodies are lists, for
      # example, posing a land mine for transformers that need to recursively
      # call `visit`.  The error needs to be raised before the exception handler
      # below is installed, because said handler will mess up if `node` is not,
      # in fact, a node.
      msg = ('invalid value for "node": expected "ast.AST", got "{}"; to'
             ' visit lists of nodes, use "visit_block" instead').format(
                 type(node))
      raise ValueError(msg)

    if anno.hasanno(node, anno.Basic.SKIP_PROCESSING):
      return node

    parent_origin = self.ctx.current_origin
    if anno.hasanno(node, anno.Basic.ORIGIN):
      self.ctx.current_origin = anno.getanno(node, anno.Basic.ORIGIN)

    try:
      processing_expr_node = isinstance(node, gast.Expr)
      if processing_expr_node:
        entry_expr_value = node.value

      result = super(Base, self).visit(node)

      # Adjust for consistency: replacing the value of an Expr with
      # an Assign node removes the need for the Expr node.
      if (processing_expr_node and isinstance(result, gast.Expr) and
          (result.value is not entry_expr_value)):
        # When the replacement is a list, it is assumed that the list came
        # from a template that contained a number of statements, which
        # themselves are standalone and don't require an enclosing Expr.
        if isinstance(result.value,
                      (list, tuple, gast.Assign, gast.AugAssign)):
          result = result.value

      # By default, all replacements receive the origin info of the replaced
      # node.
      if result is not node and result is not None:
        inherited_origin = anno.getanno(
            node, anno.Basic.ORIGIN, default=parent_origin)
        if inherited_origin is not None:
          nodes_to_adjust = result
          if isinstance(result, (list, tuple)):
            nodes_to_adjust = result
          else:
            nodes_to_adjust = (result,)
          for n in nodes_to_adjust:
            if not anno.hasanno(n, anno.Basic.ORIGIN):
              anno.setanno(n, anno.Basic.ORIGIN, inherited_origin)
    finally:
      self.ctx.current_origin = parent_origin

    return result


class CodeGenerator(NodeStateTracker, gast.NodeVisitor):
  """Base class for general-purpose Python-to-string code transformation.

  Similar to Base, but outputs arbitrary strings instead of a Python AST.

  This uses the same visitor mechanism that the standard NodeVisitor uses,
  meaning that subclasses write handlers for the different kinds of nodes.
  New code is generated using the emit method, which appends to a code buffer
  that can be afterwards obtained from code_buffer.

  Example:

    class SimpleCodeGen(CodeGenerator):

      def visitIf(self, node):
        self.emit('if ')
        self.visit(node.test)
        self.emit(' { ')
        self.visit(node.body)
        self.emit(' } else { ')
        self.visit(node.orelse)
        self.emit(' } ')

    node = ast.parse(...)
    gen = SimpleCodeGen()
    gen.visit(node)
    # gen.code_buffer contains the resulting code
  """

  def __init__(self, ctx):
    super(CodeGenerator, self).__init__(ctx)

    self._output_code = ''
    self.source_map = {}

  def emit(self, code):
    self._output_code += code

  @property
  def code_buffer(self):
    return self._output_code

  def visit(self, node):
    if anno.hasanno(node, anno.Basic.SKIP_PROCESSING):
      return

    parent_origin = self.ctx.current_origin
    eof_before = len(self._output_code)
    if anno.hasanno(node, anno.Basic.ORIGIN):
      self.ctx.current_origin = anno.getanno(node, anno.Basic.ORIGIN)

    try:
      ret = super(CodeGenerator, self).visit(node)

      # By default, all replacements receive the origin info of the replaced
      # node.
      eof_after = len(self._output_code)
      if eof_before - eof_after:
        inherited_origin = anno.getanno(
            node, anno.Basic.ORIGIN, default=parent_origin)
        if inherited_origin is not None:
          self.source_map[(eof_before, eof_after)] = inherited_origin
      return ret
    finally:
      self.ctx.current_origin = parent_origin