"""Build a demo gallery for the surgical emotion module (student model).

For each sample wav: run the student, save into
  /tmp/surgical_emotion_demo/<sample_name>/
    - <sample_name>.wav          (the input audio itself, copied)
    - <sample_name>_analysis.png (frame probs @12.5Hz, intensity vs target, teacher-vs-student bars)
    - <sample_name>_result.json  (numeric outputs summary)
plus a top-level README.md index.
"""

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

warnings.filterwarnings("ignore")
import logging

logging.disable(logging.CRITICAL)

sys.path.insert(0, "/tmp/work/s2st_emotion")
import numpy as np
import torch
import torch.nn.functional as F
import torchaudio

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

from surgical_emotion import SurgicalEmotionEncoder

FUNASR_DIR = "/tmp/work/emotion2vec_plus_base"
CKPT = "/tmp/work/surgical_out/surgical_emotion_module.pt"
GALLERY = "/tmp/surgical_emotion_demo"
DEV = "cuda:0"
LABELS = ["angry", "disgusted", "fearful", "happy", "neutral", "other", "sad", "surprised"]
EMO_FIELD = {"01": "neutral", "02": "calm", "03": "happy", "04": "sad",
             "05": "angry", "06": "fearful", "07": "disgust", "08": "surprised"}

os.makedirs(GALLERY, exist_ok=True)
enc = SurgicalEmotionEncoder.from_pretrained(FUNASR_DIR, CKPT, device=DEV)
enc.eval()
torch.set_grad_enabled(False)


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_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 pick_samples():
    files = sorted(glob.glob("/tmp/work/ravdess_kag/Actor_*/*.wav"))
    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]
    strong = [f for f in files if os.path.basename(f).split("-")[3] == "02"]  # strong intensity
    pool = strong if len(strong) >= 6 else files
    picked, seen = [], set()
    for f in pool:  # one strong example per emotion (first occurrence)
        e = os.path.basename(f).split("-")[2]
        if e in ("01", "03", "04", "05", "06", "08") and e not in seen:
            seen.add(e)
            picked.append(f)
    picked = [p for _, p in sorted(zip([p.split("-")[2] for p in picked], picked))]
    neutral = [f for f in files if os.path.basename(f).split("-")[2] == "01"]
    if neutral:
        picked = [neutral[0]] + picked  # neutral has no strong variant; add one from full pool
    extra = "/tmp/work/emotion2vec_plus_base/example/test.wav"
    return picked + [extra]


def process(path):
    wav = load_wav(path).cpu()
    dur = float(wav.numel()) / 16000
    out = enc(wav[None].to(DEV), torch.tensor([wav.numel()], device=DEV))

    states, lengths = enc.extract_teacher_states(wav[None].to(DEV),
                                                 torch.tensor([wav.numel()], device=DEV))
    t_probs = torch.softmax(enc.teacher.proj(states.mean(1)), -1)[0]

    fp = out["frame_probs"][0].cpu().numpy()          # (T*, 9)
    inten = out["frame_intensity"][0, :, 0].cpu()
    T = int(lengths[0])
    tgt = rms_env(wav, T)
    tmax = torch.softmax(enc.teacher.proj(states[0]), -1)[:, :8].max(-1).values.cpu().clamp(0, 1).sqrt()
    tgt = tmax * tgt
    tgt_ds = F.interpolate(tgt[None, None], size=inten.size(0), mode="linear")[0, 0]

    name = os.path.splitext(os.path.basename(path))[0]
    emo = EMO_FIELD.get(name.split("-")[2], "zh-angry(test)") if "-" in name else "zh-angry(test)"
    sample_dir = os.path.join(GALLERY, f"{emo}_{name[:24]}")
    os.makedirs(sample_dir, exist_ok=True)
    shutil.copy(path, os.path.join(sample_dir, f"{name[:24]}.wav"))

    # ---------------- plot ----------------
    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):
        ax[0].plot(t, fp[:, k], label=lab, lw=1.6, color=colors[k])
    ax[0].set_title(f"Surgical emotion module @ 12.5 Hz — {name}.wav  (RAVDESS label: {emo}, {dur:.1f}s)")
    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[1].plot(t, inten.numpy(), color="crimson", lw=2.0, label="student intensity")
    ax[1].plot(t, tgt_ds.numpy(), color="steelblue", lw=1.5, ls="--", label="target (conf-gated energy)")
    ax[1].set_ylabel("intensity")
    ax[1].set_ylim(0, 1.05)
    ax[1].set_xlim(0, t[-1] + 0.08)          # panels 1-2 share the time range
    ax[0].set_xlim(0, t[-1] + 0.08)
    ax[1].legend(fontsize=8, ncol=2)

    y = np.arange(len(LABELS))
    ax[2].barh(y + 0.2, t_probs[:8].cpu().numpy(), height=0.38, color="gray", label="teacher (utterance)")
    ax[2].barh(y - 0.2, out["utt_probs"][0][:8].cpu().numpy(), height=0.38, color="crimson", label="student (utterance)")
    ax[2].set_yticks(y, LABELS, fontsize=8)
    ax[2].set_xlim(0, 1.05)                   # independent probability axis
    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(sample_dir, f"{name[:24]}_analysis.png"), dpi=120, bbox_inches="tight")
    plt.close(fig)

    # ---------------- json ----------------
    res = {
        "file": name + ".wav",
        "ravdess_label": emo,
        "duration_s": round(dur, 2),
        "frame_rate_hz": 12.5,
        "n_frames": int(fp.shape[0]),
        "student_top": LABELS[int(out["utt_probs"][0].argmax())],
        "student_probs": {l: round(float(p), 4) for l, p in zip(LABELS, out["utt_probs"][0][:8])},
        "teacher_top": LABELS[int(t_probs.argmax())],
        "teacher_probs": {l: round(float(p), 4) for l, p in zip(LABELS, t_probs[:8])},
        "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)},
        "per_frame": {
            "emotion_probs_12.5hz": [[round(float(v), 4) for v in row] for row in fp[:, :8]],
            "intensity_12.5hz": [round(float(v), 4) for v in inten],
        },
    }
    with open(os.path.join(sample_dir, f"{name[:24]}_result.json"), "w") as fjson:
        json.dump(res, fjson, indent=1)
    return sample_dir, res


index = []
for path in pick_samples():
    d, res = process(path)
    index.append(res)
    print("done:", os.path.basename(d), "->", res["student_top"], res["teacher_top"], flush=True)

with open(os.path.join(GALLERY, "README.md"), "w") as f:
    f.write("# Surgical Emotion Module — Demo Gallery\n\n"
            "Student = emotion2vec_plus_base backbone (frozen) + surgical adapter -> "
            "continuous vectors @ 12.5 Hz + per-frame emotion probs + intensity.\n\n"
            "| sample | duration | RAVDESS label | teacher top | student top | intensity range |\n|---|---|---|---|---|---|\n")
    for r in index:
        ii = r["intensity"]
        f.write(f"| {r['file'][:28]} | {r['duration_s']}s | {r['ravdess_label']} | "
                f"{r['teacher_top']} | {r['student_top']} | [{ii['min']}, {ii['max']}] |\n")
print("GALLERY_DONE", flush=True)