"""低层 PPO 训练：8 方向到达任务（BC 直线专家热启动 + PPO 微调）。

产物：
    output/hrl/hrl_low.zip / hrl_low.json          PPO 模型（部署用）
    output/hrl/bc_low_expert.npz / bc_low.pt       行为克隆数据与网络
    output/hrl/hrl_low_curve_env*.monitor.csv     每个环境独立学习曲线

评估口径：到达成功率 + 路径效率（直线距离/实际路径长，1 为完美直线）。
这是独立导航实验：目前低层 PPO 仅供轨迹诊断，不影响正式 HTTP 任务。

示例：
    uv run python scripts/train_hrl_low.py --timesteps 60000
    uv run python scripts/train_hrl_low.py --skip-bc   # 纯 PPO 对照
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import time
from pathlib import Path

from mcm26b.config import ProblemConfig
from mcm26b.hrl.constants import (FEATURE_VERSION, LOW_MAX_STEPS,
                                  LOW_NUM_ACTIONS, LOW_OBS_DIM)
from mcm26b.hrl.low_level_env import LowLevelEnv

HIDDEN = (64, 64)


def _bc_module(obs_dim: int, n_actions: int, hidden: tuple[int, ...]):
    import torch

    class BCMlp(torch.nn.Module):
        def __init__(self):
            super().__init__()
            layers = []
            prev = obs_dim
            for width in hidden:
                layers += [torch.nn.Linear(prev, width), torch.nn.Tanh()]
                prev = width
            layers.append(torch.nn.Linear(prev, n_actions))
            self.net = torch.nn.Sequential(*layers)

        def forward(self, x):
            return self.net(x)

    return BCMlp


def collect(episodes: int, seed: int, config: ProblemConfig):
    import numpy as np
    env = LowLevelEnv(config=config, seed=seed)
    observations, actions, episode_ids = [], [], []
    for episode in range(episodes):
        obs, _info = env.reset(seed=seed + episode)
        for _step in range(LOW_MAX_STEPS + 2):
            observations.append(obs)
            actions.append(env.expert_action())
            episode_ids.append(episode)
            obs, _r, terminated, truncated, _info = env.step(env.expert_action())
            if terminated or truncated:
                break
    env.close()
    return (np.asarray(observations, dtype=np.float32),
            np.asarray(actions, dtype=np.int64),
            np.asarray(episode_ids, dtype=np.int64))


def bc_validation_mask(episode_ids, *, seed: int):
    """按整局留出，避免同一轨迹相邻状态跨训练/验证集。"""
    import numpy as np
    unique = np.unique(episode_ids)
    if len(unique) < 2:
        raise ValueError("BC 至少需要 2 局才能按整局划分训练/验证集")
    rng = np.random.default_rng(seed)
    validation_ids = rng.permutation(unique)[:max(1, len(unique) // 10)]
    return np.isin(episode_ids, validation_ids)


def train_bc(x, y, episode_ids, *, epochs: int, batch_size: int, lr: float,
             device: str, seed: int):
    import torch
    BCMlp = _bc_module(LOW_OBS_DIM, LOW_NUM_ACTIONS, HIDDEN)
    model = BCMlp().to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
    valid = torch.as_tensor(
        bc_validation_mask(episode_ids, seed=seed), device=device)
    x_t = torch.as_tensor(x, device=device)
    y_t = torch.as_tensor(y, device=device)
    x_v, y_v, x_tr, y_tr = x_t[valid], y_t[valid], x_t[~valid], y_t[~valid]
    loss_fn = torch.nn.CrossEntropyLoss()
    for epoch in range(epochs):
        order = torch.randperm(len(x_tr), device=device)
        total = correct = 0.0
        for start in range(0, len(order), batch_size):
            idx = order[start:start + batch_size]
            logits = model(x_tr[idx])
            loss = loss_fn(logits, y_tr[idx])
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            total += float(loss.item()) * len(idx)
            correct += int((logits.argmax(1) == y_tr[idx]).sum().item())
        if epoch == 0 or (epoch + 1) % 5 == 0:
            with torch.no_grad():
                valid_acc = float((model(x_v).argmax(1) == y_v).float().mean().item())
            print(json.dumps({"event": "bc_epoch", "epoch": epoch + 1,
                              "loss": round(total / len(x_tr), 4),
                              "train_acc": round(correct / len(x_tr), 4),
                              "valid_acc": round(valid_acc, 4)},
                             ensure_ascii=False))
    with torch.no_grad():
        valid_acc = float((model(x_v).argmax(1) == y_v).float().mean().item())
    return model, valid_acc


def copy_into_ppo(bc_model, device: str, env, *, seed: int,
                  n_steps: int = 512) -> object:
    """把 BC 权重拷入 MaskablePPO（policy_net 两层 + action_net），前向校验。"""
    import torch
    from sb3_contrib import MaskablePPO
    ppo = MaskablePPO("MlpPolicy", env, policy_kwargs={"net_arch": list(HIDDEN)},
                      n_steps=n_steps, 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=seed, device=device)
    state = {key.removeprefix("net."): value
             for key, value in bc_model.state_dict().items()}
    policy_state = {key: value for key, value in state.items()
                    if key in ("0.weight", "0.bias", "2.weight", "2.bias")}
    ppo.policy.mlp_extractor.policy_net.load_state_dict(policy_state)
    ppo.policy.action_net.load_state_dict(
        {"weight": state["4.weight"], "bias": state["4.bias"]})
    probe = torch.randn((8, LOW_OBS_DIM), device=device)
    with torch.no_grad():
        copied = ppo.policy.action_net(ppo.policy.mlp_extractor.policy_net(probe))
        if not torch.allclose(bc_model(probe), copied, atol=1e-6):
            raise RuntimeError("行为克隆与 PPO 前向映射不一致")
    return ppo


def evaluate(model, *, episodes: int, seed: int,
             config: ProblemConfig) -> dict:
    """确定性策略评估：成功率与路径效率（对照直线）。"""
    env = LowLevelEnv(config=config, seed=seed)
    successes = 0
    efficiencies = []
    for episode in range(episodes):
        obs, info = env.reset(seed=seed + episode)
        start, goal = tuple(info["position"]), tuple(info["goal"])
        straight = math.dist(start, goal)
        steps = 0
        for _ in range(LOW_MAX_STEPS + 2):
            action, _state = model.predict(obs, deterministic=True)
            obs, _r, terminated, truncated, info = env.step(int(action))
            steps += 1
            if terminated or truncated:
                break
        successes += int(terminated)
        efficiencies.append(min(1.0, straight /
                                max(info["path_length_m"], 1e-9)))
    env.close()
    return {"episodes": episodes,
            "success_rate": successes / episodes,
            "mean_path_efficiency": sum(efficiencies) / len(efficiencies)}


def main(argv=None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--timesteps", type=int, default=60_000)
    parser.add_argument("--seed", type=int, default=20260911)
    parser.add_argument("--bc-episodes", type=int, default=600)
    parser.add_argument("--bc-epochs", type=int, default=12)
    parser.add_argument("--envs", type=int, default=4)
    parser.add_argument("--rollout-steps", type=int, default=512,
                        help="每环境每次 PPO rollout 长度；短 smoke 可设 8")
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--model", default="output/hrl/hrl_low.zip")
    parser.add_argument("--eval-episodes", type=int, default=40)
    parser.add_argument("--skip-bc", action="store_true",
                        help="跳过行为克隆（纯 PPO 对照实验用）")
    parser.add_argument("--verbose", type=int, default=1)
    args = parser.parse_args(argv)
    if min(args.envs, args.timesteps, args.eval_episodes,
           args.rollout_steps) <= 0:
        parser.error("envs/timesteps/eval-episodes/rollout-steps 必须为正数")
    if args.envs * args.rollout_steps < 2:
        parser.error("PPO rollout 至少需要 2 个样本")
    if not args.skip_bc and (args.bc_episodes < 2 or args.bc_epochs < 1):
        parser.error("BC 至少需要 2 局和 1 个 epoch")

    import torch
    torch.set_num_threads(1)
    from sb3_contrib import MaskablePPO
    from stable_baselines3.common.monitor import Monitor
    from stable_baselines3.common.vec_env import DummyVecEnv
    from stable_baselines3.common.utils import set_random_seed
    set_random_seed(args.seed)

    model_path = Path(args.model)
    model_path.parent.mkdir(parents=True, exist_ok=True)
    curve_paths = [model_path.with_name(
        f"{model_path.stem}_curve_env{rank}.monitor.csv")
        for rank in range(args.envs)]

    def _factory(rank: int):
        def _make():
            env = Monitor(LowLevelEnv(seed=args.seed + rank),
                          filename=str(curve_paths[rank]))
            return env
        return _make

    vec_env = DummyVecEnv([_factory(rank) for rank in range(args.envs)])
    started = time.monotonic()
    bc_valid_acc = None
    if args.skip_bc:
        model = MaskablePPO("MlpPolicy", vec_env,
                            policy_kwargs={"net_arch": list(HIDDEN)},
                            n_steps=args.rollout_steps, 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=args.device, verbose=args.verbose)
    else:
        x, y, episode_ids = collect(args.bc_episodes, args.seed, ProblemConfig())
        bc_model, bc_valid_acc = train_bc(
            x, y, episode_ids, epochs=args.bc_epochs, batch_size=256,
            lr=3e-4, device=args.device, seed=args.seed)
        model = copy_into_ppo(bc_model, args.device, vec_env, seed=args.seed,
                              n_steps=args.rollout_steps)
        bc_dir = model_path.parent / "bc_low.pt"
        torch.save(bc_model.state_dict(), str(bc_dir))
        print(json.dumps({"event": "bc_done", "samples": int(len(x))},
                         ensure_ascii=False))
    model.learn(total_timesteps=args.timesteps,
                progress_bar=args.verbose > 0)
    elapsed_s = time.monotonic() - started

    model.save(str(model_path))
    eval_seed = args.seed + 1_000_000
    evaluation = evaluate(model, episodes=args.eval_episodes, seed=eval_seed,
                          config=ProblemConfig())
    report = {
        "level": "low", "timesteps": args.timesteps, "seed": args.seed,
        "device": args.device, "wall_seconds": round(elapsed_s, 1),
        "model": str(model_path), "algorithm": "maskable_ppo",
        "feature_version": FEATURE_VERSION, "bc_skipped": args.skip_bc,
        "bc_valid_acc": bc_valid_acc, "evaluation_seed": eval_seed,
        "monitor_files": [str(path) for path in curve_paths],
        "training_role": "standalone_navigation_only",
        "obs_dim": LOW_OBS_DIM, "act_dim": LOW_NUM_ACTIONS,
        "model_sha256": hashlib.sha256(model_path.read_bytes()).hexdigest(),
        **evaluation,
    }
    model_path.with_suffix(".json").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())