#!/usr/bin/env python3
"""
Score a JSONL of {gen_wav, target_emotion} rows with the firdhokk
Whisper-Large-v3 SER classifier and emit per-row class probabilities.

Replaces the broken in-process emotion2vec_plus_base path. firdhokk is
a Whisper-LV3 fine-tune on RAVDESS+SAVEE+TESS+URDU; multi-corpus
training avoids the saturation failure mode emotion2vec exhibits.

Verified 2026-04-30: 3/5 correct on real RAVDESS clips, 2/5 in known
confusion pairs (happy↔surprised, sad↔fearful). Probabilities are
non-saturated.

Run from any dir; the venv path is hardcoded.

Input JSONL rows (any other fields preserved/ignored):
  {"gen_wav": "/path/to/out.wav", "target_emotion": "happy"}

Output JSONL rows:
  {"gen_wav": "...", "target_emotion": "happy",
   "top_emotion": "neutral", "top_prob": 0.80,
   "target_prob": 0.006, "match": false,
   "all_probs": {"angry": 0.19, "calm": ..., ...}}

Class set (7): angry, disgust, fearful, happy, neutral, sad, surprised.
("calm" maps to neutral; "disgusted" maps to disgust.)
"""

import argparse
import json
import sys
import os

import numpy as np
import soundfile as sf
import torch
from transformers import AutoFeatureExtractor, AutoModelForAudioClassification


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--in", dest="input", required=True, help="JSONL input")
    ap.add_argument("--out", required=True, help="JSONL output")
    ap.add_argument(
        "--model",
        default="firdhokk/speech-emotion-recognition-with-openai-whisper-large-v3",
        help="HF model id; default = firdhokk Whisper-LV3 (sanity-verified)",
    )
    args = ap.parse_args()

    fe = AutoFeatureExtractor.from_pretrained(args.model)
    m = AutoModelForAudioClassification.from_pretrained(args.model)
    m.eval()
    labels = m.config.id2label
    label_to_idx = {v.lower(): k for k, v in labels.items()}

    # CLI emotion-name aliases.
    ALIAS = {
        "calm": "neutral",
        "disgusted": "disgust",
        "excited": "happy",  # firdhokk has no "excited"; closest semantic match
    }

    def find_idx(name: str):
        n = ALIAS.get(name.strip().lower(), name.strip().lower())
        return label_to_idx.get(n)

    n_total = 0
    n_match = 0
    with open(args.input) as f, open(args.out, "w") as out:
        for line in f:
            line = line.strip()
            if not line:
                continue
            row = json.loads(line)
            gen_wav = row.get("gen_wav")
            target = row.get("target_emotion") or row.get("emotion")
            if not gen_wav or not target:
                print(f"skip (missing fields): {row}", file=sys.stderr)
                continue
            if not os.path.exists(gen_wav):
                print(f"skip (missing wav): {gen_wav}", file=sys.stderr)
                continue

            audio, sr = sf.read(gen_wav)
            if audio.ndim > 1:
                audio = audio.mean(axis=1)
            if sr != 16000:
                target_len = int(len(audio) * 16000 / sr)
                idx = np.linspace(0, len(audio) - 1, target_len).astype(int)
                audio = audio[idx]

            inputs = fe(audio, sampling_rate=16000, return_tensors="pt")
            with torch.no_grad():
                logits = m(**inputs).logits
            probs = torch.softmax(logits, dim=-1)[0].cpu().numpy()
            top_idx = int(probs.argmax())
            top_label = labels[top_idx]
            target_idx = find_idx(target)
            target_prob = float(probs[target_idx]) if target_idx is not None else None
            all_probs = {labels[i]: float(probs[i]) for i in range(len(labels))}

            out_row = {
                **row,
                "top_emotion": top_label,
                "top_prob": float(probs[top_idx]),
                "target_prob": target_prob,
                "match": (target_idx == top_idx) if target_idx is not None else None,
                "all_probs": all_probs,
            }
            out.write(json.dumps(out_row) + "\n")
            n_total += 1
            if target_idx == top_idx:
                n_match += 1

            tp_str = f"{target_prob:.3f}" if target_prob is not None else "n/a"
            print(
                f"  {gen_wav.split('/')[-1]:35s} target={target:10s} → {top_label:10s} "
                f"(top={probs[top_idx]:.3f}, target={tp_str})",
                file=sys.stderr,
            )

    print(f"\n=== summary ===", file=sys.stderr)
    print(f"rows: {n_total}", file=sys.stderr)
    print(f"top-1 match: {n_match}/{n_total} = {n_match / max(n_total, 1):.1%}", file=sys.stderr)
    print(f"wrote {args.out}", file=sys.stderr)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())