import glob
import os

import numpy as np
import torch

BASE = os.environ.get("GATES_BASE", "models/abl")
VARIANTS = os.environ.get("GATES_VARIANTS", "moe,llm,both").split(",")
TASKS = ["size", "delivery", "pdi", "ee", "toxic", "biodist"]


def latest(v: str):
    dirs = sorted(glob.glob(f"{BASE}/{v}/*/"), key=os.path.getmtime)
    return dirs[-1] if dirs else None


for v in VARIANTS:
    run = latest(v)
    if run is None:
        print(f"\n=== {v}: 无 run 目录 ===")
        continue

    files = sorted(glob.glob(os.path.join(run, "**", "model.pt"), recursive=True))
    print(f"\n=== {v} ({len(files)} folds) ===")
    if not files:
        print("  未找到 model.pt（本次 run 未保存权重 -> 无法读取门控/log_vars）")
        continue

    gm, gl = [], []
    for f in files:
        sd = torch.load(f, map_location="cpu", weights_only=False)["model_state_dict"]
        for key, acc in (("fusion.g_moe", gm), ("fusion.g_llm", gl)):
            if key not in sd:
                continue
            t = sd[key]
            acc.append(abs(t.item()) if t.dim() == 0
                       else float(t.norm()) / (t.numel() ** 0.5))
    for name, acc in (("g_moe", gm), ("g_llm", gl)):
        if acc:
            print(f"  {name}: mean={np.mean(acc):.4f}  range=[{min(acc):.4f},{max(acc):.4f}]")

    sd = torch.load(files[0], map_location="cpu", weights_only=False)["model_state_dict"]
    for t in TASKS:
        lv = sd.get(f"head.log_vars.{t}")
        if lv is not None:
            print(f"  {t:9s} log_var={lv.item():+.3f}  eff_w={np.exp(-lv.item()):.3f}")