# 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.
# ==============================================================================
"""Handles control flow statements: while, for, if."""

import gast

from tensorflow.python.autograph.core import converter
from tensorflow.python.autograph.lang import directives
from tensorflow.python.autograph.pyct import anno
from tensorflow.python.autograph.pyct import cfg
from tensorflow.python.autograph.pyct import origin_info
from tensorflow.python.autograph.pyct import parser
from tensorflow.python.autograph.pyct import qual_names
from tensorflow.python.autograph.pyct import templates
from tensorflow.python.autograph.pyct.static_analysis import activity
from tensorflow.python.autograph.pyct.static_analysis import annos
from tensorflow.python.autograph.pyct.static_analysis import liveness
from tensorflow.python.autograph.pyct.static_analysis import reaching_definitions
from tensorflow.python.autograph.pyct.static_analysis import reaching_fndefs


class _Function(object):

  scope = None


class ControlFlowTransformer(converter.Base):
  """Transforms control flow structures like loops an conditionals."""

  def visit_Lambda(self, node):
    with self.state[_Function] as fn:
      fn.scope = anno.getanno(node, anno.Static.SCOPE)
      return self.generic_visit(node)

  def visit_FunctionDef(self, node):
    with self.state[_Function] as fn:
      fn.scope = anno.getanno(node, annos.NodeAnno.BODY_SCOPE)
      return self.generic_visit(node)

  def _create_nonlocal_declarations(self, vars_):
    vars_ = set(vars_)
    results = []
    global_vars = self.state[_Function].scope.globals & vars_

    if global_vars:
      results.append(gast.Global([str(v) for v in global_vars]))

    nonlocal_vars = [
        v for v in vars_ if not v.is_composite() and v not in global_vars]
    if nonlocal_vars:
      results.append(gast.Nonlocal([str(v) for v in nonlocal_vars]))

    return results

  def _create_state_functions(
      self, block_vars, nonlocal_declarations, getter_name, setter_name):
    if not block_vars:
      template = """
        def getter_name():
          return ()
        def setter_name(block_vars):
          pass
      """
      return templates.replace(
          template, getter_name=getter_name, setter_name=setter_name)

    guarded_block_vars = []
    for v in block_vars:
      if v.is_simple():
        guarded_block_vars.append(v)
      else:
        guarded_block_vars.append(
            templates.replace_as_expression(
                'ag__.ldu(lambda: var_, name)',
                var_=v,
                name=gast.Constant(str(v), kind=None)))

    template = """
      def getter_name():
        return guarded_state_vars,
      def setter_name(vars_):
        nonlocal_declarations
        state_vars, = vars_
    """
    return templates.replace(
        template,
        nonlocal_declarations=nonlocal_declarations,
        getter_name=getter_name,
        guarded_state_vars=guarded_block_vars,
        setter_name=setter_name,
        state_vars=tuple(block_vars))

  def _create_loop_options(self, node):
    if not anno.hasanno(node, anno.Basic.DIRECTIVES):
      return gast.Dict([], [])

    loop_directives = anno.getanno(node, anno.Basic.DIRECTIVES)
    if directives.set_loop_options not in loop_directives:
      return gast.Dict([], [])

    opts_dict = loop_directives[directives.set_loop_options]
    str_keys, values = zip(*opts_dict.items())
    keys = [gast.Constant(s, kind=None) for s in str_keys]
    values = list(values)  # ast and gast don't play well with tuples.
    return gast.Dict(keys, values)

  def _create_undefined_assigns(self, undefined_symbols):
    assignments = []
    for s in undefined_symbols:
      template = '''
        var = ag__.Undefined(symbol_name)
      '''
      assignments += templates.replace(
          template,
          var=s,
          symbol_name=gast.Constant(s.ssf(), kind=None))
    return assignments

  def _get_block_basic_vars(self, modified, live_in, live_out):
    nonlocals = self.state[_Function].scope.nonlocals
    basic_scope_vars = []
    for s in modified:
      if s.is_composite():
        # TODO(mdan): Raise an error when this happens for a TF scope.
        continue
      # Variables not live into or out of the scope are considered local to the
      # scope.
      if s in live_in or s in live_out or s in nonlocals:
        basic_scope_vars.append(s)
      continue
    return frozenset(basic_scope_vars)

  def _get_block_composite_vars(self, modified, live_in):
    # The scope variables corresponding to composite symbols (e.g. `self.x`).
    composite_scope_vars = []
    for s in modified:
      if not s.is_composite():
        continue
      # Mutations made to objects created inside the scope will appear as writes
      # to composite symbols. Because these mutations appear as modifications
      # made to composite symbols, we check whether the composite's parent is
      # actually live into the scope.
      # Example:
      #   while cond:
      #     x = Foo()
      #     x.foo = 2 * x.foo  # x.foo is live into the scope, but x is not.
      #
      # Note that some parents might not be symbols - for example, in x['foo'],
      # 'foo' is a parent, but it's a literal, not a symbol. We don't check the
      # liveness of literals.
      support_set_symbols = tuple(
          sss for sss in s.support_set if sss.is_symbol())
      if not all(sss in live_in for sss in support_set_symbols):
        continue
      composite_scope_vars.append(s)
    return frozenset(composite_scope_vars)

  def _get_block_vars(self, node, modified):
    """Determines the variables affected inside a control flow statement."""
    defined_in = anno.getanno(node, anno.Static.DEFINED_VARS_IN)
    live_in = anno.getanno(node, anno.Static.LIVE_VARS_IN)
    live_out = anno.getanno(node, anno.Static.LIVE_VARS_OUT)
    fn_scope = self.state[_Function].scope

    basic_scope_vars = self._get_block_basic_vars(
        modified,
        live_in,
        live_out)
    composite_scope_vars = self._get_block_composite_vars(modified, live_in)
    scope_vars = tuple(basic_scope_vars | composite_scope_vars)

    # Variables that are modified inside the scope, but not defined
    # before entering it. Only simple variables must be defined. The
    # composite ones will be implicitly checked at runtime.
    possibly_undefined = (
        modified - defined_in - fn_scope.globals - fn_scope.nonlocals)
    undefined = tuple(v for v in possibly_undefined if not v.is_composite())

    # Variables that are modified inside the scope, and depend on values outside
    # it.
    input_only = basic_scope_vars & live_in - live_out

    # Place the outputs first, then sort lexicographically.
    scope_vars = sorted(scope_vars, key=lambda v: (v in input_only, v))
    nouts = len(scope_vars) - len(input_only)

    return scope_vars, undefined, nouts

  def visit_If(self, node):
    node = self.generic_visit(node)
    body_scope = anno.getanno(node, annos.NodeAnno.BODY_SCOPE)
    orelse_scope = anno.getanno(node, annos.NodeAnno.ORELSE_SCOPE)

    cond_vars, undefined, nouts = self._get_block_vars(
        node, body_scope.bound | orelse_scope.bound)

    undefined_assigns = self._create_undefined_assigns(undefined)

    nonlocal_declarations = self._create_nonlocal_declarations(cond_vars)

    reserved = body_scope.referenced | orelse_scope.referenced
    state_getter_name = self.ctx.namer.new_symbol('get_state', reserved)
    state_setter_name = self.ctx.namer.new_symbol('set_state', reserved)
    state_functions = self._create_state_functions(
        cond_vars, nonlocal_declarations, state_getter_name, state_setter_name)

    orelse_body = node.orelse
    if not orelse_body:
      orelse_body = [gast.Pass()]

    template = """
      state_functions
      def body_name():
        nonlocal_declarations
        body
      def orelse_name():
        nonlocal_declarations
        orelse
      undefined_assigns
      ag__.if_stmt(
        test,
        body_name,
        orelse_name,
        state_getter_name,
        state_setter_name,
        (symbol_names,),
        nouts)
    """
    new_nodes = templates.replace(
        template,
        body=node.body,
        body_name=self.ctx.namer.new_symbol('if_body', reserved),
        orelse=orelse_body,
        orelse_name=self.ctx.namer.new_symbol('else_body', reserved),
        nonlocal_declarations=nonlocal_declarations,
        nouts=gast.Constant(nouts, kind=None),
        state_functions=state_functions,
        state_getter_name=state_getter_name,
        state_setter_name=state_setter_name,
        symbol_names=tuple(gast.Constant(str(s), kind=None) for s in cond_vars),
        test=node.test,
        undefined_assigns=undefined_assigns)
    origin_info.copy_origin(node, new_nodes[-1])
    return new_nodes

  def visit_While(self, node):
    node = self.generic_visit(node)
    body_scope = anno.getanno(node, annos.NodeAnno.BODY_SCOPE)

    loop_vars, undefined, _ = self._get_block_vars(node, body_scope.bound)

    undefined_assigns = self._create_undefined_assigns(undefined)

    nonlocal_declarations = self._create_nonlocal_declarations(loop_vars)

    reserved = body_scope.referenced
    state_getter_name = self.ctx.namer.new_symbol('get_state', reserved)
    state_setter_name = self.ctx.namer.new_symbol('set_state', reserved)
    state_functions = self._create_state_functions(
        loop_vars, nonlocal_declarations, state_getter_name, state_setter_name)

    opts = self._create_loop_options(node)

    template = """
      state_functions
      def body_name():
        nonlocal_declarations
        body
      def test_name():
        return test
      undefined_assigns
      ag__.while_stmt(
          test_name,
          body_name,
          state_getter_name,
          state_setter_name,
          (symbol_names,),
          opts)
    """
    new_nodes = templates.replace(
        template,
        body=node.body,
        body_name=self.ctx.namer.new_symbol('loop_body', reserved),
        nonlocal_declarations=nonlocal_declarations,
        opts=opts,
        state_functions=state_functions,
        state_getter_name=state_getter_name,
        state_setter_name=state_setter_name,
        symbol_names=tuple(gast.Constant(str(s), kind=None) for s in loop_vars),
        test=node.test,
        test_name=self.ctx.namer.new_symbol('loop_test', reserved),
        undefined_assigns=undefined_assigns)
    origin_info.copy_origin(node, new_nodes[-1])
    return new_nodes

  def visit_For(self, node):
    node = self.generic_visit(node)
    body_scope = anno.getanno(node, annos.NodeAnno.BODY_SCOPE)
    iter_scope = anno.getanno(node, annos.NodeAnno.ITERATE_SCOPE)

    loop_vars, undefined, _ = self._get_block_vars(
        node, body_scope.bound | iter_scope.bound)

    undefined_assigns = self._create_undefined_assigns(undefined)

    nonlocal_declarations = self._create_nonlocal_declarations(loop_vars)

    reserved = body_scope.referenced | iter_scope.referenced
    state_getter_name = self.ctx.namer.new_symbol('get_state', reserved)
    state_setter_name = self.ctx.namer.new_symbol('set_state', reserved)
    state_functions = self._create_state_functions(
        loop_vars, nonlocal_declarations, state_getter_name, state_setter_name)

    opts = self._create_loop_options(node)
    opts.keys.append(gast.Constant('iterate_names', kind=None))
    opts.values.append(gast.Constant(
        parser.unparse(node.target, include_encoding_marker=False), kind=None))

    if anno.hasanno(node, anno.Basic.EXTRA_LOOP_TEST):
      extra_test = anno.getanno(node, anno.Basic.EXTRA_LOOP_TEST)
      extra_test_name = self.ctx.namer.new_symbol(
          'extra_test', reserved)
      template = """
        def extra_test_name():
          nonlocal_declarations
          return extra_test_expr
      """
      extra_test_function = templates.replace(
          template,
          extra_test_expr=extra_test,
          extra_test_name=extra_test_name,
          loop_vars=loop_vars,
          nonlocal_declarations=nonlocal_declarations)
    else:
      extra_test_name = parser.parse_expression('None')
      extra_test_function = []

    # iterate_arg_name holds a single arg with the iterates, which may be a
    # tuple.
    iterate_arg_name = self.ctx.namer.new_symbol('itr', reserved)
    template = """
      iterates = iterate_arg_name
    """
    iterate_expansion = templates.replace(
        template, iterate_arg_name=iterate_arg_name, iterates=node.target)
    origin_info.copy_origin(node, iterate_expansion)

    template = """
      state_functions
      def body_name(iterate_arg_name):
        nonlocal_declarations
        iterate_expansion
        body
      extra_test_function
      undefined_assigns
      ag__.for_stmt(
          iterated,
          extra_test_name,
          body_name,
          state_getter_name,
          state_setter_name,
          (symbol_names,),
          opts)
    """
    new_nodes = templates.replace(
        template,
        body=node.body,
        body_name=self.ctx.namer.new_symbol('loop_body', reserved),
        extra_test_function=extra_test_function,
        extra_test_name=extra_test_name,
        iterate_arg_name=iterate_arg_name,
        iterate_expansion=iterate_expansion,
        iterated=node.iter,
        nonlocal_declarations=nonlocal_declarations,
        opts=opts,
        symbol_names=tuple(gast.Constant(str(s), kind=None) for s in loop_vars),
        state_functions=state_functions,
        state_getter_name=state_getter_name,
        state_setter_name=state_setter_name,
        undefined_assigns=undefined_assigns)
    origin_info.copy_origin(node, new_nodes[-1])
    return new_nodes


class AnnotatedDef(reaching_definitions.Definition):

  def __init__(self):
    super(AnnotatedDef, self).__init__()
    self.directives = {}


def transform(node, ctx):
  graphs = cfg.build(node)
  node = qual_names.resolve(node)
  node = activity.resolve(node, ctx, None)
  node = reaching_definitions.resolve(node, ctx, graphs)
  node = reaching_fndefs.resolve(node, ctx, graphs)
  node = liveness.resolve(node, ctx, graphs)

  node = ControlFlowTransformer(ctx).visit(node)
  return node