#!/usr/bin/env python3
"""Decoder temporal de fases de cerimônia (Viterbi) — v0.2 beat labeling.

Puro/testável: sem VLM, sem I/O. Prior canônico BR (mesmo do storyboard_v1.md),
probabilidades FIXAS no código (non-overfit: não tunar por casamento).

Fases ordenadas:
  pre → entrances → liturgia → rings → declaration → kiss → applause → exit
empty_wait não muda fase. Sub-beats permutam livre dentro da fase.
"""
from __future__ import annotations

from collections import Counter
from typing import Dict, List, Optional, Sequence, Tuple

# --- vocab / fases -----------------------------------------------------------

BEAT_VOCAB = frozenset({
    "pre", "empty_wait",
    "entrance_mother", "entrance_groom", "entrance_bridesmaids",
    "flower_girl", "ring_bearer", "entrance_bride",
    "officiant", "reading", "vows",
    "rings", "declaration", "kiss", "applause", "exit",
})

# fase canônica por sub-beat
BEAT_TO_PHASE: Dict[str, str] = {
    "pre": "pre",
    "empty_wait": "empty_wait",  # especial: não muda fase
    "entrance_mother": "entrances",
    "entrance_groom": "entrances",
    "entrance_bridesmaids": "entrances",
    "flower_girl": "entrances",
    "ring_bearer": "entrances",
    "entrance_bride": "entrances",
    "officiant": "liturgia",
    "reading": "liturgia",
    "vows": "liturgia",
    "rings": "rings",
    "declaration": "declaration",
    "kiss": "kiss",
    "applause": "applause",
    "exit": "exit",
}

PHASES: List[str] = [
    "pre", "entrances", "liturgia", "rings", "declaration", "kiss", "applause", "exit",
]
PHASE_INDEX = {p: i for i, p in enumerate(PHASES)}

PHASE_BEATS: Dict[str, List[str]] = {
    "pre": ["pre"],
    "entrances": [
        "entrance_mother", "entrance_groom", "entrance_bridesmaids",
        "flower_girl", "ring_bearer", "entrance_bride",
    ],
    "liturgia": ["officiant", "reading", "vows"],
    "rings": ["rings"],
    "declaration": ["declaration"],
    "kiss": ["kiss"],
    "applause": ["applause"],
    "exit": ["exit"],
}

# importance/emotion da FASE (tabela do prompt) — quebra circularidade VLM
PHASE_IMPORTANCE: Dict[str, float] = {
    "pre": 0.2,
    "entrances": 0.7,
    "liturgia": 0.8,  # vows 0.8; officiant 0.5 → usa 0.8 se votos, 0.5 se só officiant (ver abaixo)
    "rings": 0.9,
    "declaration": 0.9,
    "kiss": 0.9,
    "applause": 0.7,
    "exit": 0.75,
}

BEAT_IMPORTANCE: Dict[str, float] = {
    "pre": 0.2, "empty_wait": 0.2,
    "entrance_mother": 0.7, "entrance_groom": 0.7, "entrance_bridesmaids": 0.7,
    "flower_girl": 0.7, "ring_bearer": 0.7, "entrance_bride": 0.9,
    "officiant": 0.5, "reading": 0.5, "vows": 0.8,
    "rings": 0.9, "declaration": 0.9, "kiss": 0.9,
    "applause": 0.7, "exit": 0.75,
}

# Transições monotônicas (fixas no código)
P_STAY = 0.85
P_NEXT = 0.12
P_SKIP = 0.03
# rings→liturgia OFF no v0.2 (plano)
ALLOW_RINGS_TO_LITURGIA = False

# confusões conhecidas: crédito parcial de emissão
CONFUSION_PAIRS = {
    frozenset({"entrance_bride", "vows"}),
    frozenset({"entrance_bride", "exit"}),
    frozenset({"officiant", "vows"}),
    frozenset({"entrance_groom", "vows"}),
    frozenset({"applause", "exit"}),
}


def clamp_beat(beat: Optional[str]) -> str:
    if not beat or beat not in BEAT_VOCAB:
        return "officiant"
    return beat


def transition_logp(prev_phase: str, next_phase: str) -> float:
    """Log-prob de transição (usa log natural; valores relativos bastam).

    Avanço monotônico: ficar / +1 / +2 / +3 (pulo maior com prob menor).
    Não volta (exceto flag rings→liturgia, off no v0.2).
    """
    import math
    if prev_phase not in PHASE_INDEX or next_phase not in PHASE_INDEX:
        return -1e9
    i, j = PHASE_INDEX[prev_phase], PHASE_INDEX[next_phase]
    if j == i:
        return math.log(P_STAY)
    if j == i + 1:
        return math.log(P_NEXT)
    if j == i + 2:
        return math.log(P_SKIP)
    if j == i + 3:
        return math.log(0.01)
    if j == i + 4:
        return math.log(0.003)
    if ALLOW_RINGS_TO_LITURGIA and prev_phase == "rings" and next_phase == "liturgia":
        return math.log(0.05)
    # voltar ou pular demais
    return -1e9


def emission_score(votes: Sequence[str], phase: str) -> float:
    """Fração dos votos compatíveis com a fase (+ crédito parcial por confusão)."""
    if not votes:
        return 0.0
    phase_beats = set(PHASE_BEATS.get(phase, []))
    score = 0.0
    for v in votes:
        v = clamp_beat(v)
        if v == "empty_wait":
            score += 0.5  # neutro
            continue
        if BEAT_TO_PHASE.get(v) == phase or v in phase_beats:
            score += 1.0
            continue
        # confusão parcial: se algum beat da fase confunde com v
        partial = False
        for pb in phase_beats:
            if frozenset({pb, v}) in CONFUSION_PAIRS:
                score += 0.35
                partial = True
                break
        if not partial:
            score += 0.0
    return score / max(1, len(votes))


def viterbi_phases(window_votes: Sequence[Sequence[str]]) -> List[str]:
    """Decodifica sequência de fases para N janelas. Cada item = lista de 1–3 beats crus."""
    import math
    n = len(window_votes)
    if n == 0:
        return []
    phases = PHASES
    # DP: dp[t][p] = best logp ending in phase p at time t
    NEG = -1e18
    dp = [[NEG] * len(phases) for _ in range(n)]
    back = [[-1] * len(phases) for _ in range(n)]

    # peso de emissão > 1: confiar mais no voto local p/ avançar fase quando o sinal é claro
    EM_W = 2.5

    for pi, p in enumerate(phases):
        em = emission_score(window_votes[0], p)
        # prior: começar em pre / entrances
        start = math.log(0.7) if p == "pre" else (math.log(0.25) if p == "entrances" else math.log(0.01))
        dp[0][pi] = start + EM_W * math.log(max(em, 1e-6))

    for t in range(1, n):
        for pj, pnew in enumerate(phases):
            em = emission_score(window_votes[t], pnew)
            em_l = EM_W * math.log(max(em, 1e-6))
            best, arg = NEG, 0
            for pi, pold in enumerate(phases):
                sc = dp[t - 1][pi] + transition_logp(pold, pnew) + em_l
                if sc > best:
                    best, arg = sc, pi
            dp[t][pj] = best
            back[t][pj] = arg

    # traceback
    last = max(range(len(phases)), key=lambda i: dp[n - 1][i])
    path = [0] * n
    path[n - 1] = last
    for t in range(n - 1, 0, -1):
        path[t - 1] = back[t][path[t]]
    return [phases[i] for i in path]


def pick_subbeat(votes: Sequence[str], phase: str, prev_subbeat: Optional[str] = None) -> str:
    """Sub-beat = modo restrito à fase; senão herda vizinho."""
    allowed = PHASE_BEATS.get(phase, ["officiant"])
    filtered = [clamp_beat(v) for v in votes if clamp_beat(v) in allowed]
    if filtered:
        return Counter(filtered).most_common(1)[0][0]
    # empty_wait votes only
    if any(clamp_beat(v) == "empty_wait" for v in votes):
        return "empty_wait"
    if prev_subbeat and BEAT_TO_PHASE.get(prev_subbeat) == phase:
        return prev_subbeat
    # default representative of phase
    return allowed[0] if phase != "liturgia" else "vows"


def importance_for(beat: str, phase: str) -> float:
    return BEAT_IMPORTANCE.get(beat, PHASE_IMPORTANCE.get(phase, 0.5))


def emotion_for(votes: Sequence[str], emotions: Sequence[str], phase: str) -> str:
    """Hierarquia low < med < high sobre frames cuja vote cabe na fase."""
    order = {"low": 0, "med": 1, "high": 2}
    allowed = set(PHASE_BEATS.get(phase, []))
    best = 0
    for v, e in zip(votes, emotions):
        v = clamp_beat(v)
        if v in allowed or v == "empty_wait" or BEAT_TO_PHASE.get(v) == phase:
            best = max(best, order.get((e or "low").lower(), 0))
    for k, val in order.items():
        if val == best:
            return k
    return "low"


def decode_windows(
    window_votes: Sequence[Sequence[str]],
    window_emotions: Optional[Sequence[Sequence[str]]] = None,
) -> List[dict]:
    """Retorna lista de {phase, beat, importance, emotion} por janela."""
    phases = viterbi_phases(window_votes)
    out = []
    prev_beat = None
    for i, phase in enumerate(phases):
        votes = list(window_votes[i]) if i < len(window_votes) else []
        emos = list(window_emotions[i]) if window_emotions and i < len(window_emotions) else ["low"] * len(votes)
        beat = pick_subbeat(votes, phase, prev_beat)
        # empty_wait: mantém fase do vizinho no campo phase, mas emite empty_wait
        if all(clamp_beat(v) == "empty_wait" for v in votes) and votes:
            beat = "empty_wait"
        imp = importance_for(beat, phase)
        emo = emotion_for(votes, emos, phase)
        out.append({"phase": phase, "beat": beat, "importance": imp, "emotion": emo, "votes": votes})
        if beat != "empty_wait":
            prev_beat = beat
    return out


def merge_same_beat(rows: List[dict]) -> List[dict]:
    """Merge adjacente com mesmo sub-beat (como storyboard_gen atual)."""
    if not rows:
        return []
    out = [dict(rows[0])]
    for r in rows[1:]:
        if out[-1]["beat"] == r["beat"] and r["f0"] <= out[-1]["f1"]:
            out[-1]["f1"] = r["f1"]
            out[-1]["importance"] = max(out[-1]["importance"], r["importance"])
            if r.get("emotion") == "high":
                out[-1]["emotion"] = "high"
        else:
            out.append(dict(r))
    return out


# --- eval estrutural (sem rede) ------------------------------------------------

def eval_sequence(segments: List[dict], cov0: int, cov1: int) -> dict:
    """Métricas do plano v0.2 §3.3. segments: {f0,f1,beat}."""
    n = len(segments)
    beats = [s["beat"] for s in segments]
    # 1º rings/declaration
    first_ritual = None
    for s in segments:
        if s["beat"] in ("rings", "declaration", "kiss"):
            first_ritual = s["f0"]
            break
    entrance_after = 0
    if first_ritual is not None:
        entrance_after = sum(
            1 for s in segments
            if s["f0"] > first_ritual and str(s["beat"]).startswith("entrance")
        )
    # exit antes do último decil
    last_decile = cov0 + int(0.9 * (cov1 - cov0))
    exit_early = sum(1 for s in segments if s["beat"] == "exit" and s["f1"] < last_decile)
    # declaration/kiss antes do 1º entrance
    first_entrance = None
    for s in segments:
        if str(s["beat"]).startswith("entrance") or s["beat"] in ("flower_girl", "ring_bearer"):
            first_entrance = s["f0"]
            break
    decl_before = 0
    if first_entrance is not None:
        decl_before = sum(
            1 for s in segments
            if s["beat"] in ("declaration", "kiss") and s["f0"] < first_entrance
        )
    bride_n = sum(1 for b in beats if b == "entrance_bride")
    return {
        "n_segments": n,
        "entrance_after_ritual": entrance_after,
        "exit_before_last_decile": exit_early,
        "decl_kiss_before_entrance": decl_before,
        "entrance_bride_count": bride_n,
        "pass_n_segments": 15 <= n <= 35,
        "pass_entrance_after": entrance_after <= 3,
        "pass_exit_early": exit_early == 0,
        "pass_decl_before": decl_before == 0,
        "pass_bride_count": 1 <= bride_n <= 3,
    }
