#!/usr/bin/env python3

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""
Compare RL training losses against a committed reference file for regression
detection.  Follows the same pattern as ``scripts/loss_compare.py`` for
pre-training, but adapted for the Monarch-based RL launcher.

The script runs GRPO training with deterministic settings (seed=42,
deterministic ops) and extracts ``loss/mean`` from TensorBoard events.
It then compares the extracted losses step-by-step against a reference
file using exact floating-point equality.

The reference-file I/O helpers are reused from ``scripts/loss_compare.py``;
only the TensorBoard extraction is RL-specific (see below).  Exactly one of
``--export-result`` / ``--import-result`` must be given -- training is always
a means to either export a new reference or assert against an existing one.

Example usages:

1. Export reference losses (run once on a known-good setup):
   python scripts/rl/loss_compare.py \\
       --hf-assets-path /path/to/Qwen3-0.6B \\
       --dump-folder /tmp/rl_loss_guard \\
       --export-result tests/assets/losses/rl_grpo_cuda.txt

2. Assert losses match reference (CI mode):
   python scripts/rl/loss_compare.py \\
       --hf-assets-path /path/to/Qwen3-0.6B \\
       --dump-folder /tmp/rl_loss_guard \\
       --import-result tests/assets/losses/rl_grpo_cuda.txt \\
       --assert-equal
"""

import argparse
import os
import subprocess
import sys
import unittest
from typing import cast

LOG_PREFIX = "[RL_LOSS_COMPARE]"

TB_LOSS_TAG = "loss/mean"


def log_print(message: str = "") -> None:
    if message:
        print(f"{LOG_PREFIX} {message}")
    else:
        print(LOG_PREFIX)


# ---------------------------------------------------------------------------
# TensorBoard extraction
# ---------------------------------------------------------------------------


def extract_losses_from_tensorboard(tb_dir: str) -> dict[int, float]:
    """Extract ``loss/mean`` scalars from TensorBoard event files.

    Intentionally RL-specific (not imported from ``scripts/loss_compare.py``):
    the RL ``MetricsProcessor`` writes events directly under ``dump_folder``
    -- there is no per-run ``tb_baseline``/``tb_test`` timestamped subdir --
    and the RL loss is logged under the ``loss/mean`` tag rather than the
    pre-training ``loss_metrics/global_avg_loss`` tag.

    Args:
        tb_dir: Directory containing TensorBoard event files.  The RL
            ``MetricsProcessor`` writes events directly under ``dump_folder``,
            so this is typically the dump folder itself.

    Returns:
        Dictionary mapping step number to full-precision loss value.
    """
    from tensorboard.backend.event_processing.event_accumulator import EventAccumulator

    if not os.path.exists(tb_dir):
        raise FileNotFoundError(f"TensorBoard path does not exist: {tb_dir}")

    event_acc = EventAccumulator(tb_dir)
    event_acc.Reload()

    available_tags = cast(dict[str, list[str]], event_acc.Tags()).get("scalars", [])
    if TB_LOSS_TAG not in available_tags:
        raise KeyError(
            f"Scalar tag '{TB_LOSS_TAG}' not found in TensorBoard events at "
            f"{tb_dir}. Available tags: {available_tags}"
        )

    scalars = event_acc.Scalars(TB_LOSS_TAG)
    losses = {scalar.step: scalar.value for scalar in scalars}

    log_print(f"Extracted {len(losses)} steps from TensorBoard events")
    return losses


# Reference-file I/O is reused from ``scripts/loss_compare.py`` -- imported
# lazily in ``main`` so this module only needs the repo root on ``sys.path``
# when actually run. The shared helpers support multiple metrics, while this
# guard reads and writes only the ``loss`` metric.


# ---------------------------------------------------------------------------
# Training
# ---------------------------------------------------------------------------


def run_training(
    hf_assets_path: str,
    dump_folder: str,
    module: str,
    config: str,
) -> None:
    """Run a deterministic RL GRPO configuration as a subprocess."""
    env = os.environ.copy()
    env["TORCHTITAN_RL_LOSS_MODULE"] = module
    env["TORCHTITAN_RL_LOSS_CONFIG"] = config
    env["TORCHTITAN_RL_LOSS_HF_ASSETS_PATH"] = hf_assets_path
    cmd = [
        sys.executable,
        "-m",
        "torchtitan.rl.train",
        "--module",
        "scripts.rl.loss_compare",
        "--config",
        "rl_loss_compare_config",
        "--output-dir",
        dump_folder,
    ]

    log_print(f"Running: {' '.join(cmd)}")

    result = subprocess.run(cmd, env=env)
    if result.returncode != 0:
        log_print(f"Training failed with return code {result.returncode}")
        sys.exit(1)

    log_print("Training completed successfully")


def rl_loss_compare_config():
    """Build the deterministic RL configuration used by this script."""
    from torchtitan.config import ConfigLoader
    from torchtitan.rl.controller import Controller

    config = ConfigLoader._load_config(
        os.environ["TORCHTITAN_RL_LOSS_MODULE"],
        os.environ["TORCHTITAN_RL_LOSS_CONFIG"],
    )
    if not isinstance(config, Controller.Config):
        raise TypeError("RL loss comparison requires a Controller.Config recipe.")
    config.hf_assets_path = os.environ["TORCHTITAN_RL_LOSS_HF_ASSETS_PATH"]
    config.trainer.debug.deterministic = True
    config.generator.debug.deterministic = True
    config.trainer.debug.seed = 42
    config.generator.debug.seed = 42
    config.num_generators = 1
    config.trainer.parallelism.tensor_parallel_degree = 4
    config.generator.parallelism.tensor_parallel_degree = 4
    config.async_loop.num_training_steps = 10
    # Lockstep removes timing-dependent policy lag from the loss curve.
    config.async_loop.target_offpolicy_steps = 0
    config.async_loop.validation.num_samples = 0
    config.metrics.enable_tensorboard = True
    config.metrics.enable_wandb = False
    return config


# ---------------------------------------------------------------------------
# Assertion
# ---------------------------------------------------------------------------


def assert_losses_equal(
    actual: dict[int, float],
    expected: dict[int, float],
) -> None:
    """Assert exact floating-point equality between actual and expected losses."""
    log_print(
        f"Comparing {len(actual)} actual steps against {len(expected)} expected steps"
    )

    class LossEqualityTest(unittest.TestCase):
        def test_losses_equal(self):
            actual_steps = set(actual.keys())
            expected_steps = set(expected.keys())
            self.assertEqual(
                actual_steps,
                expected_steps,
                f"Steps mismatch: actual has {len(actual_steps)} steps, "
                f"expected has {len(expected_steps)} steps",
            )

            for step in sorted(actual_steps):
                self.assertEqual(
                    actual[step],
                    expected[step],
                    f"Loss mismatch at step {step}: "
                    f"actual={repr(actual[step])}, expected={repr(expected[step])}",
                )

    suite = unittest.TestLoader().loadTestsFromTestCase(LossEqualityTest)
    runner = unittest.TextTestRunner(verbosity=2)
    result = runner.run(suite)

    if not result.wasSuccessful():
        log_print("Loss assertion FAILED!")
        log_print()
        log_print(
            "Actual losses (use these to update the reference file if "
            "the change is intentional):"
        )
        log_print("WARNING: verify the loss curve change is not a regression first!")
        for step in sorted(actual.keys()):
            print(f"{step} {repr(actual[step])}")
        log_print()
        sys.exit(1)
    else:
        log_print("All losses match. Assertion passed!")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Compare RL training losses against a reference file.",
    )
    parser.add_argument(
        "--hf-assets-path",
        required=True,
        help="Path to HF model checkpoint (weights, tokenizer, config).",
    )
    parser.add_argument(
        "--dump-folder",
        required=True,
        help="Directory for training outputs and TensorBoard events.",
    )
    parser.add_argument(
        "--module",
        default="torchtitan_recipes.rl.alphabet_sort",
        help="RL config module (default: torchtitan_recipes.rl.alphabet_sort).",
    )
    parser.add_argument(
        "--config",
        default="rl_grpo_qwen3_0_6b_varlen_batch_invariant",
        help="RL config name (default: rl_grpo_qwen3_0_6b_varlen_batch_invariant).",
    )
    parser.add_argument(
        "--export-result",
        default="",
        help="Export losses to this file path after training.",
    )
    parser.add_argument(
        "--import-result",
        default="",
        help="Reference file to compare against (requires --assert-equal).",
    )
    parser.add_argument(
        "--assert-equal",
        action="store_true",
        help="Assert losses match the reference file exactly.",
    )
    args = parser.parse_args()

    # scripts/loss_compare.py (repo root) isn't on the editable-install
    # sys.path; add it, then reuse its reference-file I/O helpers.
    repo_root = os.path.abspath(
        os.path.join(os.path.dirname(__file__), *([os.pardir] * 3))
    )
    if repo_root not in sys.path:
        sys.path.insert(0, repo_root)
    from scripts.loss_compare import export_metrics_to_file, read_metrics_from_file

    export_result = args.export_result or None
    import_result = args.import_result or None

    if not export_result and not import_result:
        parser.error("one of --export-result / --import-result is required")

    if args.assert_equal and not import_result:
        parser.error("--assert-equal requires --import-result")

    if import_result:
        args.assert_equal = True

    # Run training
    run_training(args.hf_assets_path, args.dump_folder, args.module, args.config)

    # Extract losses from TensorBoard
    actual_losses = extract_losses_from_tensorboard(args.dump_folder)

    if not actual_losses:
        log_print("Error: no losses found in TensorBoard events")
        sys.exit(1)

    # Export if requested
    if export_result:
        export_metrics_to_file({"loss": actual_losses}, export_result)

    # Compare if requested
    if args.assert_equal:
        assert import_result is not None
        expected_losses = read_metrics_from_file(import_result)["loss"]
        assert_losses_equal(actual_losses, expected_losses)


if __name__ == "__main__":
    main()
