"""BayesLoop 的 Optuna 离线调参（TPE），无 Optuna 时退化为随机搜索。

目标：场景子集的平均虚拟时间（秒）+ 10^6 × 失败惩罚（未全清/未完成/
异常/缺席误判任一即失败）。场景与误差种子固定但多样，防止对单一
随机种子过拟合；最优参数写入 JSON，人工复跑 benchmark_q3_bayes.py
确认后再回填 config.BayesConfig 默认值。

Run: uv run python scripts/tune_q3_bayes.py --n-trials 100
"""
from __future__ import annotations

import argparse
from dataclasses import replace
import json
import math
from pathlib import Path
from random import Random
from statistics import mean
from time import perf_counter

from mcm26b.config import ProblemConfig
from mcm26b.planner.bayes_loop import BayesLoop
from mcm26b.protocol.client import SimulatorClient
from mcm26b.harness import LocalSimTransport, LocalSimulator
from referee import Source
from mcm26b.harness import random_omni_scene

ROOT = Path(__file__).resolve().parents[1]
FAIL_PENALTY = 1.0e6

BASE = replace(ProblemConfig(),
               search=replace(ProblemConfig().search, persistent_search=False))

# 调参场景：多种子 × 两档源数 × 边界场景（误差种子与场景种子解耦）
SCENES = [(f"tune-{n}-{s}", random_omni_scene(s, n, config=BASE), 20260910 + s)
          for n in (10, 16) for s in (0, 1, 2)]
SCENES.append(("tune-boundary-10",
               [Source(c + 1, 1800 * math.cos(math.pi / 6 + 2 * math.pi * c / 10),
                       1800 * math.sin(math.pi / 6 + 2 * math.pi * c / 10), 1000)
                for c in range(10)], 20260977))


def evaluate(bayes_params: dict) -> tuple[float, dict]:
    config = replace(BASE, bayes=replace(BASE.bayes, **bayes_params))
    times = []
    failures = 0
    for name, sources, error_seed in SCENES:
        sim = LocalSimulator(sources, robot_id="TUNE", config=config,
                             error_seed=error_seed)
        client = SimulatorClient("TUNE", config=config,
                                 transport=LocalSimTransport(sim))
        summary = BayesLoop(client, config).run()
        truth = {s.channel for s in sources}
        ok = (sorted(summary["cleared_channels"]) == sorted(truth)
              and summary["completed"]
              and not summary["anomalies"]
              and set(summary["absent_channels"]) <= set(range(1, 21)) - truth)
        if not ok:
            failures += 1
            continue
        times.append(summary["virtual_time_s"])
    objective = mean(times) if times else FAIL_PENALTY
    return objective + FAIL_PENALTY * failures, {
        "failures": failures, "mean_virtual_s": objective, "times": times}


def suggest(trial, rng: Random):
    """搜索空间：Optuna trial 或（退化时）rng 随机采样，二者同构。"""
    if trial is not None:
        return {
            "lambda_move_s": trial.suggest_float("lambda_move_s", 0.01, 0.6, log=True),
            "mc_representatives": trial.suggest_int("mc_representatives", 3, 12),
            "score_subsample": trial.suggest_categorical("score_subsample",
                                                         [300, 500, 800, 1200]),
            "particle_spacing_m": trial.suggest_float("particle_spacing_m", 12.0, 40.0),
            "station_ring_radius_m": trial.suggest_float("station_ring_radius_m",
                                                         1100.0, 1450.0),
            "two_opt_passes": trial.suggest_int("two_opt_passes", 1, 6),
            "score_mode": trial.suggest_categorical("score_mode",
                                                    ["radius", "entropy"]),
        }
    return {
        "lambda_move_s": math.exp(rng.uniform(math.log(0.01), math.log(0.6))),
        "mc_representatives": rng.randint(3, 12),
        "score_subsample": rng.choice([300, 500, 800, 1200]),
        "particle_spacing_m": rng.uniform(12.0, 40.0),
        "station_ring_radius_m": rng.uniform(1100.0, 1450.0),
        "two_opt_passes": rng.randint(1, 6),
        "score_mode": rng.choice(["radius", "entropy"]),
    }


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--n-trials", type=int, default=60)
    parser.add_argument("--out", type=Path,
                        default=ROOT / "output/q3/bayes-tuning")
    args = parser.parse_args()
    args.out.mkdir(parents=True, exist_ok=True)
    start = perf_counter()

    try:
        import optuna
        optuna.logging.set_verbosity(optuna.logging.WARNING)

        def objective(trial):
            value, _detail = evaluate(suggest(trial, None))
            return value

        study = optuna.create_study(direction="minimize",
                                    sampler=optuna.samplers.TPESampler(seed=7))
        study.optimize(objective, n_trials=args.n_trials)
        best_params, best_value = study.best_params, study.best_value
        sampler = "tpe"
    except ImportError:
        rng = Random(7)
        best_params, best_value, sampler = None, float("inf"), "random_fallback"
        for _ in range(args.n_trials):
            params = suggest(None, rng)
            value, _detail = evaluate(params)
            if value < best_value:
                best_params, best_value = params, value

    result = {"sampler": sampler, "n_trials": args.n_trials,
              "best_value_s": best_value, "best_params": best_params,
              "scenes": [name for name, _s, _e in SCENES],
              "wall_s": perf_counter() - start,
              "note": "回填前先以 benchmark_q3_bayes.py 全量复跑确认。"}
    (args.out / "best_params.json").write_text(
        json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(result, ensure_ascii=False, indent=2))
    return 0


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