#!/usr/bin/env python3
"""End-to-end demo: audio file(s) -> all 61 VoiceClap attribute scores.

    python example_inference.py examples/clip_a.mp3 examples/clip_b.mp3
"""
import sys

from voiceclap_heads import AttributeScorer

GROUPS = {
    "speaker / voice": ["Gender", "Age", "High-Pitched_vs._Low-Pitched", "Monotone_vs._Expressive",
                        "Soft_vs._Harsh", "Warm_vs._Cold", "Confident_vs._Hesitant",
                        "Submissive_vs._Dominant", "Serious_vs._Humorous",
                        "Vulnerable_vs._Emotionally_Detached", "talking_speed", "duration"],
    "recording quality": ["Recording_Quality", "Background_Noise", "Authenticity",
                          "score_overall_quality", "score_speech_quality",
                          "score_background_quality", "score_content_enjoyment"],
    "core affect": ["Valence", "Arousal"],
}


def main(paths):
    scorer = AttributeScorer()          # downloads heads.pt + the frozen encoder
    print(f"loaded {len(scorer.dims)} heads on {scorer.device}\n")

    for p in paths:
        s = scorer.score(p)             # dict: 61 dimension -> float
        print("=" * 78)
        print(p)
        print("=" * 78)
        shown = set()
        for title, dims in GROUPS.items():
            print(f"-- {title}")
            for d in dims:
                print(f"   {d:<38s} {s[d]:7.3f}")
                shown.add(d)
        rest = sorted((k for k in s if k not in shown), key=lambda k: -s[k])
        print("-- top 8 emotion dimensions")
        for d in rest[:8]:
            print(f"   {d:<38s} {s[d]:7.3f}")
        print()

    # ---- batched variant: one encoder pass, one head pass, for many clips
    if len(paths) > 1:
        batch = scorer.score(paths, dims=["Valence", "Arousal", "Gender", "Age"], batch_size=32)
        print("batched:", {k: [round(float(x), 3) for x in v] for k, v in batch.items()})


if __name__ == "__main__":
    main(sys.argv[1:] or ["examples/clip_a.mp3", "examples/clip_b.mp3"])