"""Standalone v5 demo pass: one STRONG utterance per emotion + neutral + zh clip."""
import glob, json, os, sys
import torch
import torch.nn.functional as F
sys.path.insert(0, "/tmp/work/s2st_emotion")
from compare_v5 import (demo_plot, new_target, load_wav, EMO_FIELD)

def main():
    from funasr import AutoModel
    from surgical_emotion import SurgicalConfig, SurgicalEmotionEncoder
    DEV = "cuda:0"
    fun = AutoModel(model="/tmp/work/emotion2vec_plus_base", trust_remote_code=False,
                    device=DEV, disable_update=True)
    teacher = fun.model.eval()
    payload = torch.load("/tmp/work/surgical_out_v5/surgical_emotion_module.pt",
                         map_location=DEV, weights_only=False)
    cfg = SurgicalConfig()
    cfg.group_size = int((payload.get("config") or {}).get("group_size", 8))
    enc = SurgicalEmotionEncoder(teacher, cfg).to(DEV)
    enc.adapter.load_state_dict(payload["adapter_state"]); enc.eval()
    torch.set_grad_enabled(False)

    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)
    demo_files = []
    for e in sorted(by_emo):
        if e == "01": continue
        strong = [f for f in by_emo[e] if os.path.basename(f).split("-")[3] == "02"]
        demo_files.append(strong[0] if strong else by_emo[e][0])
    neutral = by_emo.get("01", [])
    if neutral: demo_files = [neutral[0]] + demo_files
    for f in demo_files:
        wav = load_wav(f).cpu()
        b = os.path.basename(f)
        out = enc(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]
        tgt50 = new_target(probs_t9, wav)
        emo = EMO_FIELD.get(b.split("-")[2], "?")
        ilbl = b.split("-")[3]
        demo_plot(f"{emo}_{b[:20]}", wav, out, probs_t9, tgt50, cfg.frame_rate_hz,
                  {"ravdess_label": emo,
                   "intensity_label": "strong" if ilbl == "02" else "normal"})
        print("demo done:", emo, flush=True)
    print("DEMOS_DONE", flush=True)

if __name__ == "__main__":
    main()