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

"""Utils for orbax tests."""

from concurrent import futures
import contextlib
import functools
import inspect
import string
import time
import typing
from typing import Any, List, Optional, Set, Tuple
from unittest import mock

from absl import logging
from etils import epath
import jax
from jax.experimental import pjit
from jax.experimental.array_serialization import serialization
import jax.numpy as jnp
import numpy as np
from orbax.checkpoint import async_checkpoint_handler
from orbax.checkpoint import checkpoint_args
from orbax.checkpoint import multihost
from orbax.checkpoint import pytree_checkpoint_handler
from orbax.checkpoint import tree as tree_utils
from orbax.checkpoint import type_handlers
from orbax.checkpoint import utils
from orbax.checkpoint.multihost import multislice


def sync_global_processes(name: str):
  jax.experimental.multihost_utils.sync_global_devices(name)


def save_fake_tmp_dir(
    directory: epath.Path,
    step: int,
    item: str,
    subdirs: Optional[List[str]] = None,
    step_prefix: Optional[str] = None,
) -> epath.Path:
  """Saves a directory with a tmp folder to simulate preemption."""
  subdirs = subdirs or []
  if not step_prefix:
    step_prefix = ''
  step_tmp_dir = utils.create_tmp_directory(
      directory / (step_prefix + str(step))
  )
  item_tmp_dir = utils.create_tmp_directory(step_tmp_dir / item)
  if multihost.process_index() == 0:
    for sub in subdirs:
      (item_tmp_dir / sub).mkdir(parents=True)
  sync_global_processes('save_fake_tmp_dir')
  return item_tmp_dir


def replicate_sharded_array(arr: jax.Array):
  """Returns the input array, but replicated across all devices."""
  mesh = jax.sharding.Mesh(np.asarray(jax.devices()), ('x',))
  replicated_sharding = jax.sharding.NamedSharding(
      mesh,
      jax.sharding.PartitionSpec(
          None,
      ),
  )
  return pjit.pjit(lambda x: x, out_shardings=replicated_sharding)(arr)


def apply_function(tree, function):
  """Applies the given function to every leaf in tree.

  Args:
    tree: a nested dict where every leaf is a sharded jax.Array.
    function: a function accepting an array and returning jax.Array.

  Returns:
    a transformed sharded array tree.
  """

  def f(arr):
    return pjit.pjit(function)(arr)

  return jax.tree.map(f, tree)


def assert_array_equal(testclass, v_expected, v_actual):
  """Asserts that two arrays are equal."""
  testclass.assertIsInstance(v_actual, type(v_expected))
  if isinstance(v_expected, jax.Array):
    testclass.assertEqual(
        len(v_expected.addressable_shards), len(v_actual.addressable_shards)
    )
    for shard_expected, shard_actual in zip(
        v_expected.addressable_shards, v_actual.addressable_shards
    ):
      np.testing.assert_array_equal(shard_expected.data, shard_actual.data)
  elif isinstance(v_expected, (np.ndarray, jnp.ndarray)):
    np.testing.assert_array_equal(v_expected, v_actual)
  else:
    testclass.assertEqual(v_expected, v_actual)


def assert_tree_equal(testclass, expected, actual):
  """Asserts that two PyTrees are equal."""
  expected_flat = tree_utils.to_flat_dict(expected)
  actual_flat = tree_utils.to_flat_dict(actual)
  testclass.assertSameElements(expected_flat.keys(), actual_flat.keys())
  jax.tree.map(
      functools.partial(assert_array_equal, testclass), expected, actual
  )


def setup_pytree(add: int = 0):
  """Creates a numpy PyTree for testing."""
  pytree = {
      'a': np.arange(8) * 1,
      'b': np.arange(16) * 2,
      'c': {
          'a': np.arange(8).reshape((2, 4)) * 3,
          'e': np.arange(16).reshape((4, 4)) * 4,
      },
  }
  pytree = jax.tree.map(lambda x: x + add, pytree, is_leaf=is_leaf)
  return pytree


def setup_sharded_pytree(
    pytree: Optional[pytree_checkpoint_handler.PyTree] = None,
    reverse_devices: bool = False,
):
  """Creates a PyTree of sharded arrays for testing."""
  devices = jax.devices()
  num_devices = len(devices)
  if reverse_devices:
    devices = np.asarray(list(reversed(devices)))
  else:
    devices = np.asarray(devices)

  mesh_2d = jax.sharding.Mesh(
      devices.reshape((2, num_devices // 2)), ('x', 'y')
  )
  mesh_axes_2d = jax.sharding.PartitionSpec('x', 'y')
  mesh_1d = jax.sharding.Mesh(devices, ('x',))
  mesh_axes_1d = jax.sharding.PartitionSpec(
      'x',
  )
  mesh_0d = jax.sharding.Mesh(devices, ('x',))
  mesh_axes_0d = jax.sharding.PartitionSpec(
      None,
  )

  if pytree is None:
    pytree = setup_pytree()
  mesh_tree = {
      'a': mesh_0d,
      'b': mesh_1d,
      'c': {
          'a': mesh_2d,
          'e': mesh_2d,
      },
  }
  axes_tree = {
      'a': mesh_axes_0d,
      'b': mesh_axes_1d,
      'c': {
          'a': mesh_axes_2d,
          'e': mesh_axes_2d,
      },
  }

  pytree = jax.tree.map(
      create_sharded_array, pytree, mesh_tree, axes_tree, is_leaf=is_leaf
  )
  return pytree, mesh_tree, axes_tree


def setup_replica_sharded_arrays(
    arrays: List[jax.Array],
    mesh_shape: Tuple[int, ...],
    is_replica_first: Optional[bool] = True,
):
  """Creates a tuple of sharded arrays for testing."""
  devices = jax.devices()
  devices = np.asarray(devices)

  dim = len(mesh_shape)
  axis_names = list(string.ascii_lowercase)
  mesh = jax.sharding.Mesh(devices.reshape(mesh_shape), axis_names[-dim:])
  if is_replica_first:
    mesh_axes = jax.sharding.PartitionSpec(None, axis_names[-dim+1:])
  else:
    mesh_axes = jax.sharding.PartitionSpec(axis_names[-dim:-1], None)

  sharded_arrs = [create_sharded_array(arr, mesh, mesh_axes) for arr in arrays]
  return sharded_arrs, mesh, mesh_axes


def get_fake_global_mesh_for_slices(
    slice_processes: List[Set[int]],
) -> jax.sharding.Mesh:
  """Creates a "multi-slice" global mesh for testing.

  Args:
    slice_processes: List of sets of process indices, where each element in the
      list is a set of processes that are active in a single slice.

  Returns:
    A global mesh.
  """
  devices = jax.devices()
  slice_devices = []
  devices_per_slices = None
  all_processes = set()
  for processes in slice_processes:
    all_processes |= processes
    slice_devices.append([
        d
        for d in devices
        if multihost.utils.runtime_to_distributed_process_id(d.process_index)
        in processes
    ])
    devices_per_slices = devices_per_slices or len(slice_devices[-1])
    if len(slice_devices[-1]) != devices_per_slices:
      raise ValueError('All slices must have the same number of devices.')
  if len(all_processes) != jax.process_count():
    raise ValueError('All processes must be accounted for.')

  slice_devices = np.asarray(slice_devices)
  return jax.sharding.Mesh(slice_devices, ('replica', 'data'))


def select_single_replica(
    arrays: List[jax.Array], global_mesh: jax.sharding.Mesh
) -> List[jax.Array]:
  """Returns arrays sharded over single slice."""
  slice_devices = multislice.local_slice_devices(global_mesh.devices)
  # slice_devices = global_mesh.devices[0]
  single_slice_mesh = jax.sharding.Mesh(
      slice_devices, global_mesh.axis_names[1:]
  )

  def _get_single_slice_sharding(arr: jax.Array):
    sharding = typing.cast(jax.sharding.NamedSharding, arr.sharding)
    return jax.sharding.NamedSharding(
        single_slice_mesh,
        jax.sharding.PartitionSpec(*sharding.spec),  # exclude 'replica' axis
    )

  single_slice_shardings = jax.tree.map(
      _get_single_slice_sharding,
      arrays,
  )

  def _make_single_slice_array(
      arr: jax.Array, single_slice_sharding: jax.sharding.NamedSharding
  ):
    data = [
        arr.addressable_data(idx) for idx in range(len(arr.addressable_shards))
    ]
    return jax.make_array_from_single_device_arrays(
        arr.shape, single_slice_sharding, data
    )

  return jax.tree.map(
      _make_single_slice_array, arrays, single_slice_shardings
  )


def is_leaf(x):
  return (
      isinstance(x, np.ndarray)
      or isinstance(x, jax.sharding.Mesh)
      or isinstance(x, jax.sharding.PartitionSpec)
      or isinstance(x, pytree_checkpoint_handler.ParamInfo)
  )


def create_sharded_array(arr, mesh, mesh_axes):
  """Create sharded jax.Array."""
  if isinstance(arr, (int, float)):
    arr = np.asarray(arr)
  return jax.make_array_from_callback(
      arr.shape,
      jax.sharding.NamedSharding(mesh, mesh_axes),
      lambda idx: arr[idx],
  )


def print_directory(directory: epath.PathLike, level: int = 0):
  """Prints a directory tree for debugging purposes."""
  directory = epath.Path(directory)
  assert directory.is_dir()
  level_str = '..' * level
  if level == 0:
    logging.info('Printing directory tree: %s/', directory)
  else:
    logging.info('%s%s/', level_str, directory.name)

  level_str = '..' * (level + 1)
  for p in directory.iterdir():
    if p.is_dir():
      print_directory(p, level=level + 1)
    else:
      logging.info('%s%s', level_str, p.name)


def erase_and_create_empty(directory: epath.PathLike) -> epath.Path:
  """Creates an empty directory, erasing existing if present.

  Generally should be avoided in production code, as the heedless deletion is
  not very safe.

  Args:
    directory: directory to create.

  Returns:
    The created directory.
  """
  directory = epath.Path(directory)
  if directory.exists():
    directory.rmtree()
  directory.mkdir()
  return directory


def empty_directory(directory: epath.PathLike) -> bool:
  directory = epath.Path(directory)
  if not directory.exists():
    return False
  for p in directory.iterdir():
    if p.is_dir():
      p.rmtree()
    else:
      p.unlink()
  return True


def set_tensorstore_driver_for_test():
  # Sets TS driver for testing. Within Google, this defaults to `gfile`, which
  # results in issues writing to the OCDBT manifest. When using `gfile` on the
  # local filesystem, write operations are not atomic.
  serialization._DEFAULT_DRIVER = 'file'  # pylint: disable=protected-access


class ErrorCheckpointHandler(async_checkpoint_handler.AsyncCheckpointHandler):
  """Wrapper for PyTreeCheckpointHandler that has an error during save."""

  def __init__(
      self,
      handler: pytree_checkpoint_handler.PyTreeCheckpointHandler,
      executor: futures.ThreadPoolExecutor,
  ):
    self._handler = handler
    self._executor = executor

  def save(self, directory: epath.Path, args: 'ErrorSaveArgs'):
    return self._handler.save(directory, args=args)

  def restore(self, directory: epath.Path, args: 'ErrorRestoreArgs'):
    return self._handler.restore(directory, args=args)

  async def async_save(self, directory: epath.Path, args: 'ErrorSaveArgs'):
    commit_futures = await self._handler.async_save(directory, args=args)

    def error_commit():
      time.sleep(3)  # Pretend to write data.
      a = 1
      b = 2
      if a != b:
        raise SystemError()
      return 42

    return commit_futures + [self._executor.submit(error_commit)]

  def finalize(self, directory: epath.Path):
    self._handler.finalize(directory)


@checkpoint_args.register_with_handler(ErrorCheckpointHandler, for_save=True)
class ErrorSaveArgs(pytree_checkpoint_handler.PyTreeSaveArgs):
  pass


@checkpoint_args.register_with_handler(ErrorCheckpointHandler, for_restore=True)
class ErrorRestoreArgs(pytree_checkpoint_handler.PyTreeSaveArgs):
  pass


@contextlib.contextmanager
def register_type_handler(ty, handler, func):
  """Registers new func for type, and restores original handler when done."""
  original_handler = type_handlers.get_type_handler(ty)
  type_handlers.register_type_handler(ty, handler, func=func, override=True)
  try:
    yield
  finally:
    type_handlers.register_type_handler(
        ty, original_handler, func=func, override=True
    )


@contextlib.contextmanager
def ocdbt_checkpoint_context(use_ocdbt: bool, ts_context: Any):
  """Use OCDBT driver within context."""
  original_registry = list(
      type_handlers.GLOBAL_TYPE_HANDLER_REGISTRY._type_registry  # pylint: disable=protected-access
  )
  if use_ocdbt:
    type_handlers.register_standard_handlers_with_options(
        use_ocdbt=use_ocdbt, ts_context=ts_context
    )
  try:
    yield
  finally:
    type_handlers.GLOBAL_TYPE_HANDLER_REGISTRY._type_registry = (  # pylint: disable=protected-access
        original_registry
    )


def _get_wrapper_function(func):
  """Creates a function to wrap a test method and return a unique barrier key."""

  def _get_unique_barrier_key(key: str, func_id: str) -> str:
    return f'{key}.{func_id}'

  def wrapper(self, *args, **kwargs):
    with mock.patch.object(
        multihost.utils,
        '_unique_barrier_key',
        new=functools.partial(_get_unique_barrier_key, func_id=self.id()),
    ):
      return func(self, *args, **kwargs)

  return wrapper


def subset_barrier_compatible_test(cls):
  """A decorator to be used with a test class.

  E.g.

  @subset_barrier_compatible_test
  class MyTest(googletest.TestCase):
    def test_foo(self):
      ...

  The point of this decorator is to modify all functions to mock the private
  method `multihost.utils._unique_barrier_key`, to append the test case name.
  This allows multiple test cases to reuse barrier names that would otherwise
  conflict.

  Args:
    cls: the test class to decorate.

  Returns:
    The decorated class.
  """
  for name, func in inspect.getmembers(cls, predicate=inspect.isfunction):
    if name.startswith('test'):
      setattr(cls, name, _get_wrapper_function(func))
  return cls
