# Copyright 2024 The Flax Authors.
#
# 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.

"""Wrapper around jax.lax.scan with in_axes/out_axes API."""
import functools
from typing import Any, Optional
from collections.abc import Callable

import jax
import jax.numpy as jnp
import numpy as np
from jax import core, lax
from jax.extend import linear_util as lu
from jax.interpreters import partial_eval as pe

ScanAxis = Optional[int]


class _Broadcast:
  pass


broadcast = _Broadcast()


def scan(
  fn: Callable[..., Any],
  in_axes: Any,
  out_axes: Any,
  length: int | None = None,
  reverse: bool = False,
  unroll: int = 1,
  _split_transpose: bool = False
):
  """A wrapper around `jax.lax.scan` with in_axes/out_axes api.

  Example::
    def body_fn(b, c, x):
      return b + 2, c + 1, 2 * x

    loop = scan(body_fn, in_axes=0, out_axes=0)
    broadcast_in = 1
    carry = 2
    xs = jnp.arange(3)
    broadcast_out, carry, ys = loop(broadcast_in, carry, xs)
    print(broadcast_out)  # prints: 3
    print(carry)  # prints: 5
    print(ys)  # prints: [0, 2, 4]


  Args:
    fn: the body function of the scan loop of the form
      `(broadcast_in, carry, *args) -> (broadcast_out, carry, scan_out)`.
      the broadcast argument allows for loop independent inputs/outputs to
      be computed inside `fn`. `fn` will be called once to compute
      `broadcast_out`. The actual loop will receive `broadcast_out` as the new
      `broadcast_in`. This is useful for initializing values inside the loop.
    in_axes: specifies the axis along which arguments are scanned.
      Use `broadcast` to use the same value across iterations.
    out_axes: specifies the axis along which outputs are concatenated.
      Use `broadcast` if a return value should not be concatenated and
      is independent of the loop body.
    length: number of iterations. Only needs to be specified if there
      is no scan axis from which it can be derived.
    reverse: scan in reverse order from end to start.
    unroll: how many scan iterations to unroll within a single
      iteration of a loop (default: 1).
    _split_transpose: An experimental feature to split the transpose of scan
       into a scan and a map, backed by an experimental Jax lax.scan() feature.
  Returns:
     the function that performs the scan of the form:
     (broadcast_in, carry_in, *args) -> (broadcast_out, carry_out, scan_out).
  """

  def transpose_to_front(ax, xs):
    if ax is broadcast:
      return ()
    if ax == 0:
      return xs

    def trans(x):
      perm = tuple(range(x.ndim))
      perm = (ax,) + tuple(np.delete(perm, ax))
      return jnp.transpose(x, perm)

    return jax.tree_util.tree_map(trans, xs)

  def transpose_from_front(ax, xs):
    if ax is broadcast:
      return ()
    if ax == 0:
      return xs

    def trans(x):
      if ax < 0:
        pax = x.ndim + ax
      else:
        pax = ax
      assert pax < x.ndim
      perm = tuple(range(1, pax + 1)) + (0,) + tuple(range(pax + 1, x.ndim))
      return jnp.transpose(x, perm)

    return jax.tree_util.tree_map(trans, xs)

  def scan_fn(broadcast_in, init, *args):
    xs = jax.tree_util.tree_map(transpose_to_front, in_axes, args)

    def body_fn(c, xs, init_mode=False):
      # inject constants
      xs = jax.tree_util.tree_map(
        lambda ax, arg, x: (arg if ax is broadcast else x), in_axes, args, xs
      )
      broadcast_out, c, ys = fn(broadcast_in, c, *xs)

      if init_mode:
        ys = jax.tree_util.tree_map(
          lambda ax, y: (y if ax is broadcast else ()), out_axes, ys
        )
        return broadcast_out, ys
      else:
        ys = jax.tree_util.tree_map(
          lambda ax, y: (() if ax is broadcast else y), out_axes, ys
        )
        return c, ys

    broadcast_body = functools.partial(body_fn, init_mode=True)

    carry_avals = jax.tree_util.tree_map(
      lambda x: core.ShapedArray(jnp.shape(x), jnp.result_type(x)), init
    )
    scan_avals = jax.tree_util.tree_map(
      lambda x: core.ShapedArray(jnp.shape(x)[1:], jnp.result_type(x)), xs
    )
    input_avals = (carry_avals, scan_avals)

    in_avals, in_tree = jax.tree_util.tree_flatten(input_avals)
    f_flat, out_tree = jax.api_util.flatten_fun_nokwargs(
      lu.wrap_init(broadcast_body), in_tree
    )
    in_pvals = list(map(pe.PartialVal.unknown, in_avals))
    _, out_pvals, _ = pe.trace_to_jaxpr_nounits(f_flat, in_pvals)

    out_flat = []
    for pv, const in out_pvals:
      if pv is not None:
        raise ValueError(
          'broadcasted variable has a data dependency on the scan body.'
        )
      out_flat.append(const)
    broadcast_in, constants_out = jax.tree_util.tree_unflatten(
      out_tree(), out_flat
    )

    if jax.version.__version_info__ > (0, 4, 25):
      c, ys = lax.scan(
        body_fn, init, xs, length=length, reverse=reverse, unroll=unroll,
        _split_transpose=_split_transpose
      )
    else:
      c, ys = lax.scan(
        body_fn, init, xs, length=length, reverse=reverse, unroll=unroll
      )
    ys = jax.tree_util.tree_map(transpose_from_front, out_axes, ys)
    ys = jax.tree_util.tree_map(
      lambda ax, const, y: (const if ax is broadcast else y),
      out_axes,
      constants_out,
      ys,
    )
    return broadcast_in, c, ys

  return scan_fn
