"""Validation suite v5 (6.25 Hz + weak label supervision).

Checks: 6.25 Hz rate contract, teacher fidelity, intensity corr, the acceptance
gate (normal 01 vs strong 02 separation over ALL 336 emotional held-out files,
plus the target ceiling), NEW per-emotion separation consistency, and NEW
Spearman rank correlation with CREMA-D LO/MD/HI labels (in-train diagnostic).
Writes /tmp/work/surgical_out_v5/eval_report_v5.json + dynamics_plot_v5.png
"""
from __future__ import annotations
import glob, json, os, random, sys
import numpy as np
import torch
import torch.nn.functional as F
import torchaudio

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from surgical_emotion import SurgicalEmotionEncoder

FUNASR_DIR = "/tmp/work/emotion2vec_plus_base"
OUT_DIR = "/tmp/work/surgical_out_v5"
CKPT = os.path.join(OUT_DIR, "surgical_emotion_module.pt")
RAVDESS_DIR = "/tmp/work/ravdess_kag"
CREMAD_DIR = "/tmp/work/cremad"
EMO_IDX = [0, 1, 2, 3, 6, 7]
GATE_FLOOR = 0.02
EMO_FIELD = {"01": "neutral", "02": "calm", "03": "happy", "04": "sad",
             "05": "angry", "06": "fearful", "07": "disgust", "08": "surprised"}


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 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 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 new_target(enc, wav, states):
    probs = torch.softmax(enc.teacher.proj(states), -1)[0]
    p6 = probs[:, EMO_IDX].max(-1).values.clamp(min=GATE_FLOOR).cpu()
    env = (rms_raw(wav, states.size(1)) / META["global_ref"]).clamp(0, 1)
    return p6 * env, probs


def main():
    device = "cuda:0"
    enc = SurgicalEmotionEncoder.from_pretrained(FUNASR_DIR, CKPT, device=device)
    enc.eval()
    report = {"frame_rate_hz": enc.cfg.frame_rate_hz, "checks": {}, "intensity": {}}

    ok = True
    for sec in [1.0, 2.5, 5.0, 10.0]:
        n = int(sec * 16000)
        wav = (torch.randn(1, n, device=device)); wav = wav / wav.abs().max().clamp(min=1e-3)
        out = enc(wav, torch.tensor([n], device=device))
        expect = int(np.ceil(sec * enc.cfg.frame_rate_hz))
        good = abs(int(out["frame_embs"].shape[1]) - expect) <= 1
        ok &= good
        report["checks"][f"{sec}s"] = {"expected~": expect, "got": int(out["frame_embs"].shape[1]), "ok": good}
    report["rate_contract_ok"] = bool(ok)
    print("rate contract ok:", ok)

    files = sorted(glob.glob(os.path.join(RAVDESS_DIR, "**", "*.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]
    emotional = [f for f in speech if os.path.basename(f).split("-")[2] != "01"]
    random.seed(0)
    sel = random.sample(speech, min(60, len(speech)))

    agree, jss, corrs = [], [], []
    plot_sample = None
    with torch.no_grad():
        for f in sel:
            wav = load_wav(f)
            if wav.numel() < 8000: continue
            out = enc(wav[None].to(device), torch.tensor([wav.numel()], device=device))
            states = enc.teacher.extract_features(
                F.layer_norm(wav[None].to(device), wav[None].to(device).shape),
                padding_mask=None, mask=False)["x"]
            tp = torch.softmax(enc.teacher.proj(states.mean(1)), -1)[0]
            sp = out["utt_probs"][0]
            agree.append(int(sp.argmax().item() == tp.argmax().item()))
            jss.append(js_div(sp[None], tp[None]))
            tgt50, _ = new_target(enc, wav, states)
            si = out["frame_intensity"][0, :, 0].cpu()
            tgt = F.interpolate(tgt50[None, None], size=si.size(0), mode="area")[0, 0]
            if si.var() > 1e-8 and tgt.var() > 1e-8:
                corrs.append(float(np.corrcoef(si.numpy(), tgt.numpy())[0, 1]))
            if plot_sample is None and "-05-" in f:
                plot_sample = (os.path.basename(f), out, tgt50)
    report["intensity"]["fidelity_agreement"] = float(np.mean(agree))
    report["intensity"]["fidelity_mean_js"] = float(np.mean(jss))
    report["intensity"]["mean_corr_new_target"] = float(np.mean(corrs))
    print(f"agreement={report['intensity']['fidelity_agreement']:.3f}  "
          f"JS={report['intensity']['fidelity_mean_js']:.4f}  "
          f"int_corr={report['intensity']['mean_corr_new_target']:.3f}")

    # ground-truth separation + per-emotion consistency + target ceiling
    groups = {"01": [], "02": []}
    groups_tgt = {"01": [], "02": []}
    per_emo = {}
    with torch.no_grad():
        for f in emotional:
            wav = load_wav(f)
            if wav.numel() < 8000: continue
            out = enc(wav[None].to(device), torch.tensor([wav.numel()], device=device))
            states = enc.teacher.extract_features(
                F.layer_norm(wav[None].to(device), wav[None].to(device).shape),
                padding_mask=None, mask=False)["x"]
            tgt50, _ = new_target(enc, wav, states)
            b = os.path.basename(f)
            lbl = b.split("-")[3]; ecode = b.split("-")[2]
            mi = float(out["frame_intensity"][0, :, 0].mean())
            groups[lbl].append(mi)
            groups_tgt[lbl].append(float(tgt50.mean()))
            per_emo.setdefault(ecode, {"01": [], "02": []})[lbl].append(mi)
    a, b = np.array(groups["01"]), np.array(groups["02"])
    d = (b.mean() - a.mean()) / np.sqrt((a.var() + b.var()) / 2 + 1e-9)
    report["intensity"]["sep_normal_mean"] = float(a.mean())
    report["intensity"]["sep_strong_mean"] = float(b.mean())
    report["intensity"]["sep_diff"] = float(b.mean() - a.mean())
    report["intensity"]["sep_cohen_d"] = float(d)
    ta, tb = np.array(groups_tgt["01"]), np.array(groups_tgt["02"])
    td = (tb.mean() - ta.mean()) / np.sqrt((ta.var() + tb.var()) / 2 + 1e-9)
    report["intensity"]["target_sep_cohen_d"] = float(td)
    print(f"intensity separation: normal {a.mean():.3f} vs strong {b.mean():.3f}  "
          f"diff {b.mean()-a.mean():+.3f}  Cohen-d {d:.2f}  (target ceiling d={td:.2f})")

    emo_cons = {}
    for ecode, g in sorted(per_emo.items()):
        if len(g["01"]) and len(g["02"]):
            diff = float(np.mean(g["02"]) - np.mean(g["01"]))
            emo_cons[EMO_FIELD.get(ecode, ecode)] = round(diff, 4)
    report["intensity"]["per_emotion_strong_minus_normal"] = emo_cons
    n_pos = sum(1 for v in emo_cons.values() if v > 0)
    report["intensity"]["per_emotion_consistency"] = f"{n_pos}/{len(emo_cons)}"
    print(f"per-emotion strong>normal: {n_pos}/{len(emo_cons)}  {emo_cons}")

    # CREMA-D Spearman (LO/MD/HI) — in-train diagnostic
    cremad = [f for f in sorted(glob.glob(os.path.join(CREMAD_DIR, "**", "*.wav"), recursive=True))
              if parse_label(os.path.basename(f)) is not None]
    preds, labs = [], []
    with torch.no_grad():
        for f in cremad:
            wav = load_wav(f)
            if wav.numel() < 8000: continue
            out = enc(wav[None].to(device), torch.tensor([wav.numel()], device=device))
            preds.append(float(out["frame_intensity"][0, :, 0].mean()))
            labs.append(parse_label(os.path.basename(f)))
    sp = spearman(preds, labs)
    report["intensity"]["cremad_spearman"] = sp
    report["intensity"]["cremad_n"] = len(labs)
    print(f"CREMA-D Spearman (LO/MD/HI): {sp:.3f}  (n={len(labs)}, in-train diagnostic)")

    # dynamics plot
    try:
        import matplotlib; matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        name, out, tgt50 = plot_sample
        fp = out["frame_probs"][0].cpu().numpy()
        si = out["frame_intensity"][0, :, 0].cpu()
        tgt = F.interpolate(tgt50[None, None], size=si.size(0), mode="area")[0, 0]
        t = np.arange(si.size(0)) * 0.16
        fig, ax = plt.subplots(2, 1, figsize=(11, 6))
        for k, lab in enumerate(["angry", "disgusted", "fearful", "happy",
                                 "neutral", "other", "sad", "surprised"]):
            ax[0].plot(t, fp[:, k], label=lab, lw=1.2)
        ax[0].set_title(f"v5 frame emotion probs @6.25Hz — {name}")
        ax[0].legend(fontsize=7, ncol=4); ax[0].set_ylabel("p")
        ax[1].plot(t, si, color="crimson", lw=1.8, label="v5 intensity")
        ax[1].plot(t, tgt, color="steelblue", lw=1.4, ls="--", label="target (pre-calibration)")
        ax[1].set_ylabel("intensity"); ax[1].set_ylim(0, 1.05); ax[1].legend(fontsize=8)
        ax[1].set_xlabel("time (s) — frame = 160 ms | 6.25 Hz (1 frame / dots.tts token)")
        fig.tight_layout()
        fig.savefig(os.path.join(OUT_DIR, "dynamics_plot_v5.png"), dpi=110)
        print("plot saved")
    except Exception as e:
        print("plot skipped:", e)

    with open(os.path.join(OUT_DIR, "eval_report_v5.json"), "w") as f:
        json.dump(report, f, indent=1)
    print("EVAL_DONE")


if __name__ == "__main__":
    with open(os.path.join(OUT_DIR, "target_meta.json")) as f:
        META = json.load(f)
    main()