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

Checkpoints: v1 & v2 downloaded from HF (amir0907/student, amir0907/student-v2),
v3 freshly trained at /tmp/work/surgical_out.
Test set: RAVDESS actors 19-24 speech (never in ANY training pool) + the funasr
Chinese clip. Also emits a per-sample demo (audio + plot + json) for v3.
Saves /tmp/work/surgical_out/v1_v2_v3_compare.json
"""

import glob
import json
import os
import shutil
import sys
import time
import warnings

warnings.filterwarnings("ignore")
import logging

logging.disable(logging.CRITICAL)

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",
    "v3": "/tmp/work/surgical_out/surgical_emotion_module.pt",
}
OUT = "/tmp/work/surgical_out/v1_v2_v3_compare.json"
DEMO = "/tmp/work/surgical_out/demo_v3"
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"}

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


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 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 rms_env(wav, T):
    frames = torch.zeros(T)
    for i in range(T):
        seg = wav[i * 320: (i + 1) * 320]
        frames[i] = float(seg.pow(2).mean().sqrt()) if seg.numel() else (frames[i - 1] if i else 0.0)
    if T >= 9:
        frames = F.conv1d(frames[None, None], torch.ones(9)[None, None] / 9.0, padding=4)[0, 0]
    lo, hi = frames.min(), frames.max()
    return torch.full_like(frames, 0.5) if hi - lo < 1e-6 else (frames - lo) / (hi - lo)


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():
        enc = SurgicalEmotionEncoder(teacher, SurgicalConfig()).to(DEV)
        enc.adapter.load_state_dict(torch.load(ck, map_location=DEV, weights_only=False)["adapter_state"])
        enc.eval()
        students[name] = enc
        print(f"loaded {name}", flush=True)

    files = sorted(glob.glob("/tmp/work/ravdess_kag/**/*.wav", recursive=True))
    files = [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 files:
        by_emo.setdefault(os.path.basename(f).split("-")[2], []).append(f)
    bench = []
    for e in sorted(by_emo):
        bench.extend(by_emo[e][:4])           # 32 held-out files
    bench.append("/tmp/work/emotion2vec_plus_base/example/test.wav")  # cross-lingual
    print(f"comparison inputs: {len(bench)} files", flush=True)

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

    def demo_plot(name_tag, src_path, wav, out, states, lengths, tp, tgt):
        fp = out["frame_probs"][0].cpu().numpy()
        inten = out["frame_intensity"][0, :, 0].cpu()
        T = int(lengths[0])
        tgt_ds = F.interpolate(tgt[None, None], size=inten.size(0), mode="linear")[0, 0]
        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]) * 0.08
        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"v3 surgical module @ 12.5 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] + 0.08)
        ax[1].plot(t, inten.numpy(), color="crimson", lw=2.0, label="v3 intensity")
        ax[1].plot(t, tgt_ds.numpy(), color="steelblue", lw=1.5, ls="--", label="target")
        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="v3")
        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 = 80 ms | {fp.shape[0]} frames @ 12.5 Hz")
        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),
               "n_frames": int(fp.shape[0]),
               "v3_top": LABELS[int(out["utt_probs"][0].argmax())],
               "teacher_top": LABELS[int(tp.argmax())],
               "v3_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)

    for f in bench:
        wav = load_wav(f).cpu()
        dur = wav.numel() / 16000
        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)
        r = teacher.extract_features(src, padding_mask=None, mask=False)
        states = r["x"].detach()
        lengths = torch.tensor([states.size(1)], device=DEV)
        tp = torch.softmax(teacher.proj(states.mean(1)), -1)[0]

        T = int(lengths[0])
        tgt = rms_env(wav, T)
        tmax = torch.softmax(teacher.proj(states[0]), -1)[:, :8].max(-1).values.cpu().clamp(0, 1).sqrt()
        tgt = tmax * tgt

        is_test_wav = f.endswith("example/test.wav")
        bname = os.path.basename(f)
        emo = "zh-angry(test)" if is_test_wav else EMO_FIELD.get(bname.split("-")[2], "?")
        row = {"dur_s": round(dur, 2), "corpus": "funasr-zh" if is_test_wav else "RAVDESS-heldout",
               "teacher_top": LABELS[int(tp.argmax())]}
        for name in students:
            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_ds = F.interpolate(tgt[None, None], size=si.size(0), mode="linear")[0, 0]
            ic = float(np.corrcoef(si.numpy(), tgt_ds.numpy())[0, 1]) if si.var() > 1e-6 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]["inten_std"].append(float(si.std()))
            if ic is not None:
                stats[name]["inten_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)}
        report["files"][bname] = row

        # demo for v3: every emotion variant + the chinese clip
        if is_test_wav or bname.split("-")[3] == "02":
            tag = ("zh-angry-test" if is_test_wav else f"{emo}_{bname[:20]}")
            demo_plot(tag, f, wav, outs["v3"], states, lengths, tp, tgt)

    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_std": float(np.mean(s["inten_std"])),
            "mean_intensity_corr": float(np.mean(s["inten_corr"])) if s["inten_corr"] else None,
        }
    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()