"""问题3/4 强化学习选点策略训练（PPO + CUDA）。

训练环境 mcm26b.rl.env.LocalizationEnv 与运行闭环共享 features 观测/候选
布局。模型输出到 output/rl/q3_ppo.zip（问题3）或 q4_ppo.zip（问题4），
由 `mcm26b run --strategy rl` / `--q4-strategy rl` 加载。

示例：
    uv sync --group rl --group dev
    uv run python scripts/train_rl.py --problem 3 --timesteps 30000
    uv run python scripts/train_rl.py --problem 4 --timesteps 30000
    uv run mcm26b run --robot-id 2026001 --self-test --strategy rl
    uv run mcm26b run --robot-id 2026001 --self-test --problem 4 --q4-strategy rl
"""

from __future__ import annotations

import argparse
import json
import time
from dataclasses import asdict
import hashlib
from pathlib import Path

from stable_baselines3 import PPO
from sb3_contrib import MaskablePPO
from stable_baselines3.common.vec_env import DummyVecEnv

from mcm26b.config import ProblemConfig
from mcm26b.rl.env import LocalizationEnv


def _make_env(problem: int, seed: int, config: ProblemConfig,
              directional_probability: float):
    def _factory():
        env = LocalizationEnv(config, problem=problem,
                              directional_probability=directional_probability,
                              real_budget_s=3600.0)
        env.reset(seed=seed)
        return env
    return _factory


def evaluate(model, problem: int, *, episodes: int = 24,
             directional_probability: float = 0.5,
             config: ProblemConfig = ProblemConfig()) -> dict:
    """独立场景评估：确定性策略的成功率、步数与平均回报。"""
    env = LocalizationEnv(config, problem=problem,
                          directional_probability=directional_probability)
    total_reward = 0.0
    successes = 0
    steps = []
    costs = []
    for episode in range(episodes):
        obs, _info = env.reset(seed=1_000_000 + episode)
        episode_reward = 0.0
        for step in range(env.max_episode_steps + 2):
            kwargs = {'action_masks': env.action_masks()} if isinstance(model, MaskablePPO) else {}
            action, _state = model.predict(obs, deterministic=True, **kwargs)
            obs, reward, terminated, truncated, info = env.step(int(action))
            episode_reward += float(reward)
            if terminated or truncated:
                break
        total_reward += episode_reward
        cleared = bool(info.get("cleared"))
        successes += int(cleared)
        steps.append(step + 1)
        costs.append(info['virtual_cost_s'])
    env.close()
    return {
        "episodes": episodes,
        "success_rate": successes / episodes,
        "mean_steps": sum(steps) / episodes,
        "mean_episode_reward": total_reward / episodes,
        "successes": successes,
        "mean_virtual_cost_s": sum(costs) / episodes,
        "max_virtual_cost_s": max(costs),
    }


def main(argv=None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--problem", type=int, choices=(3, 4), default=3,
                        help="训练问题：3=全向源选点，4=含定向遮挡梯子")
    parser.add_argument("--timesteps", type=int, default=30000,
                        help="PPO 总环境步数")
    parser.add_argument("--seed", type=int, default=20260911)
    parser.add_argument("--envs", type=int, default=2, help="并行环境份数（CPU 侧）")
    parser.add_argument("--directional-probability", type=float, default=None,
                        help="问题4训练场景中定向源比例（默认 1.0，先学遮挡分支）")
    parser.add_argument("--device", default="auto",
                        help="auto=优先 CUDA；或 cuda/cpu")
    parser.add_argument("--model", default=None,
                        help="模型输出路径（默认 output/rl/q<problem>_ppo.zip）")
    parser.add_argument("--init", default=None,
                        help="初始模型路径（如 output/rl/q3_ppo_bc_init.zip；"
                             "加载后继续 PPO 微调）")
    parser.add_argument("--eval-episodes", type=int, default=24)
    parser.add_argument("--verbose", type=int, default=1)
    parser.add_argument('--algorithm', choices=('ppo', 'maskable_ppo'), default='maskable_ppo')
    parser.add_argument('--torch-threads', type=int, default=1)
    args = parser.parse_args(argv)
    directional_probability = (args.directional_probability
                               if args.directional_probability is not None
                               else (1.0 if args.problem == 4 else 0.0))
    if not 0.0 <= directional_probability <= 1.0:
        parser.error("--directional-probability 须在 [0,1]")

    import torch
    torch.set_num_threads(args.torch_threads)
    if args.device == "auto":
        device = "cuda" if torch.cuda.is_available() else "cpu"
    else:
        device = args.device
    if device == "cuda" and not torch.cuda.is_available():
        parser.error("--device cuda 但当前环境检测不到 CUDA GPU")
    print(json.dumps({
        "event": "train_start", "problem": args.problem,
        "timesteps": args.timesteps, "device": device,
        "directional_probability": directional_probability,
        "torch": torch.__version__,
        "gpu": torch.cuda.get_device_name(0) if device == "cuda" else None,
        "cuda_available": torch.cuda.is_available(),
    }, ensure_ascii=False))

    config = ProblemConfig()
    model_cls = MaskablePPO if args.algorithm == 'maskable_ppo' else PPO
    vec_env = DummyVecEnv(
        [_make_env(args.problem, args.seed + rank, config,
                   directional_probability)
         for rank in range(max(1, args.envs))])
    if args.init:
        metadata = json.loads(Path(args.init).with_suffix('.json').read_text(encoding='utf-8'))
        if (metadata.get('algorithm', 'ppo') != args.algorithm
                or metadata.get('feature_version') != 2):
            parser.error('初始化模型的算法/特征版本与当前训练不符，请重训或提供v2初始化模型')
        model = model_cls.load(args.init, env=vec_env, device=device)
    else:
        model = model_cls(
            "MlpPolicy", vec_env,
            policy_kwargs={"net_arch": [128, 128]},
            n_steps=512,
            batch_size=128,
            n_epochs=6,
            gamma=1.0,
            gae_lambda=0.95,
            learning_rate=3e-4,
            clip_range=0.2,
            ent_coef=0.005,
            vf_coef=0.5,
            max_grad_norm=0.5,
            seed=args.seed,
            device=device,
            verbose=args.verbose,
        )

    started = time.monotonic()
    model.learn(total_timesteps=args.timesteps,
                progress_bar=args.verbose > 0)
    elapsed_s = time.monotonic() - started

    model_path = Path(args.model) if args.model \
        else Path("output/rl") / f"q{args.problem}_ppo.zip"
    model_path.parent.mkdir(parents=True, exist_ok=True)
    model.save(str(model_path))

    evaluation = evaluate(model, args.problem, episodes=args.eval_episodes,
                          directional_probability=directional_probability,
                          config=config)
    report = {
        "problem": args.problem,
        "timesteps": args.timesteps,
        "seed": args.seed,
        "device": device,
        "directional_probability": directional_probability,
        "wall_seconds": round(elapsed_s, 1),
        "model": str(model_path),
        "init": args.init,
        "algorithm": args.algorithm,
        "feature_version": 2,
        "reward": 'negative_virtual_cost_div100_success20_failure100',
        "gamma": 1.0,
        "config": asdict(config),
        "model_sha256": hashlib.sha256(model_path.read_bytes()).hexdigest(),
        **evaluation,
    }
    report_path = model_path.with_suffix(".json")
    report_path.write_text(json.dumps(report, ensure_ascii=False, indent=2),
                           encoding="utf-8")
    print(json.dumps(report, ensure_ascii=False, indent=2))
    vec_env.close()
    return 0


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