"""Compare v1 / v2 / v4 / v5 on identical held-out inputs vs the frozen teacher.

Bench: 32 held-out RAVDESS files (4 per emotion) + funasr zh clip.
Gates: ground-truth 01/02 separation over ALL 336 emotional held-out files, and
Spearman rank corr with CREMA-D LO/MD/HI labels (in-train diagnostic).
Demo: one STRONG utterance per emotion + neutral + zh clip (audio+plot+json).
Saves /tmp/work/surgical_out_v5/v1_v2_v4_v5_compare.json
"""
import glob, json, os, sys
import numpy as np
import torch
import torch.nn.functional as F
import torchaudio
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

sys.path.insert(0, "/tmp/work/s2st_emotion")
from surgical_emotion import SurgicalConfig, SurgicalEmotionEncoder

FUNASR_DIR = "/tmp/work/emotion2vec_plus_base"
CKPTS = {"v1": "/tmp/work/v1_student/surgical_emotion_module.pt",
         "v2": "/tmp/work/v2_student/surgical_emotion_module.pt",
         "v4": "/tmp/work/surgical_out_v4/surgical_emotion_module.pt",
         "v5": "/tmp/work/surgical_out_v5/surgical_emotion_module.pt"}
OUT = "/tmp/work/surgical_out_v5/v1_v2_v4_v5_compare.json"
DEMO = "/tmp/work/surgical_out_v5/demo_v5"
DEV = "cuda:0"
LABELS = ["angry", "disgusted", "fearful", "happy", "neutral", "other", "sad", "surprised", "unk"]
EMO_FIELD = {"01": "neutral", "02": "calm", "03": "happy", "04": "sad",
             "05": "angry", "06": "fearful", "07": "disgust", "08": "surprised"}
EMO_IDX = [0, 1, 2, 3, 6, 7]
GATE_FLOOR = 0.02
with open("/tmp/work/surgical_out_v5/target_meta.json") as f:
    META = json.load(f)

torch.set_grad_enabled(False)
os.makedirs(DEMO, exist_ok=True)


def parse_label(fname):
    stem = os.path.splitext(fname)[0]
    if stem.count("-") >= 6:
        p = stem.split("-")
        if p[1] in ("01", "02"):
            return {"01": 0.0, "02": 1.0}.get(p[3])
    if "_" in stem:
        p = stem.split("_")
        if len(p) == 4:
            return {"LO": 0.0, "MD": 0.5, "HI": 1.0}.get(p[3])
    return None


def rankdata(x):
    order = np.argsort(x, kind="mergesort")
    sx = x[order]
    r = np.empty(len(x), dtype=np.float64)
    i = 0
    while i < len(sx):
        j = i
        while j + 1 < len(sx) and sx[j + 1] == sx[i]:
            j += 1
        r[i:j + 1] = (i + j) / 2.0 + 1.0
        i = j + 1
    out = np.empty(len(x), dtype=np.float64)
    out[order] = r
    return out


def spearman(a, b):
    ra, rb = rankdata(np.asarray(a, dtype=np.float64)), rankdata(np.asarray(b, dtype=np.float64))
    return float(np.corrcoef(ra, rb)[0, 1])


def load_wav(path, max_s=12.0):
    wav, sr = torchaudio.load(path)
    wav = wav.mean(0)
    if sr != 16000:
        wav = torchaudio.functional.resample(wav, sr, 16000)
    wav = wav[: int(max_s * 16000)]
    return wav / wav.abs().max().clamp(min=1e-4)


def rms_raw(wav, T):
    L = wav.numel(); n_full = (L + 319) // 320
    w = torch.zeros(T * 320); w[: min(L, T * 320)] = wav[: min(L, T * 320)]
    ms = w.view(T, 320).pow(2).mean(1).sqrt()
    if T > n_full: ms[n_full:] = ms[n_full - 1]
    if T >= 9:
        ms = F.conv1d(ms[None, None], torch.ones(9)[None, None] / 9.0, padding=4)[0, 0]
    return ms


def js_div(p, q, eps=1e-8):
    p = p.clamp(min=eps); q = q.clamp(min=eps)
    m = 0.5 * (p + q)
    kl = lambda a, b: float((a * (a.log() - b.log())).sum(-1))
    return 0.5 * kl(p, m) + 0.5 * kl(q, m)


def new_target(probs_t9, wav):
    p6 = probs_t9[:, EMO_IDX].max(-1).values.clamp(min=GATE_FLOOR).cpu()
    return p6 * (rms_raw(wav, probs_t9.size(0)) / META["global_ref"]).clamp(0, 1)


def demo_plot(name_tag, wav, out, probs_t9, tgt50, rate_hz, extra):
    fp = out["frame_probs"][0].cpu().numpy()
    inten = out["frame_intensity"][0, :, 0].cpu()
    tgt = F.interpolate(tgt50[None, None], size=inten.size(0), mode="area")[0, 0]
    tp = torch.softmax(probs_t9.mean(1), -1)
    d = os.path.join(DEMO, name_tag)
    os.makedirs(d, exist_ok=True)
    torchaudio.save(os.path.join(d, "audio.wav"), wav[None], 16000)
    t = np.arange(fp.shape[0]) * (1.0 / rate_hz)
    fig, ax = plt.subplots(3, 1, figsize=(12, 9), gridspec_kw={"height_ratios": [3, 1.6, 1.2]})
    colors = plt.cm.tab10(np.linspace(0, 1, 10))
    for k, lab in enumerate(LABELS[:8]):
        ax[0].plot(t, fp[:, k], label=lab, lw=1.6, color=colors[k])
    ax[0].set_title(f"v5 surgical module @ 6.25 Hz — {name_tag}")
    ax[0].set_ylabel("frame emotion probability")
    ax[0].legend(fontsize=8, ncol=8, loc="upper center", bbox_to_anchor=(0.5, 1.18))
    ax[0].set_ylim(-0.03, 1.05); ax[0].set_xlim(0, t[-1] + 1.0 / rate_hz)
    ax[1].plot(t, inten.numpy(), color="crimson", lw=2.0, label="v5 intensity")
    ax[1].plot(t, tgt.numpy(), color="steelblue", lw=1.5, ls="--", label="target (pre-calibration)")
    ax[1].set_ylabel("intensity"); ax[1].set_ylim(0, 1.05); ax[1].legend(fontsize=8)
    y = np.arange(8)
    ax[2].barh(y + 0.2, tp[:8].cpu().numpy(), height=0.38, color="gray", label="teacher")
    ax[2].barh(y - 0.2, out["utt_probs"][0][:8].cpu().numpy(), height=0.38,
               color="crimson", label="v5")
    ax[2].set_yticks(y, LABELS[:8], fontsize=8); ax[2].set_xlim(0, 1.05)
    ax[2].invert_yaxis(); ax[2].legend(fontsize=8, loc="lower right")
    ax[2].set_xlabel("utterance-level probability")
    ax[1].set_xlabel(f"time (s) — frame = 160 ms | {fp.shape[0]} frames @ 6.25 Hz (1 frame / token)")
    fig.tight_layout()
    fig.savefig(os.path.join(d, "analysis.png"), dpi=120, bbox_inches="tight")
    plt.close(fig)
    res = {"file": name_tag, "duration_s": round(wav.numel() / 16000, 2),
           "frame_rate_hz": rate_hz, "n_frames": int(fp.shape[0]),
           **extra,
           "v5_top": LABELS[int(out["utt_probs"][0].argmax())],
           "teacher_top": LABELS[int(tp.argmax())],
           "v5_probs": {l: round(float(p), 4) for l, p in zip(LABELS, out["utt_probs"][0])},
           "teacher_probs": {l: round(float(p), 4) for l, p in zip(LABELS, tp)},
           "intensity": {"mean": round(float(inten.mean()), 4),
                         "std": round(float(inten.std()), 4),
                         "min": round(float(inten.min()), 4),
                         "max": round(float(inten.max()), 4)}}
    with open(os.path.join(d, "result.json"), "w") as fj:
        json.dump(res, fj, indent=1)


def main():
    from funasr import AutoModel
    fun = AutoModel(model=FUNASR_DIR, trust_remote_code=False, device=DEV, disable_update=True)
    teacher = fun.model.eval()

    students = {}
    for name, ck in CKPTS.items():
        payload = torch.load(ck, map_location=DEV, weights_only=False)
        cfg = SurgicalConfig()
        cfg.group_size = int((payload.get("config") or {}).get("group_size", 4))
        enc = SurgicalEmotionEncoder(teacher, cfg).to(DEV)
        enc.adapter.load_state_dict(payload["adapter_state"])
        enc.eval()
        students[name] = enc
        print(f"loaded {name} (G={cfg.group_size}, {cfg.frame_rate_hz} Hz)", flush=True)

    files = sorted(glob.glob("/tmp/work/ravdess_kag/**/*.wav", recursive=True))
    speech = [f for f in files if os.path.basename(f).split("-")[1] == "01"
              and 19 <= int(os.path.basename(f).split("-")[-1].split(".")[0]) <= 24]
    by_emo = {}
    for f in speech:
        by_emo.setdefault(os.path.basename(f).split("-")[2], []).append(f)
    bench = []
    for e in sorted(by_emo):
        bench.extend(by_emo[e][:4])
    bench.append("/tmp/work/emotion2vec_plus_base/example/test.wav")
    print(f"comparison inputs: {len(bench)} files", flush=True)

    stats = {k: {"agree": [], "js": [], "pool_cos": [], "frame_cos": [],
                 "int_corr": [], "int_std": []} for k in students}
    report = {"n_files": len(bench), "per_student": {}, "files": {}}

    for f in bench:
        wav = load_wav(f).cpu()
        outs = {}
        for name in students:
            outs[name] = students[name](wav[None].to(DEV),
                                        torch.tensor([wav.numel()], device=DEV))
        src = F.layer_norm(wav[None], wav[None].shape).to(DEV)
        states = teacher.extract_features(src, padding_mask=None, mask=False)["x"]
        probs_t9 = torch.softmax(teacher.proj(states), -1)[0]
        tp = torch.softmax(teacher.proj(states.mean(1)), -1)
        tgt50 = new_target(probs_t9, wav)

        is_test = f.endswith("example/test.wav")
        bname = os.path.basename(f)
        emo = "zh-angry(test)" if is_test else EMO_FIELD.get(bname.split("-")[2], "?")
        row = {"dur_s": round(wav.numel() / 16000, 2),
               "corpus": "funasr-zh" if is_test else "RAVDESS-heldout",
               "teacher_top": LABELS[int(tp.argmax())]}
        for name, enc in students.items():
            o = outs[name]
            sp = o["utt_probs"][0]
            a = int(sp.argmax().item() == tp.argmax().item())
            j = js_div(sp[None], tp[None])
            pc = float(F.cosine_similarity(o["utt_emb"][0], states.mean(1)[0], dim=0))
            up = F.interpolate(o["frame_embs"].transpose(1, 2), size=states.size(1),
                               mode="nearest").transpose(1, 2)
            fc = float(F.cosine_similarity(up[0], states[0], dim=-1).mean())
            si = o["frame_intensity"][0, :, 0].cpu()
            tgt = F.interpolate(tgt50[None, None], size=si.size(0), mode="area")[0, 0]
            ic = (float(np.corrcoef(si.numpy(), tgt.numpy())[0, 1])
                  if si.var() > 1e-8 and tgt.var() > 1e-8 else None)
            stats[name]["agree"].append(a); stats[name]["js"].append(j)
            stats[name]["pool_cos"].append(pc); stats[name]["frame_cos"].append(fc)
            stats[name]["int_std"].append(float(si.std()))
            if ic is not None:
                stats[name]["int_corr"].append(ic)
            row[name] = {"top": LABELS[int(sp.argmax())], "p": round(float(sp.max()), 3),
                         "js": round(j, 5), "frame_cos": round(fc, 4),
                         "int_corr": round(ic, 3) if ic is not None else None,
                         "int_mean": round(float(si.mean()), 3)}
        report["files"][bname] = row

        if is_test:
            demo_plot("zh-angry-test", wav, outs["v5"], probs_t9, tgt50,
                      students["v5"].cfg.frame_rate_hz,
                      {"ravdess_label": emo, "intensity_label": "-"})

    # ground-truth separation for ALL students + per-emotion consistency for v5
    emotional = [f for f in speech if os.path.basename(f).split("-")[2] != "01"]
    sep = {k: {"01": [], "02": []} for k in students}
    per_emo_v5 = {}
    for f in emotional:
        wav = load_wav(f)
        b = os.path.basename(f)
        lbl = b.split("-")[3]
        for name, enc in students.items():
            o = enc(wav[None].to(DEV), torch.tensor([wav.numel()], device=DEV))
            sep[name][lbl].append(float(o["frame_intensity"][0, :, 0].mean()))
            if name == "v5":
                per_emo_v5.setdefault(b.split("-")[2], {"01": [], "02": []})[lbl].append(
                    float(o["frame_intensity"][0, :, 0].mean()))
    for name in students:
        a, b = np.array(sep[name]["01"]), np.array(sep[name]["02"])
        d = (b.mean() - a.mean()) / np.sqrt((a.var() + b.var()) / 2 + 1e-9)
        stats[name]["sep"] = {"normal_mean": round(float(a.mean()), 4),
                              "strong_mean": round(float(b.mean()), 4),
                              "diff": round(float(b.mean() - a.mean()), 4),
                              "cohen_d": round(float(d), 3)}
        print(f"sep {name}: normal {a.mean():.3f} strong {b.mean():.3f} "
              f"diff {b.mean()-a.mean():+.3f} d={d:.2f}", flush=True)
    emo_cons = {EMO_FIELD.get(e, e): round(float(np.mean(g["02"]) - np.mean(g["01"])), 4)
                for e, g in sorted(per_emo_v5.items()) if g["01"] and g["02"]}
    n_pos = sum(1 for v in emo_cons.values() if v > 0)
    print(f"v5 per-emotion strong>normal: {n_pos}/{len(emo_cons)}  {emo_cons}", flush=True)

    # CREMA-D Spearman for every student
    cremad = [f for f in sorted(glob.glob("/tmp/work/cremad/**/*.wav", recursive=True))
              if parse_label(os.path.basename(f)) is not None]
    print(f"CREMA-D labeled: {len(cremad)}", flush=True)
    for name, enc in students.items():
        preds, labs = [], []
        for f in cremad:
            wav = load_wav(f)
            if wav.numel() < 8000: continue
            o = enc(wav[None].to(DEV), torch.tensor([wav.numel()], device=DEV))
            preds.append(float(o["frame_intensity"][0, :, 0].mean()))
            labs.append(parse_label(os.path.basename(f)))
        stats[name]["spearman"] = round(spearman(preds, labs), 4)
        print(f"spearman {name}: {stats[name]['spearman']} (n={len(labs)})", flush=True)

    for name in students:
        s = stats[name]
        report["per_student"][name] = {
            "top1_agreement": float(np.mean(s["agree"])),
            "mean_js": float(np.mean(s["js"])),
            "mean_pooled_cos": float(np.mean(s["pool_cos"])),
            "mean_frame_cos": float(np.mean(s["frame_cos"])),
            "mean_intensity_corr_newtarget": float(np.mean(s["int_corr"])) if s["int_corr"] else None,
            "mean_intensity_std": float(np.mean(s["int_std"])),
            "intensity_separation": s["sep"],
            "cremad_spearman": s.get("spearman"),
        }
    report["v5_per_emotion_strong_minus_normal"] = emo_cons
    report["v5_per_emotion_consistency"] = f"{n_pos}/{len(emo_cons)}"
    with open(OUT, "w") as fp:
        json.dump(report, fp, indent=1)
    print(json.dumps(report["per_student"], indent=1), flush=True)
    print("COMPARE_DONE", flush=True)


if __name__ == "__main__":
    main()