# Copyright 2025 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.

"""Performance benchmarks for orbax.checkpoint.CheckpointManager."""

from collections.abc import Sequence
import dataclasses
import time
from typing import Any

import jax
import numpy as np
import orbax.checkpoint as ocp
from orbax.checkpoint._src.checkpoint_managers import preservation_policy as preservation_policy_lib
from orbax.checkpoint._src.checkpoint_managers import save_decision_policy as save_decision_policy_lib
from orbax.checkpoint._src.testing.benchmarks.core import core as benchmarks_core
from orbax.checkpoint._src.testing.benchmarks.core import metric as metric_lib


# ==============================================================================
# 1. Define the Options Dataclass for this specific benchmark
# ==============================================================================
@dataclasses.dataclass(frozen=True)
class CheckpointManagerPerfBenchmarkOptions(benchmarks_core.BenchmarkOptions):
  """Configuration options for benchmarks targeting CheckpointManager.

  Each attribute can be a single value or a list of values to create
  a parameter sweep.

  Attributes:
    use_ocdbt: Whether to use OCPDBT for checkpointing.
    use_zarr3: Whether to use Zarr3 for checkpointing.
  """

  use_ocdbt: bool | Sequence[bool] = True
  use_zarr3: bool | Sequence[bool] = False
  train_steps: int = 20


# ==============================================================================
# 2. Implement the Benchmark Generator
# ==============================================================================
@benchmarks_core.benchmark_options(CheckpointManagerPerfBenchmarkOptions)
class CheckpointManagerPerfBenchmark(benchmarks_core.BenchmarksGenerator):
  """A concrete generator for `orbax.checkpoint.CheckpointManager`.

  This class provides the specific test logic for benchmarking the
  CheckpointManager with various configurations.
  """

  def _clear_pytree(self, pytree: Any) -> Any:
    """Clears the pytree to free up memory."""
    return jax.tree.map(
        lambda x: x.delete() if isinstance(x, jax.Array) else None, pytree
    )

  def _train_step(self, pytree):
    def increment(tree):
      return jax.tree.map(lambda x: x + 0.5, tree)

    train_fn = jax.jit(increment, donate_argnums=(0,))
    return train_fn(pytree)

  def test_fn(
      self, context: benchmarks_core.TestContext
  ) -> benchmarks_core.TestResult:
    """The core test logic for a single save/restore cycle.

    This function is called for each combination of options generated by the
    framework. It uses the `context.options` to configure the handler
    dynamically for each run.

    Args:
      context: The test context containing the pytree, path, and options.

    Returns:
      The test result containing the metrics.
    """
    metrics = metric_lib.Metrics()
    pytree = context.pytree
    save_path = context.path / 'checkpoint_manager_perf'
    options = context.options
    assert isinstance(options, CheckpointManagerPerfBenchmarkOptions)

    mngr = ocp.CheckpointManager(
        save_path,
        item_handlers={
            'pytree_item': ocp.PyTreeCheckpointHandler(
                use_ocdbt=options.use_ocdbt,
                use_zarr3=options.use_zarr3,
            )
        },
        options=ocp.CheckpointManagerOptions(
            save_decision_policy=save_decision_policy_lib.ContinuousCheckpointingPolicy(),
            preservation_policy=preservation_policy_lib.LatestN(n=1),
        ),
    )

    save_times = []
    total_save_times = []
    for i in range(options.train_steps):
      with metrics.time(f'save_{i}'):
        save_start = time.time()
        mngr.save(
            i,
            args=ocp.args.Composite(
                pytree_item=ocp.args.StandardSave(item=pytree)
            ),
        )
        save_times.append(time.time() - save_start)
        mngr.wait_until_finished()
        total_save_times.append(time.time() - save_start)
      with metrics.time(f'train_step_{i}'):
        pytree = self._train_step(pytree)

    save_times = np.array(save_times)
    total_save_times = np.array(total_save_times)

    # Exclude step 0 from assertions; setup may take extra time.
    asserting_save_times = save_times[1:]
    asserting_total_save_times = total_save_times[1:]

    mean_save_time = np.mean(asserting_save_times)
    mean_total_save_time = np.mean(asserting_total_save_times)

    assert np.all(asserting_save_times <= 2 * mean_save_time), (
        f'Save times={asserting_save_times}, mean save time={mean_save_time}'
    )
    assert np.all(asserting_total_save_times <= 2 * mean_total_save_time), (
        f'Total save times={asserting_total_save_times}, mean total save'
        f' time={mean_total_save_time}'
    )

    abstract_pytree = jax.tree.map(
        lambda x: ocp.utils.to_shape_dtype_struct(x)
        if isinstance(x, jax.Array)
        else x,
        pytree,
    )
    context.pytree = self._clear_pytree(context.pytree)

    with metrics.time('restore'):
      mngr.restore(
          mngr.latest_step(),
          args=ocp.args.Composite(
              pytree_item=ocp.args.StandardRestore(
                  item=abstract_pytree,
              ),
          ),
      )

    mngr.close()
    return benchmarks_core.TestResult(metrics=metrics)
