#!/usr/bin/env python3
"""emotion_ontology — CLAP emotion scoring via synonym CLUSTERS on the V/A circumplex.

A single prompt per emotion is brittle: "euphoric" and "ecstatic festival rush" and
"blissful and radiant" all point at the same feeling, but any one phrase lands in a
slightly different spot in CLAP's text space, so a lone prompt over/under-fires. So we
give each emotion a CLUSTER of angles — adjective pairs, scene metaphors, bodily/spatial
framings, EN + FR — and marginalize the cluster down to one robust per-emotion score
(mean cosine over the cluster, size-invariant). Same trick the family resolver uses to
fold descriptors → family ([[reference_audio_feature_stack]]); here it folds synonyms →
emotion.

Each emotion also carries a **valence** (−1 unpleasant … +1 pleasant) and **arousal**
(−1 calm … +1 energetic) coordinate (Russell's circumplex). From the emotion distribution
we read a continuous V/A point — the scalar pair a visuals engine (hydra-live-hexa) wants
to drive palette/intensity, alongside the discrete top-emotion label.

This is the per-window kernel: feed it ONE CLAP audio embedding and it returns the
emotion read. The track-timeline builder slides it over a master to emit an emotion lane.

    python3 emotion_ontology.py wheel                  # print the ontology + V/A map
    python3 emotion_ontology.py tag <folder|file>      # emotion read for a sound
    python3 emotion_ontology.py probe                  # archetype sanity check (katana)

CLAP proposes, PLN's ear confirms ([[feedback_build_katana_first]])."""
import sys
from pathlib import Path

import numpy as np

import sample_semantics as SEM

# ── the emotion wheel: name → (valence, arousal, synonym/angle cluster) ─────────
# Prompts are phrased to complete sample_semantics' "this is the sound of ___" frame
# naturally (same frame the audio embeds are compared against — consistency > a
# marginally better standalone frame). Angles per emotion deliberately VARY: a plain
# adjective pair, a scene/metaphor, a bodily/spatial image, and a FR synonym, so the
# cluster brackets the feeling from several directions instead of restating one word.
EMOTIONS = {
    "euphoric": (0.90, 0.85, [
        "euphoric uplifting music", "ecstatic hands-in-the-air rave energy",
        "blissful elated and radiant", "a triumphant festival anthem",
        "joyful exhilarating release", "musique lumineuse et euphorique"]),
    "triumphant": (0.80, 0.60, [
        "triumphant and epic music", "heroic victorious and grand",
        "a soaring cinematic climax", "majestic and powerful",
        "an anthemic uplifting swell", "musique héroïque et triomphante"]),
    "playful": (0.70, 0.50, [
        "playful and quirky music", "cheeky bouncy and fun",
        "a cartoonish whimsical romp", "lighthearted and mischievous",
        "goofy upbeat energy", "musique espiègle et ludique"]),
    "warm": (0.60, -0.20, [
        "warm and nostalgic music", "cozy intimate and tender",
        "a golden vintage glow", "soulful comforting and heartfelt",
        "mellow and affectionate", "musique chaleureuse et tendre"]),
    "serene": (0.50, -0.80, [
        "serene and peaceful music", "calm tranquil and still",
        "a meditative gentle float", "blissful quiet stillness",
        "soothing restful calm", "musique paisible et sereine"]),
    "dreamy": (0.40, -0.45, [
        "dreamy floating music", "ethereal washed-out and hazy",
        "a weightless ambient reverie", "soft shimmering and otherworldly",
        "hypnagogic and faraway", "musique onirique et planante"]),
    "hypnotic": (0.00, 0.35, [
        "hypnotic driving music", "a trance-inducing repetitive groove",
        "mesmerizing and entrancing", "a relentless motorik pulse",
        "looping meditative momentum", "musique hypnotique et répétitive"]),
    "cold": (-0.20, -0.30, [
        "cold and clinical music", "detached robotic and sterile",
        "an icy futuristic chrome surface", "alien emotionless precision",
        "glacial minimal austerity", "musique froide et clinique"]),
    "tense": (-0.30, 0.60, [
        "tense and suspenseful music", "anxious nervous and uneasy",
        "ominous building pressure before a drop", "edgy paranoid dread",
        "a ticking nerve-wracking countdown", "musique tendue et angoissante"]),
    "aggressive": (-0.40, 0.90, [
        "aggressive intense music", "violent pounding and relentless",
        "a furious distorted assault", "harsh militant and brutal",
        "raw adrenaline-fueled rage", "musique agressive et brutale"]),
    "melancholic": (-0.65, -0.40, [
        "melancholic and sad music", "wistful sorrowful and tender",
        "a lonely rainy-day lament", "bittersweet longing and regret",
        "quiet heartbreak and grief", "musique mélancolique et nostalgique"]),
    "dark": (-0.75, 0.20, [
        "dark brooding music", "a sinister menacing atmosphere",
        "shadowy ténébreux and ominous", "bleak dystopian and foreboding",
        "a descent into industrial gloom", "musique sombre et inquiétante"]),
}
NAMES = list(EMOTIONS)

_E = {}


def _matrix():
    """Cache: (names, prompt_embeds (P×512 np), group (P,) emotion idx, VA (E×2 np)).
    All cluster prompts embedded ONCE; group maps each prompt row to its emotion."""
    if "ready" not in _E:
        prompts, group, va = [], [], []
        for ei, name in enumerate(NAMES):
            v, a, cluster = EMOTIONS[name]
            va.append((v, a))
            for p in cluster:
                prompts.append(p)
                group.append(ei)
        te = SEM._embed_texts(prompts).cpu().numpy().astype("float32")  # (P×512), normalized
        _E.update(ready=True, names=NAMES, te=te,
                  group=np.array(group), va=np.array(va, dtype="float32"),
                  counts=np.bincount(group, minlength=len(NAMES)))
    return _E


def score(ae_row, tau=0.05):
    """One normalized CLAP audio embed (512,) → emotion read:
      {scores}  mean cosine per emotion (size-invariant cluster marginalization)
      {dist}    softmax(scores/tau) across emotions — a probability distribution
      top       [(emotion, prob)…] sorted
      valence, arousal  continuous circumplex point = dist-weighted mean of emotion V/A
      confidence  top prob minus second (how decisive the read is)
    """
    E = _matrix()
    ae = np.asarray(ae_row, dtype="float32")
    cos = E["te"] @ ae                                   # (P,) cosine per prompt
    # marginalize cluster → emotion by MEAN (size-invariant; a 6-angle emotion isn't
    # advantaged over a 5-angle one, unlike summing prompt-softmax mass).
    sums = np.zeros(len(E["names"]), dtype="float32")
    np.add.at(sums, E["group"], cos)
    mean_cos = sums / np.maximum(E["counts"], 1)
    z = mean_cos / max(tau, 1e-6)
    z -= z.max()
    dist = np.exp(z); dist /= dist.sum()
    order = np.argsort(-dist)
    val = float(dist @ E["va"][:, 0])
    aro = float(dist @ E["va"][:, 1])
    top = [(E["names"][i], round(float(dist[i]), 3)) for i in order]
    conf = float(dist[order[0]] - (dist[order[1]] if len(order) > 1 else 0.0))
    return {
        "scores": {E["names"][i]: round(float(mean_cos[i]), 4) for i in range(len(E["names"]))},
        "dist": {n: p for n, p in top},
        "top": top,
        "valence": round(val, 3),
        "arousal": round(aro, 3),
        "confidence": round(conf, 3),
    }


def score_folder(name, max_files=None):
    """Mean CLAP embed over a folder's files → emotion read (a folder-level vibe)."""
    files = SEM.CLF.folder_files(name)
    if not files:
        return None
    cap = max_files or SEM.CLF.MAX_FILES
    ae, _ = SEM.embed_audios(files[:cap])
    if not len(ae):
        return None
    mean = ae.mean(0); mean /= (np.linalg.norm(mean) + 1e-9)
    return score(mean)


# ── CLI ─────────────────────────────────────────────────────────────────────────
def _bar(v, lo=-1.0, hi=1.0, w=21):
    """Tiny ASCII gauge for a value in [lo,hi]; marks the midpoint."""
    pos = int(round((v - lo) / (hi - lo) * (w - 1)))
    pos = max(0, min(w - 1, pos))
    cells = ["·"] * w
    cells[w // 2] = "│"
    cells[pos] = "●"
    return "".join(cells)


def cmd_wheel():
    print("⛵ emotion wheel — valence × arousal circumplex, with cluster sizes\n")
    print(f"  {'emotion':<12} {'val':>5} {'aro':>5}  angles")
    for n in sorted(NAMES, key=lambda n: (-EMOTIONS[n][0], -EMOTIONS[n][1])):
        v, a, cl = EMOTIONS[n]
        print(f"  {n:<12} {v:+.2f} {a:+.2f}  {len(cl)}  e.g. “{cl[1]}”")
    print(f"\n  {len(NAMES)} emotions · {sum(len(EMOTIONS[n][2]) for n in NAMES)} prompts total")


def _print_read(label, r):
    print(f"🎭 {label}")
    print(f"   valence {r['valence']:+.2f} {_bar(r['valence'])}  (− unpleasant · pleasant +)")
    print(f"   arousal {r['arousal']:+.2f} {_bar(r['arousal'])}  (− calm · energetic +)")
    top = ", ".join(f"{n} {p:.2f}" for n, p in r["top"][:4])
    print(f"   top: {top}   (conf {r['confidence']:.2f})")


def cmd_tag(target):
    p = Path(target)
    if p.is_file():
        ae, _ = SEM.embed_audios([p])
        if not len(ae):
            sys.exit("decode failed")
        r = score(ae[0])
    else:
        r = score_folder(target)
        if not r:
            sys.exit(f"no folder/files for {target!r}")
    _print_read(target, r)


def cmd_probe():
    """Katana check: archetype TEXT vibes should land where the circumplex says they
    should (no audio needed — validates the ontology geometry + scoring, fast)."""
    print("⛵ probe — archetype text → emotion read (ontology self-consistency)\n")
    archetypes = {
        "a bright joyful major-key piano fanfare": "euphoric/triumphant, +V",
        "a slow sad lonely cello in the rain": "melancholic, −V −A",
        "a brutal distorted industrial techno onslaught": "aggressive/dark, −V +A",
        "a calm ambient meditation drone": "serene/dreamy, +V −A",
        "a relentless hypnotic acid loop": "hypnotic, ~0V +A",
        "an icy detached robotic synth": "cold, −V −A",
    }
    for phrase, expect in archetypes.items():
        emb = SEM._embed_texts([phrase]).cpu().numpy()[0]   # text-as-pseudo-audio probe
        r = score(emb)
        _print_read(f"“{phrase}”  (expect {expect})", r)
        print()


def main():
    args = sys.argv[1:]
    if args and args[0] == "wheel":
        cmd_wheel()
    elif args and args[0] == "tag" and len(args) > 1:
        cmd_tag(args[1])
    elif args and args[0] == "probe":
        cmd_probe()
    else:
        sys.exit("usage: emotion_ontology.py [wheel | tag <folder|file> | probe]")


if __name__ == "__main__":
    main()
