"""Track scoring: lyrics in, match score out.

This is the whole of /api/score after lyrics have been fetched. It lives here
rather than in api.py so the deployed endpoint and the local benchmark run the
same code — a candidate checkpoint can then be compared against the live model
on identical logic before anything ships, and the helpers can be tested without
importing FastAPI.
"""

from __future__ import annotations

import re
from typing import Dict, Optional


# English function words, excluding forms shared with Spanish, Portuguese,
# Italian or romanized Hindi (a, me, no, te, se, ho, main...), so a line in
# those languages scores near zero rather than borrowing English matches.
_EN_WORDS = frozenset("""the and you i'm to my it in of is your that on for be we all this with but
what can don't just like know when so they are was there if at up one now it's i've i'll you're
never from how will she he her his out them not can't oh our do have""".split())


def lyrics_language_ok(text: str) -> bool:
    """False when lyrics are clearly not English.

    The classifier was trained on English only; on other languages its output
    is noise (Despacito and Sigur Ros both read as "sarcasm"), which the fusion
    would then treat as evidence. English lyrics put ~50% of their words in the
    list above; Spanish and romanized Hindi put 0%. Non-English: fewer than half
    the letters are Latin (Devanagari, CJK...), or at least 20 words with under
    10% English function words. Short texts are given the benefit of the doubt.
    """
    letters = re.findall(r"[^\W\d_]", text)
    if letters and sum(1 for c in letters if c.isascii()) / len(letters) < 0.5:
        return False
    words = re.findall(r"[a-z']+", text.lower())
    if len(words) < 20:
        return True
    return sum(w in _EN_WORDS for w in words) / len(words) >= 0.10


def lyric_chunks(lyrics: str) -> list:
    """The stanzas production scores: section markers stripped, split on blank
    lines (or into 5-line groups when there are none), repeated choruses
    dropped, each capped at 350 characters. Shared with lyric fine-tuning so the
    model trains on exactly what it is later asked to read."""
    # Strip section markers like [Verse 1], [Chorus], [Bridge]
    cleaned = re.sub(r'^\s*\[.*?\]\s*$', '', lyrics, flags=re.MULTILINE)

    # Split on blank lines first (stanza boundaries), then fallback to 5-line groups
    raw_blocks = [b.strip() for b in re.split(r'\n{2,}', cleaned) if b.strip()]
    if len(raw_blocks) <= 1:
        lines = [l.strip() for l in cleaned.split('\n') if l.strip()]
        raw_blocks = ['\n'.join(lines[i:i + 5]) for i in range(0, len(lines), 5)]

    # Deduplicate: skip blocks whose first 50 chars already appeared (repeated choruses)
    seen: set = set()
    chunks: list = []
    for block in raw_blocks:
        fp = block[:50].lower().strip()
        if fp not in seen and len(block) > 15:
            seen.add(fp)
            chunks.append(block[:350])   # DeBERTa max_length=256 tokens; 350 chars ≈ 100 tokens
    return chunks


def chunked_lyric_probs(predictor, lyrics: str) -> dict:
    """Score song lyrics by stanza — far more accurate than single-pass.

    Song lyrics use figurative language that DeBERTa (trained on social media
    text) reads as 'neutral' when presented as one long block. Splitting into
    stanzas and averaging lets emotional peaks in each verse pull the score up.

    Returns: weighted-average probability dict across up to 5 unique stanzas.
    """
    chunks = lyric_chunks(lyrics)

    if not chunks:
        result = predictor.predict(lyrics[:400], create_chart=False)
        return result.get('probabilities', {})

    # Score up to 5 chunks — balances signal vs latency. One batched pass when
    # the predictor supports it (2x faster on CPU); per chunk otherwise, where
    # a failing chunk is skipped rather than failing the track.
    chunk_results: list = []
    batched = None
    if hasattr(predictor, "predict_probs_many"):
        try:
            batched = predictor.predict_probs_many(chunks[:5])
        except Exception:
            batched = None
    if batched is not None:
        chunk_results = [(r.get('probabilities', {}), r.get('confidence', 0.5)) for r in batched]
    else:
        for chunk in chunks[:5]:
            try:
                r = predictor.predict(chunk, create_chart=False)
                chunk_results.append((r.get('probabilities', {}), r.get('confidence', 0.5)))
            except Exception:
                pass

    if not chunk_results:
        result = predictor.predict(lyrics[:400], create_chart=False)
        return result.get('probabilities', {})

    # Confidence-weighted average: high-certainty stanzas count more
    all_keys: set = set()
    for probs, _ in chunk_results:
        all_keys.update(probs.keys())

    avg: dict = {k: 0.0 for k in all_keys}
    total_w = 0.0
    for probs, conf in chunk_results:
        for k in all_keys:
            avg[k] += probs.get(k, 0.0) * conf
        total_w += conf

    if total_w > 0:
        for k in avg:
            avg[k] /= total_w

    return avg


def cosine(u: dict, t: dict) -> float:
    keys = list(set(u) | set(t))
    uv = [float(u.get(k, 0.0)) for k in keys]
    tv = [float(t.get(k, 0.0)) for k in keys]
    dot = sum(a * b for a, b in zip(uv, tv))
    mag_u = sum(a ** 2 for a in uv) ** 0.5
    mag_t = sum(b ** 2 for b in tv) ** 0.5
    return dot / (mag_u * mag_t) if (mag_u > 0 and mag_t > 0) else 0.0


def emotion_match(user: dict, track: dict, neutral_key: str = "neutral") -> float:
    """Similarity between a user's emotion vector and a track's, 0-1.

    Plain cosine over the full 8-class distribution has a pathology that matters
    a great deal here: it treats `neutral` as just another emotion, maximally
    distant from `sadness`. But `neutral` is not an emotion, it is the model
    declining to name one. Measured on a sadness-dominant user vector:

        track 90% neutral   -> 8
        track 90% joy       -> 8      the opposite emotion, same score
        uniform (no signal) -> 50     more than either

    Since the classifier retreats to `neutral` on most figurative lyrics, every
    such song was scoring like a genuinely mismatched one — "Holocene" ranked
    alongside a happy song rather than above it.

    Instead, drop `neutral` from both vectors and compare what remains, which is
    the actual emotional content. A track read as mostly neutral still carries a
    faint direction in its residual mass, and that direction is informative:
    neutral-leaning-sad should outrank neutral-leaning-joyful.

    That residual is also noisy, so the result is scaled by how much
    non-neutral mass the track has. A 90%-neutral track cannot score like a
    confident match, but it still sorts above one leaning the wrong way:

        confident sad 95  >  neutral-but-sad 55  >  flat 28  >  joy-leaning 8

    This is independent of the model — it corrects how predictions are compared,
    not how they are produced.
    """
    u = {k: float(v) for k, v in user.items() if k != neutral_key}
    t = {k: float(v) for k, v in track.items() if k != neutral_key}

    u_mass, t_mass = sum(u.values()), sum(t.values())
    if u_mass <= 0 or t_mass <= 0:
        # One side is entirely neutral: there is nothing to compare.
        return 0.0

    # Confidence is the *fraction* of the track's mass that is non-neutral, not
    # the raw sum. Using the raw sum assumes the input is already normalized;
    # an unnormalized distribution then inflates the factor without bound and
    # produces scores above 100%.
    track_total = sum(float(v) for v in track.values())
    if track_total <= 0:
        return 0.0
    emotional_fraction = min(1.0, t_mass / track_total)

    similarity = cosine(
        {k: v / u_mass for k, v in u.items()},
        {k: v / t_mass for k, v in t.items()},
    )
    # The floor keeps ordering among neutral-heavy tracks intact rather than
    # collapsing it to zero.
    return similarity * (0.5 + 0.5 * emotional_fraction)


def score_probs(track_probs: Dict[str, float], user_probabilities: Dict[str, float]) -> dict:
    """The /api/score lyric fields for an already-computed lyric distribution.

    Split from score_lyrics so the distribution, which is the expensive part
    (up to five DeBERTa passes), can be cached per track and matched against
    any listener for free.
    """
    dominant = max(track_probs, key=track_probs.get) if track_probs else "neutral"
    match = emotion_match(user_probabilities, track_probs)

    # Report the strongest *emotion* too. When the model abstains, the dominant
    # class is `neutral`, which tells the UI nothing about the song.
    emotional = {k: v for k, v in track_probs.items() if k != "neutral"}
    lean = max(emotional, key=emotional.get) if emotional else dominant

    return {
        "match_score": round(match * 100),
        "track_emotion": dominant,
        "track_lean": lean,
        "emotional_mass": round(sum(emotional.values()), 4) if emotional else 0.0,
    }


def score_lyrics(predictor, lyrics: str, user_probabilities: Dict[str, float]) -> dict:
    """The /api/score response for already-fetched lyrics."""
    return score_probs(chunked_lyric_probs(predictor, lyrics), user_probabilities)


# Lyrics' share of the fused score. Fixed at 50/50 before any evaluation and
# not tuned afterwards; see FINDINGS.md for the measurements behind it.
FUSION_WEIGHT = 0.5


def combine(lyric_result: Optional[dict], audio: Optional[dict],
            audio_match: Optional[float], weight: float = FUSION_WEIGHT) -> Optional[dict]:
    """Merge the lyric and acoustic readings of a track into one response.

    Either side may be missing: instrumental tracks and songs no provider has
    lyrics for are scored by sound alone, and tracks without a preview keep the
    lyric-only score they had before audio existed. `basis` says which applied,
    so a score is never presented as more informed than it is.

    lyric_result: score_lyrics() output, or None.
    audio: the acoustic analysis to expose to the client, or None.
    audio_match: 0-1 agreement from AcousticHead.match, or None.
    """
    lyric = (lyric_result or {}).get("match_score")
    if lyric is None and audio_match is None:
        return None

    out = dict(lyric_result or {})
    out["lyric_match"] = lyric
    out["audio_match"] = None if audio_match is None else round(audio_match * 100)
    if lyric is not None and audio_match is not None:
        out["match_score"] = round(weight * lyric + (1 - weight) * audio_match * 100)
        out["basis"] = "lyrics+audio"
    elif lyric is not None:
        out["basis"] = "lyrics"
    else:
        out["match_score"] = out["audio_match"]
        out["basis"] = "audio"
    out["audio"] = audio
    return out


def fused_response(fusion, user_probabilities: Dict[str, float],
                   track_probs: Optional[Dict[str, float]], audio: Optional[dict],
                   lyrics_english: bool = True) -> Optional[dict]:
    """The /api/score response from the four-modality VAD fusion (Objective 8).

    user sentiment (user_probabilities), lyrical emotion (track_probs), the
    acoustic embedding and structured audio features (audio["vad_embedding"],
    audio["vad_structured"]) become one fused VAD point for the track and a
    calibrated probability that it fits the listener's mood. Returns None when
    there is nothing to score, or the listener's text carries no emotion, so the
    caller can fall back to the previous scoring.
    """
    from .fusion import DIMS, lyric_vad, user_vad

    you = user_vad(user_probabilities)
    lv, mass = lyric_vad(track_probs)
    emb = audio.get("vad_embedding") if audio else None
    struct = audio.get("vad_structured") if audio else None
    t = fusion.track_vad(emb, struct, lv, mass)
    if you is None or t is None:
        return None
    cosine = emotion_match(user_probabilities, track_probs) if track_probs else None
    p = fusion.match(user_probabilities, t, cosine)
    if p is None:
        return None

    lyric = score_probs(track_probs, user_probabilities) if track_probs else None
    out = dict(lyric or {})
    heard = emb is not None
    out["lyric_match"] = lyric["match_score"] if lyric else None
    out["audio_match"] = None
    out["match_score"] = round(p * 100)
    out["basis"] = "lyrics+audio" if (track_probs and heard) else ("audio" if heard else "lyrics")
    out["calibrated"] = True

    def point(v):
        return None if v is None else {d: round(float(x), 2) for d, x in zip(DIMS, v)}
    out["vad"] = {"track": point(t), "you": point(you), "lyrics": point(lv),
                  "embedding": point(emb), "structured": point(struct)}
    out["explanation"] = fusion.explain(
        user_probabilities, t, p,
        lyric_lean=(lyric or {}).get("track_lean"), lyric_mass=mass,
        sound=(audio or {}).get("sound"), lyrics_english=lyrics_english)
    if track_probs and not lyrics_english:
        out["lyrics_language"] = "non-English"
    view = None
    if audio:
        view = {k: audio.get(k) for k in ("valence", "arousal", "sound", "descriptors", "preview_url")}
    out["audio"] = view
    return out