UmbraAudio/umbraaudio/analysis/ambiguity.py

from __future__ import annotations

import logging
import re
from dataclasses import dataclass
from typing import Any

logger = logging.getLogger(__name__)

HIGH_CONFIDENCE = 0.9
MEDIUM_CONFIDENCE = 0.6

SPEECH_VERBS = {
    "said",
    "asked",
    "replied",
    "whispered",
    "shouted",
    "muttered",
    "called",
    "exclaimed",
    "answered",
    "continued",
    "added",
    "noted",
    "stated",
    "remarked",
    "yelled",
    "screamed",
    "laughed",
    "sighed",
    "says",
    "tells",
    "told",
    "orders",
    "ordered",
    "demands",
    "demanded",
    "insists",
    "insisted",
    "begs",
    "begged",
    "pleads",
    "pleaded",
    "commands",
    "commanded",
    "snapped",
    "snarled",
    "hissed",
    "growled",
    "barked",
    "roared",
    "croaked",
    "grunted",
    "murmured",
    "purred",
    "intoned",
    "declared",
    "announced",
    "proclaimed",
    "responds",
    "repeats",
    "mentions",
    "offers",
}

FIRST_SPEAKER_PATTERNS = [
    re.compile(r"for the first time", re.IGNORECASE),
    re.compile(r"spoke (up|first)", re.IGNORECASE),
    re.compile(r"broke the silence", re.IGNORECASE),
    re.compile(r"said (quietly|softly|finally|at last)", re.IGNORECASE),
    re.compile(r"after a (long |moment's |brief )?pause", re.IGNORECASE),
    re.compile(r"stepped forward", re.IGNORECASE),
    re.compile(r"cleared (his|her|their) throat", re.IGNORECASE),
]


@dataclass
class AmbiguityFlag:
    utterance_index: int
    speaker_name: str
    reason: str
    confidence: float
    context_before: str
    context_after: str
    suggested_action: str = "review"

    def to_dict(self) -> dict[str, Any]:
        return {
            "utterance_index": self.utterance_index,
            "speaker_name": self.speaker_name,
            "reason": self.reason,
            "confidence": self.confidence,
            "context_before": self.context_before,
            "context_after": self.context_after,
            "suggested_action": self.suggested_action,
        }


def detect_ambiguities(
    utterances: list[dict[str, Any]],
    known_characters: list[str],
    chapter_title: str = "",
) -> list[AmbiguityFlag]:
    flags: list[AmbiguityFlag] = []

    for i, utt in enumerate(utterances):
        if utt.get("type") != "dialogue":
            continue

        text = utt.get("text", "")
        speaker = utt.get("character_name") or utt.get("speaker_name", "")
        context_before = _get_context(utterances, i, -1)
        context_after = _get_context(utterances, i, 1)

        confidence, reasons = _score_attribution_confidence(
            text, speaker, context_before, context_after, known_characters, i, utterances
        )

        if confidence < HIGH_CONFIDENCE:
            reason_str = "; ".join(reasons) if reasons else "low attribution confidence"
            flags.append(
                AmbiguityFlag(
                    utterance_index=utt.get("index", i),
                    speaker_name=speaker,
                    reason=reason_str,
                    confidence=confidence,
                    context_before=context_before,
                    context_after=context_after,
                    suggested_action="review" if confidence < MEDIUM_CONFIDENCE else "flag",
                )
            )

    return flags


def _score_attribution_confidence(
    text: str,
    speaker: str,
    context_before: str,
    context_after: str,
    known_characters: list[str],
    utterance_index: int,
    all_utterances: list[dict[str, Any]],
) -> tuple[float, list[str]]:
    reasons: list[str] = []

    if speaker == "narrator" or not speaker:
        return 0.0, ["unattributed dialogue"]

    text_lower = text.lower()
    context_lower = (context_before + " " + context_after).lower()
    speaker_lower = speaker.lower()

    has_direct_tag = any(
        re.search(rf"\b{re.escape(speaker)}\b.*?\b{verb}\b", context_before, re.IGNORECASE)
        or re.search(rf"\b{verb}\b.*?\b{re.escape(speaker)}\b", context_after, re.IGNORECASE)
        for verb in SPEECH_VERBS
    )

    if has_direct_tag:
        return 0.95, ["dialogue tag present"]

    has_indirect_tag = any(
        re.search(rf"\b{re.escape(name)}\b", context_before, re.IGNORECASE)
        for name in known_characters
    )
    has_verb = any(verb in context_lower for verb in SPEECH_VERBS)

    score = 0.5

    if has_indirect_tag and has_verb:
        score = 0.8
    elif has_indirect_tag:
        score = 0.7
        reasons.append("character named nearby but no speech verb")
    elif has_verb:
        score = 0.6
        reasons.append("speech verb present but speaker name not found")

    pronoun_match = re.match(r"^(i|we)\b", text_lower)
    if pronoun_match and speaker_lower not in context_lower:
        score -= 0.2
        reasons.append("first-person dialogue without clear speaker context")

    for pattern in FIRST_SPEAKER_PATTERNS:
        if pattern.search(context_before):
            score = min(score, 0.5)
            reasons.append("possible first-speaker introduction - verify")
            break

    scene_participants = _detect_scene_participants(all_utterances, utterance_index)
    if scene_participants and speaker not in scene_participants:
        score -= 0.15
        reasons.append(
            f"speaker {speaker} not in detected scene participants: {scene_participants}"
        )

    consecutive_by_same = 0
    for j in range(utterance_index - 1, max(utterance_index - 4, -1), -1):
        if j >= 0 and all_utterances[j].get("character_name") == speaker:
            consecutive_by_same += 1
        else:
            break
    if consecutive_by_same >= 3:
        score -= 0.1
        reasons.append(f"high consecutive dialogue count ({consecutive_by_same}) for same speaker")

    return max(0.0, min(1.0, score)), reasons


def _get_context(utterances: list[dict[str, Any]], index: int, direction: int) -> str:
    target = index + direction
    if 0 <= target < len(utterances):
        return utterances[target].get("text", "")
    return ""


def _detect_scene_participants(utterances: list[dict[str, Any]], current_index: int) -> set[str]:
    participants: set[str] = set()
    for j in range(max(0, current_index - 10), min(len(utterances), current_index + 10)):
        utt = utterances[j]
        if utt.get("type") == "scene_break":
            participants.clear()
        name = utt.get("character_name") or utt.get("speaker_name", "")
        if name and name != "narrator":
            participants.add(name)
    return participants