"""用自建模拟器的既有专家逻辑做行为克隆，初始化 PPO 策略网络。

思路：自建 referee 的观测模型与真值逻辑既快又可复现。先在它上面用
手写专家（Q3: worst-posterior 选点器；Q4: 遮挡梯子选点器）生成
(观测, 专家动作) 数据集，训练一个与 PPO policy_net 同构的 MLP 分类器，
再把权重拷入 PPO 模型作为初始化；随后 scripts/train_rl.py --init 继续
PPO 微调。这样 RL 从强先验出发，而不是从随机策略开始。

产物：
    output/rl/bc/q3_expert.npz  output/rl/bc/q3_bc.pt  output/rl/q3_ppo_bc_init.zip
    output/rl/bc/q4_expert.npz  output/rl/bc/q4_bc.pt  output/rl/q4_ppo_bc_init.zip
"""
from __future__ import annotations

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

import numpy as np
import torch
from stable_baselines3 import PPO

from mcm26b.config import ProblemConfig
from mcm26b.planner import safe_disk_radius, select_measurement_point
from mcm26b.rl._constants import NUM_ACTIONS, OBS_DIM
from mcm26b.rl.env import LocalizationEnv
from mcm26b.rl.features import candidate_actions

HIDDEN = (128, 128)


def expert_action(env: LocalizationEnv) -> int:
    """当前状态的专家动作编号（与训练环境候选同源，就近映射）。"""
    config = env.config
    channel = env._channel
    region = env._belief.candidate_set(channel)
    position = env._client.state.position
    blocked = env._blocked
    candidates = candidate_actions(region, position, config, blocked)

    if env.problem == 3:
        expert_point = select_measurement_point(
            region, position, config).position
    else:
        center, radius = region.mec()
        safe = safe_disk_radius(region, config)
        if not blocked:
            expert_point = select_measurement_point(
                region, position, config).position
        elif len(blocked) == 1 and safe >= 0:
            first = blocked[0]
            dx, dy = center[0] - first[0], center[1] - first[1]
            norm = math.hypot(dx, dy)
            if norm <= 1e-9:
                expert_point = (center[0] + safe, center[1])
            else:
                scale = min(1.0, safe / norm)
                expert_point = (center[0] + dx * scale,
                                center[1] + dy * scale)
        else:
            first = blocked[0]
            base = math.atan2(first[1] - center[1], first[0] - center[0])
            angle = base + math.radians(45.0 * (len(blocked) - 1))
            expert_point = (center[0] + radius * math.cos(angle),
                            center[1] + radius * math.sin(angle))
    return min(range(len(candidates)),
               key=lambda k: math.dist(candidates[k], expert_point))


class BCMlp(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.net = torch.nn.Sequential(
            torch.nn.Linear(OBS_DIM, HIDDEN[0]), torch.nn.Tanh(),
            torch.nn.Linear(HIDDEN[0], HIDDEN[1]), torch.nn.Tanh(),
            torch.nn.Linear(HIDDEN[1], NUM_ACTIONS))

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


def collect(problem: int, episodes: int, seed: int, config: ProblemConfig):
    env = LocalizationEnv(
        config, problem=problem,
        directional_probability=1.0 if problem == 4 else 0.0,
        max_episode_steps=16)
    observations, actions, groups = [], [], []
    for episode in range(episodes):
        obs, _info = env.reset(seed=seed + episode)
        for _step in range(env.max_episode_steps + 2):
            action = expert_action(env)
            observations.append(obs)
            actions.append(action)
            groups.append(episode)
            obs, _reward, terminated, truncated, _info = env.step(action)
            if terminated or truncated:
                break
    return np.asarray(observations, dtype=np.float32), \
        np.asarray(actions, dtype=np.int64), np.asarray(groups, dtype=np.int64)


def train_bc(x, y, *, epochs: int, batch_size: int, lr: float,
             device: str, groups, valid_frac: float = 0.1):
    model = BCMlp().to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
    episodes = np.unique(groups)
    if len(episodes) < 2:
        raise ValueError('行为克隆至少需要两个独立episode才能划分验证集')
    valid_groups = np.random.default_rng(0).permutation(episodes)
    valid_groups = valid_groups[:max(1, int(len(episodes)*valid_frac))]
    valid = torch.as_tensor(np.isin(groups, valid_groups), device=device)
    x_t, y_t = torch.as_tensor(x, device=device), torch.as_tensor(y, device=device)
    x_v, y_v = x_t[valid], y_t[valid]
    x_tr, y_tr = 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())
        train_acc = correct / len(x_tr)
        with torch.no_grad():
            valid_acc = float((model(x_v).argmax(1) == y_v).float().mean().item())
        if epoch == 0 or (epoch + 1) % 5 == 0:
            print(json.dumps({"event": "bc_epoch", "epoch": epoch + 1,
                              "loss": round(total / len(x_tr), 4),
                              "train_acc": round(train_acc, 4),
                              "valid_acc": round(valid_acc, 4)},
                             ensure_ascii=False))
    return model, float(valid_acc)


def copy_into_ppo(problem: int, model: BCMlp, device: str) -> Path | None:
    """把 BC 网络权重按 SB3 的 policy_net/action_net 拆分拷入 PPO。"""
    # 当前奖励/特征已修订，重新初始化critic，不能沿用旧任务的价值预测。
    env = LocalizationEnv(problem=problem)
    ppo = PPO('MlpPolicy', env, policy_kwargs={'net_arch': list(HIDDEN)},
              n_steps=512, batch_size=128, gamma=1.0, device=device)
    state = {key.removeprefix("net."): value
             for key, value in model.state_dict().items()}
    # SB3 MlpPolicy：policy_net 只含前两个 Linear（Tanh 激活），
    # 输出层在 action_net。BC 的 Sequential 键为 0/2/4。
    policy_state = {key: value for key, value in state.items()
                    if key in ("0.weight", "0.bias", "2.weight", "2.bias")}
    action_state = {
        "weight": state["4.weight"], "bias": state["4.bias"]}
    ppo.policy.mlp_extractor.policy_net.load_state_dict(policy_state)
    ppo.policy.action_net.load_state_dict(action_state)
    # 数值检查整个前向映射，而不只是线性层形状。
    probe = torch.randn((8, OBS_DIM), device=device)
    with torch.no_grad():
        copied = ppo.policy.action_net(ppo.policy.mlp_extractor.policy_net(probe))
        if not torch.allclose(model(probe), copied, atol=1e-6):
            raise RuntimeError('行为克隆与PPO激活函数/前向映射不一致')
    init_path = Path(f"output/rl/q{problem}_ppo_bc_init.zip")
    ppo.save(str(init_path))
    env.close()
    init_path.with_suffix('.json').write_text(json.dumps(
        dict(algorithm='ppo', feature_version=2, problem=problem)), encoding='utf-8')
    return init_path


def main(argv=None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--problem", type=int, choices=(3, 4), required=True)
    parser.add_argument("--episodes", type=int, default=3000)
    parser.add_argument("--epochs", type=int, default=20)
    parser.add_argument("--batch-size", type=int, default=256)
    parser.add_argument("--lr", type=float, default=3e-4)
    parser.add_argument("--seed", type=int, default=20260911)
    parser.add_argument("--device", default="auto")
    parser.add_argument("--copy-only", action="store_true",
                        help="跳过采集/训练，只把已有 q<problem>_bc.pt 拷入 PPO 初始化")
    args = parser.parse_args(argv)

    if args.device == "auto":
        device = "cuda" if torch.cuda.is_available() else "cpu"
    else:
        device = args.device
    print(json.dumps({"event": "bc_start", "problem": args.problem,
                      "episodes": args.episodes, "device": device,
                      "cuda": torch.cuda.is_available()},
                     ensure_ascii=False))

    if args.copy_only:
        bc_dir = Path("output/rl/bc")
        checkpoint = torch.load(bc_dir / f"q{args.problem}_bc.pt",
                                map_location="cpu")
        if checkpoint.get('activation') != 'Tanh' or checkpoint.get('feature_version') != 2:
            parser.error('旧ReLU/旧特征BC检查点不能直接拷贝，请重新采集并训练')
        model = BCMlp()
        model.load_state_dict(checkpoint["state_dict"])
        init_path = copy_into_ppo(args.problem, model, device)
        dataset = np.load(bc_dir / f"q{args.problem}_expert.npz",
                          allow_pickle=False)
        report = {
            "problem": args.problem,
            "samples": int(len(dataset["obs"])),
            "valid_accuracy": checkpoint.get("valid_acc"),
            "dataset": str(bc_dir / f"q{args.problem}_expert.npz"),
            "bc_model": str(bc_dir / f"q{args.problem}_bc.pt"),
            "ppo_init": str(init_path) if init_path is not None else None,
        }
        (bc_dir / f"q{args.problem}_bc.json").write_text(
            json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
        print(json.dumps({"event": "bc_copy", "problem": args.problem,
                          "ppo_init": str(init_path), **report},
                         ensure_ascii=False, indent=2))
        return 0 if init_path is not None else 2

    config = ProblemConfig()
    started = time.monotonic()
    torch.manual_seed(args.seed)
    x, y, groups = collect(args.problem, args.episodes, args.seed, config)
    collect_s = time.monotonic() - started
    bc_dir = Path("output/rl/bc")
    bc_dir.mkdir(parents=True, exist_ok=True)
    np.savez_compressed(bc_dir / f"q{args.problem}_expert.npz", obs=x, action=y, episode=groups)

    model, valid_acc = train_bc(x, y, epochs=args.epochs,
                                batch_size=args.batch_size, lr=args.lr,
                                device=device, groups=groups)
    torch.save({"state_dict": model.state_dict(), "obs_dim": OBS_DIM,
                "actions": NUM_ACTIONS, "valid_acc": valid_acc,
                "activation": "Tanh", "feature_version": 2},
               bc_dir / f"q{args.problem}_bc.pt")

    # 拷入现有 PPO 的 policy_net/action_net 作为初始化（critic 保持原权重）。
    init_path = copy_into_ppo(args.problem, model, device)
    report = {
        "problem": args.problem, "episodes": args.episodes,
        "samples": int(len(x)), "valid_accuracy": valid_acc,
        "collect_seconds": round(collect_s, 1),
        "train_seconds": round(time.monotonic() - started - collect_s, 1),
        "dataset": str(bc_dir / f"q{args.problem}_expert.npz"),
        "bc_model": str(bc_dir / f"q{args.problem}_bc.pt"),
        "ppo_init": str(init_path) if init_path is not None else None,
    }
    (bc_dir / f"q{args.problem}_bc.json").write_text(
        json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(report, ensure_ascii=False, indent=2))
    return 0


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