# Copyright 2021 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.
# ==============================================================================
"""Jax recurrent layer implementation.

The main interface of this module is recurrent_func().
This expects the caller to describe the recurrent neural net by specifying:

  - theta: the "weights" each RNN uses.
  - states_0: the initial state of each RNN.
  - cell_fn: A python function describing RNN cell. It must have the following
    signature::

        cell_fn: (theta, states_0, inputs) -> states_1

    states_1 is the next RNN state.

recurrent_func computes, roughly::

    state = states_0
    t = 0
    while t < seq_length:
      state = cell_fn(theta, state, inputs[t, :])
      accumulate_state[t, :] = state
      t += 1
    return state, accumulate_state
"""

import enum
import functools
from typing import Callable, Optional, Tuple

import jax
from jax import ad_checkpoint
from jax import numpy as jnp
from lingvo.jax import base_layer
from lingvo.jax import py_utils
from lingvo.jax import pytypes
import tensorflow.compat.v2 as tf

NestedMap = py_utils.NestedMap
WeightInit = py_utils.WeightInit

ParamsT = pytypes.ParamsT
JTensor = pytypes.JTensor
CallableOrNone = Optional[Callable]
NestedMapOrNone = Optional[NestedMap]


@enum.unique
class AutodiffCheckpointType(str, enum.Enum):
  """jax.checkpoint policy types."""
  SAVE_NOTHING = 'save_nothing'
  SAVE_EVERYTHING = 'save_everything'
  SAVE_QKV_OUT_PROJ = 'save_qkv_out_proj'
  SAVE_OUT_PROJ = 'save_out_proj'
  SAVE_CONTEXT = 'save_context'
  SAVE_CONTEXT_AND_OUT_PROJ = 'save_encoded_and_out_proj'
  SAVE_DOT_ONLY = 'save_dot_only'
  SAVE_DOT_WITH_NO_BATCH_DIM = 'save_dot_with_no_batch_dims'
  SAVE_DOT_FOR_MLPERF_200B = 'save_dot_for_mlperf_200b'


def recurrent_func(theta: NestedMap, states_0: NestedMap, inputs: NestedMap,
                   cell_fn: Callable[[NestedMap, NestedMap, NestedMap],
                                     NestedMap]):
  """Computes a recurrent neural net.

  Args:
    theta: weights. A `.NestedMap`.
    states_0: initial state. A `.NestedMap`.
    inputs: inputs. A `.NestedMap`.
    cell_fn: A python function which computes::
        states_1 = cell_fn(theta, states_0, inputs[t, :])

  Returns:
    `accumulate_state` and the final state.
  """
  input_seq_len = inputs.Flatten()[0].shape[0]

  def assert_not_none(x):
    assert x is not None

  tf.nest.map_structure(assert_not_none, states_0)
  tf.nest.map_structure(assert_not_none, inputs)
  tf.nest.map_structure(assert_not_none, theta)

  def new_cum_state(x):
    x1 = jnp.expand_dims(x, 0)
    # +1 so that we can store initial_states at position 0.
    return jnp.tile(x1, [input_seq_len + 1] + [1] * x.ndim)

  cumulative_states = states_0.Transform(new_cum_state)

  prng_key = base_layer.next_prng_key()
  global_step = base_layer.cur_global_step()

  start_time = jnp.array(0, dtype=jnp.uint32)
  fwd_initial_loop_vars = NestedMap(
      cur_time=start_time,
      theta=theta,
      states_0=states_0,
      cumulative_states=cumulative_states,
      inputs=inputs)

  def same_type_shape(x, y):
    assert x.dtype == y.dtype, (x.dtype, y.dtype)
    assert x.shape == y.shape, (x.shape, y.shape)

  def wrapped_cell_fn(fn_in):
    # fn_in is NestedMap containing the following elements:
    #    - t
    #    - theta
    #    - states_0
    #    - inputs_t
    # Start a chain of prng key that also takes into account of time steps.
    t = fn_in.t
    theta = fn_in.theta
    states_0 = fn_in.states_0
    inputs_t = fn_in.inputs_t
    with base_layer.JaxContext.new_context(
        prng_key=jax.random.fold_in(prng_key, t), global_step=global_step):
      # NO side-effect ops are allowed as the enclosing JaxContext is not bound
      # to any layer.
      states_1 = cell_fn(theta, states_0, inputs_t)

      tf.nest.assert_same_structure(states_0, states_1)
      tf.nest.map_structure(same_type_shape, states_0, states_1)
    return states_1

  def wrapped_cell_fn_grad(fn_in, d_fn_out):
    # This is roughly the following:
    #
    # fn_out = wrapped_cell_fn(fn_in)
    # d_fn_in = tf.gradient(fn_out, fn_in, d_fn_out)
    # return d_fn_in
    #
    assert isinstance(fn_in, NestedMap)
    fn_out, vjp_fn = jax.vjp(wrapped_cell_fn, fn_in)
    del fn_out
    d_fn_in = vjp_fn(d_fn_out)
    assert isinstance(d_fn_in, tuple)
    assert len(d_fn_in) == 1
    d_fn_in_0 = d_fn_in[0]
    # Over-write gradient for t, the time step.
    d_fn_in_0.t = jnp.zeros_like(fn_in.t)
    tf.nest.assert_same_structure(fn_in, d_fn_in_0)
    tf.nest.map_structure(same_type_shape, fn_in, d_fn_in_0)
    return d_fn_in_0

  def fwd_comp_fn(loop_vars):
    # loop_vars is a NestedMap containing the following elements:
    #   - cur_time
    #   - theta
    #   - inputs
    #   - cumulative_states
    #   - states_0
    t = loop_vars.cur_time
    theta = loop_vars.theta
    inputs = loop_vars.inputs
    cumulative_states = loop_vars.cumulative_states
    states_0 = loop_vars.states_0
    inputs_t = inputs.Transform(lambda x: x[t])

    states_1 = wrapped_cell_fn(
        NestedMap(t=t, theta=theta, states_0=states_0, inputs_t=inputs_t))

    def set_t(x, x_t):
      return x.at[t + 1].set(x_t)

    cumulative_states = tf.nest.map_structure(set_t, cumulative_states,
                                              states_1)
    loop_out = NestedMap(
        cur_time=t + 1,
        theta=theta,
        inputs=inputs,
        states_0=states_1,
        cumulative_states=cumulative_states)
    return loop_out

  def fwd_continue_fn(loop_vars):
    return loop_vars.cur_time < input_seq_len

  # This custom_vjp implementation follows examples here:
  # https://jax.readthedocs.io/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html
  @jax.custom_vjp
  def fwd_loop(loop_vars):
    final_loop_vars = jax.lax.while_loop(fwd_continue_fn, fwd_comp_fn,
                                         loop_vars)
    return NestedMap(
        final_states=final_loop_vars.states_0,
        cumulative_states=final_loop_vars.cumulative_states)

  def loop_fn_vjp_fwd(loop_vars):
    loop_fn_out = fwd_loop(loop_vars)
    return loop_fn_out, (loop_vars, loop_fn_out.cumulative_states)

  def loop_fn_vjp_bwd(res, d_out):
    fwd_loop_vars, cumulative_states = res
    d_final_states = d_out.final_states
    d_cumulative_states = d_out.cumulative_states

    start_time = input_seq_len - 1
    d_states_1 = tf.nest.map_structure(lambda x, y: x[start_time + 1] + y,
                                       d_cumulative_states, d_final_states)
    bwd_loop_vars = NestedMap(
        cur_time=start_time,
        theta=fwd_loop_vars.theta,
        inputs=fwd_loop_vars.inputs,
        cumulative_states=cumulative_states,
        d_cumulative_states=d_cumulative_states,
        d_theta=fwd_loop_vars.theta.Transform(jnp.zeros_like),
        d_inputs=fwd_loop_vars.inputs.Transform(jnp.zeros_like),
        d_states_1=d_states_1)

    def bwd_comp_fn(loop_vars):
      t = loop_vars.cur_time
      inputs = loop_vars.inputs
      inputs_t = inputs.Transform(lambda x: x[t])
      states_0 = loop_vars.cumulative_states.Transform(lambda x: x[t])
      d_cell_in = wrapped_cell_fn_grad(
          NestedMap(
              t=t, theta=loop_vars.theta, states_0=states_0, inputs_t=inputs_t),
          loop_vars.d_states_1)
      d_theta = tf.nest.map_structure(lambda x, y: x + y, loop_vars.d_theta,
                                      d_cell_in.theta)
      d_states_0 = tf.nest.map_structure(lambda x, y: x + y[t],
                                         d_cell_in.states_0,
                                         loop_vars.d_cumulative_states)

      def set_t(x, x_t):
        return x.at[t].set(x_t)

      d_inputs = tf.nest.map_structure(set_t, loop_vars.d_inputs,
                                       d_cell_in.inputs_t)
      loop_vars_out = loop_vars.Transform(lambda x: x)
      loop_vars_out.d_inputs = d_inputs
      loop_vars_out.d_states_1 = d_states_0
      loop_vars_out.d_theta = d_theta
      loop_vars_out.cur_time = t - 1
      return loop_vars_out

    def bwd_continue_fn(loop_vars):
      return loop_vars.cur_time >= 0

    bwd_final_loop_vars = jax.lax.while_loop(bwd_continue_fn, bwd_comp_fn,
                                             bwd_loop_vars)
    d_out = fwd_loop_vars.Transform(jnp.zeros_like)

    tf.nest.map_structure(same_type_shape, d_out.states_0,
                          bwd_final_loop_vars.d_states_1)
    tf.nest.map_structure(same_type_shape, d_out.theta,
                          bwd_final_loop_vars.d_theta)
    tf.nest.map_structure(same_type_shape, d_out.inputs,
                          bwd_final_loop_vars.d_inputs)

    d_out.states_0 = bwd_final_loop_vars.d_states_1
    d_out.theta = bwd_final_loop_vars.d_theta
    d_out.inputs = bwd_final_loop_vars.d_inputs
    return (d_out,)

  fwd_loop.defvjp(loop_fn_vjp_fwd, loop_fn_vjp_bwd)

  # Finally, let's simply run the forward loop fn.
  fwd_final_loop_vars = fwd_loop(fwd_initial_loop_vars)
  fwd_cumulative_states = fwd_final_loop_vars.cumulative_states.Transform(
      lambda x: x[1:])
  return fwd_final_loop_vars.final_states, fwd_cumulative_states


def recurrent_static(theta: NestedMap,
                     states_0: NestedMap,
                     inputs: NestedMap,
                     cell_fn: Callable[[NestedMap, NestedMap, NestedMap],
                                       NestedMap],
                     root_layer: Optional[base_layer.BaseLayer] = None):
  """A simpler form of Recurrent where num of steps is known statically.

  Back-prop is availale through auto-diff.

  'padding' in inputs is used to skip certain steps dynamically. If the
  'padding' tensor exists, it is expected of a binary 0/1 tensor.

  Args:
    theta: weights. A `.NestedMap`.
    states_0: initial state. A `.NestedMap`.
    inputs: inputs. A `.NestedMap`. All inputs in time-major.
    cell_fn: A python function which computes::
        states_1 = cell_fn(theta, states_0, inputs[t, :])
    root_layer: The root layer within which this recurrent_static recurrent loop
      is carried out.

  Returns:
    `accumulate_state` and the final state.
  """

  assert 'time_step' not in states_0
  # The initial time step.
  time_step = jnp.array(0, dtype=jnp.uint32)
  # Make a copy of states_0 structure.
  states_0 = tf.nest.map_structure(lambda x: x, states_0)
  states_0.time_step = time_step

  prng_key = base_layer.next_prng_key()
  global_step = base_layer.cur_global_step()

  # TODO(zhangqiaorjc): Switch to ad_checkpoint.checkpoint after mattjj bug fix.
  @jax.checkpoint
  def comp_fn(states_0, inputs_t):
    # Start a new prng_key branch that also depends on the time step.
    if root_layer is not None:
      forward_updated_vars_before = tf.nest.map_structure(
          lambda x: x, root_layer.forward_updated_vars)
    prng_key_t = jax.random.fold_in(prng_key, states_0.time_step)
    with base_layer.JaxContext.new_context(
        prng_key=prng_key_t, global_step=global_step):
      # NO side-effect ops are allowed as the enclosing JaxContext is not bound
      # to any layer.
      #
      # Whether or not we should skip this time step.
      if 'padding' in inputs_t:
        # We skip if all are padded steps.
        skip = jnp.all(inputs_t.padding > 0.5)
      else:
        skip = jnp.array(False)

      def carry_over(args):
        states_0, inputs_t = args
        del inputs_t
        # We simply carry over the states for this time step.
        states_1 = tf.nest.map_structure(lambda x: x, states_0)
        states_1.time_step = states_0.time_step + 1
        return states_1

      def do_compute(args):
        states_0, inputs_t = args
        # Actually carry out the computation.
        states_1 = cell_fn(theta, states_0, inputs_t)
        states_1.time_step = states_0.time_step + 1
        return states_1

      if 'padding' in inputs_t:
        states_1 = jax.lax.cond(skip, carry_over, do_compute,
                                (states_0, inputs_t))
      else:
        states_1 = do_compute((states_0, inputs_t))
      tf.nest.assert_same_structure(states_0, states_1)

      if root_layer is not None:
        forward_updated_vars_after = tf.nest.map_structure(
            lambda x: x, root_layer.forward_updated_vars)

        def assert_no_change(x, y):
          assert (x is None and y is None) or (x is not None and y is not None)

        tf.nest.map_structure(assert_no_change, forward_updated_vars_before,
                              forward_updated_vars_after)

      return states_1, states_1

  final_states, cumulative_states = jax.lax.scan(comp_fn, states_0, inputs)
  del final_states.time_step
  del cumulative_states.time_step
  return final_states, cumulative_states


def scan(carry_init: NestedMap,
         xs: NestedMap,
         fn: Callable[[NestedMap, NestedMap], Tuple[NestedMap, NestedMap]],
         root_layer: Optional[base_layer.BaseLayer] = None,
         checkpoint_policy: AutodiffCheckpointType = AutodiffCheckpointType
         .SAVE_NOTHING):
  """A simple wrap around jax.lax.scan.

  Back-prop is availale through auto-diff.

  Args:
    carry_init: initial state. A `.NestedMap`.
    xs: inputs. A `.NestedMap`. All inputs in time-major.
    fn: A python function which computes:
        carry, ys[t] = fn(carry, xs[t, :])
    root_layer: The root layer within which this jax.lax.scan based while_loop
      is carried out. If root_layer is provided, some basic-effort check is
      performed to make sure fn is side-effect free. Otherwise, no such checks
      are performed.
    checkpoint_policy: A AutodiffCheckpointType. How to checkpoint for BProp:
      SAVE_NOTHING, SAVE_DOT_ONLY, SAVE_DOT_WITH_NO_BATCH_DIM.

  Returns:
    (final 'carry', 'ys', stacked summaries).
  """
  del root_layer
  assert isinstance(carry_init, py_utils.NestedMap)
  assert isinstance(xs, py_utils.NestedMap)
  # Make a copy of carry_init structure.
  carry_init = tf.nest.map_structure(lambda x: x, carry_init)
  # "carry" will be augmented with the following three tensors, so make sure
  # they don't already exist in the NestedMap.
  assert 'time_step' not in carry_init
  assert 'prng_key' not in carry_init
  assert 'global_step' not in carry_init

  def custom_policy(checkpoint_policy: AutodiffCheckpointType):

    # TODO(zhangqiaorjc): Configure custom checkpoint policy in expt config
    # without introducing enum.
    if checkpoint_policy == AutodiffCheckpointType.SAVE_EVERYTHING:
      return jax.checkpoint_policies.everything_saveable
    if checkpoint_policy == AutodiffCheckpointType.SAVE_DOT_ONLY:
      return jax.checkpoint_policies.checkpoint_dots
    if checkpoint_policy == AutodiffCheckpointType.SAVE_DOT_WITH_NO_BATCH_DIM:
      return jax.checkpoint_policies.checkpoint_dots_with_no_batch_dims
    if checkpoint_policy == AutodiffCheckpointType.SAVE_QKV_OUT_PROJ:
      return jax.checkpoint_policies.save_only_these_names(
          'combined_qkv_proj', 'out_proj')
    if checkpoint_policy == AutodiffCheckpointType.SAVE_CONTEXT:
      return jax.checkpoint_policies.save_only_these_names('context')
    if checkpoint_policy == AutodiffCheckpointType.SAVE_OUT_PROJ:
      return jax.checkpoint_policies.save_only_these_names('out_proj')
    if checkpoint_policy == AutodiffCheckpointType.SAVE_CONTEXT_AND_OUT_PROJ:
      return jax.checkpoint_policies.save_only_these_names(
          'context', 'out_proj')
    if checkpoint_policy == AutodiffCheckpointType.SAVE_DOT_FOR_MLPERF_200B:
      return jax.checkpoint_policies.save_only_these_names(
          'combined_qkv_proj', 'query_proj', 'value_proj', 'key_proj',
          'context', 'out_proj')
    assert checkpoint_policy == AutodiffCheckpointType.SAVE_NOTHING
    return jax.checkpoint_policies.nothing_saveable

  @functools.partial(
      ad_checkpoint.checkpoint,
      prevent_cse=False,
      policy=custom_policy(checkpoint_policy))
  def fn_wrap(carry, xs_t):
    # carry is augmented with time_step, prng_key, global_step three additional
    # tensors to make fn_wrap fully functional.
    # Start a new prng_key branch that also depends on the time step.
    prng_key_t = jax.random.fold_in(carry.prng_key, carry.time_step)
    with base_layer.JaxContext.new_context(
        prng_key=prng_key_t, global_step=carry.global_step):

      carry_new, ys_t = fn(carry, xs_t)
      carry_new.time_step = carry.time_step + 1
      # copy over prng_key and global_step
      carry_new.prng_key = carry.prng_key
      carry_new.global_step = carry.global_step

      tf.nest.assert_same_structure(carry_new, carry)
      summaries = base_layer.all_summaries()

    return carry_new, (ys_t, summaries)

  # The initial time step.
  time_step = jnp.array(0, dtype=jnp.uint32)
  prng_key = base_layer.next_prng_key()
  global_step = base_layer.cur_global_step()
  carry_init.time_step = time_step
  carry_init.prng_key = prng_key
  carry_init.global_step = global_step

  carry_final, (ys, summaries) = jax.lax.scan(fn_wrap, carry_init, xs)

  del carry_final.time_step
  del carry_final.global_step
  del carry_final.prng_key
  return carry_final, ys, summaries
