# 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 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,
        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.measure(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.measure(f'train_step_{i}'):
        pytree = self._train_step(pytree)

    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.measure('restore'):
      mngr.restore(
          mngr.latest_step(),
          args=ocp.args.Composite(
              pytree_item=ocp.args.StandardRestore(
                  item=abstract_pytree,
              ),
          ),
      )

    # TODO(nikhilbansall) : Add assertions for this test.

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