"""问题3/4 RL 策略与既有策略的配对自建消融。

同一种子、同一场景分别跑基线（Q3 bayes / Q4 directional）与 RL 变体，
比较全部清除成功、虚拟总时间、检测次数与请求总数。结果写入
output/rl/benchmark.json。

用法：
    uv sync --group rl --group dev
    uv run python scripts/benchmark_rl.py --seeds 5
"""

from __future__ import annotations

import argparse
import json
import statistics
import time
from pathlib import Path

from mcm26b.config import ProblemConfig
from mcm26b.harness import (LocalSimTransport, mixed_scene, random_omni_scene,
                            referee_spec)
from mcm26b.planner import DirectionalExtension
from mcm26b.planner.bayes_loop import BayesLoop
from mcm26b.planner.rl_loop import RLBayesLoop, RLDirectionalLoop
from mcm26b.protocol.client import SimulatorClient
from referee import LocalSimulator


def run_one(loop_cls, sources, *, config, robot_id: str, **kwargs) -> dict:
    simulator = LocalSimulator(sources, robot_id=robot_id,
                               spec=referee_spec(config, max_real_duration_s=3600.0),
                               stress_mode=True)
    client = SimulatorClient(robot_id, config=config,
                             transport=LocalSimTransport(simulator))
    started = time.monotonic()
    summary = loop_cls(client, config, **kwargs).run() \
        if kwargs else loop_cls(client, config).run()
    return {
        "completed": bool(summary["completed"]),
        "virtual_time_s": float(summary["virtual_time_s"]),
        "measures": int(summary["measures"]),
        "total_requests": int(summary["total_requests"]),
        "move_distance_m": float(summary["move_distance_m"]),
        "wall_s": round(time.monotonic() - started, 3),
    }


def main(argv=None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--seeds", type=int, default=5)
    parser.add_argument("--sources", type=int, default=10)
    parser.add_argument("--directional", type=int, default=3)
    parser.add_argument("--output", default="output/rl/benchmark.json")
    args = parser.parse_args(argv)
    config = ProblemConfig()
    robot_id = "2026001"
    report = {"config": {"seeds": args.seeds, "sources": args.sources,
                         "directional": args.directional},
              "q3": [], "q4": []}

    for seed in range(20260910, 20260910 + args.seeds):
        omni = random_omni_scene(seed=seed, count=args.sources, config=config)
        mixed = mixed_scene(seed=seed, count=args.sources,
                            directional=args.directional, config=config)
        q3_base = run_one(BayesLoop, omni, config=config, robot_id=robot_id)
        q3_rl = run_one(RLBayesLoop, omni, config=config, robot_id=robot_id)
        q4_base = run_one(DirectionalExtension, mixed, config=config,
                          robot_id=robot_id)
        q4_rl = run_one(RLDirectionalLoop, mixed, config=config,
                        robot_id=robot_id)
        report["q3"].append({"seed": seed, "bayes": q3_base, "rl": q3_rl})
        report["q4"].append({"seed": seed, "directional": q4_base, "rl": q4_rl})
        print(json.dumps({"event": "seed_done", "seed": seed,
                          "q3_base_s": q3_base["virtual_time_s"],
                          "q3_rl_s": q3_rl["virtual_time_s"],
                          "q4_base_s": q4_base["virtual_time_s"],
                          "q4_rl_s": q4_rl["virtual_time_s"]},
                         ensure_ascii=False))

    for key, baseline in (("q3", "bayes"), ("q4", "directional")):
        rows = report[key]
        deltas = [row["rl"]["virtual_time_s"] - row[baseline]["virtual_time_s"]
                  for row in rows]
        cleared = all(row["rl"]["completed"] and row[baseline]["completed"]
                      for row in rows)
        report[key + "_summary"] = {
            "mean_virtual_s": {
                baseline: statistics.mean(row[baseline]["virtual_time_s"] for row in rows),
                "rl": statistics.mean(row["rl"]["virtual_time_s"] for row in rows),
            },
            "mean_delta_s": statistics.mean(deltas),
            "relative_delta_pct": 100.0 * statistics.mean(deltas)
            / statistics.mean(row[baseline]["virtual_time_s"] for row in rows),
            "mean_measures": {
                baseline: statistics.mean(row[baseline]["measures"] for row in rows),
                "rl": statistics.mean(row["rl"]["measures"] for row in rows),
            },
            "all_cleared": cleared,
        }
    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(json.dumps(report, ensure_ascii=False, indent=2),
                      encoding="utf-8")
    print(json.dumps({"event": "benchmark_done", "path": str(output),
                      "q3": report["q3_summary"],
                      "q4": report["q4_summary"]},
                     ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())