import glob
import os

import numpy as np
import torch

BASE = "models/abl"
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 ["moe", "llm", "both"]:
    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"]
        if "fusion.g_moe" in sd:
            gm.append(sd["fusion.g_moe"].item())
        if "fusion.g_llm" in sd:
            gl.append(sd["fusion.g_llm"].item())
    if gm:
        print(f"  g_moe: mean={np.mean(gm):+.4f}  |abs|=[{min(map(abs, gm)):.4f},{max(map(abs, gm)):.4f}]")
    if gl:
        print(f"  g_llm: mean={np.mean(gl):+.4f}  |abs|=[{min(map(abs, gl)):.4f},{max(map(abs, gl)):.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}")