"""Streaming keyword spotting over CTC phoneme posteriors.

Keyword-filler decoding with two robustness measures beyond the textbook
version (both port directly to C on the ESP32):

  - minimum phone duration: every phone is expanded to two chained states,
    so an alignment must spend >= 2 frames (40 ms) per phone. Kills
    spurious single-frame matches.
  - duration-normalized scoring: each Viterbi token carries the frame at
    which its path entered the keyword, and the detection statistic is
    (keyword_score - filler_score) / path_duration -- an average
    per-frame deficit, comparable across phrase lengths and speaking
    rates. Alignments faster than 2 frames/phone or slower than
    MAX_FRAMES_PER_PHONE are rejected outright.

Enrollment of a new wake word is just: phones = text_to_phones(phrase).
State chain per phone i: A_i -> B_i (same label), optional blank after.
"""

import numpy as np

from .phones import BLANK, PHONE_TO_ID, text_to_phones

NEG_INF = -1e30
# 200 ms per phone upper bound: fluent speech runs 40-120 ms/phone, and
# looser caps let garbage alignments crawl across continuous speech
# (observed 1.5-2.3 s "matches" on background TV at 400 ms/phone)
MAX_FRAMES_PER_PHONE = 10

# Highly confusable phone pairs: an alignment may match either member.
# Text pronunciations (CMUdict/G2P) use canonical vowels, but real and TTS
# speech often reduces them -- e.g. "orbit" is listed as AO R B AH T yet
# actually said as AO R B IH T. Without this, a sharper acoustic model
# *punishes* the mismatch harder.
CONFUSABLE = {
    "AH": ("AH", "IH", "ER"),   # schwa reduces/r-colors: orbit -> orbERt
    "IH": ("IH", "AH"),
    "ER": ("ER", "AH"),
    "AO": ("AO", "AA"),         # cot-caught merger and accent variation
    "AE": ("AE", "AA"),         # trap-father variation (sakura, tanaka)
    "UH": ("UH", "UW"),         # lax/tense u (book/boot neighbors)
    "K": ("K", "G"),            # stops voice between vowels: nakuma->naguma
    "G": ("G", "K"),
    "T": ("T", "D"),            # also covers tapped/rolled R heard as D
    "D": ("D", "T"),
}


class KeywordSpotter:
    def __init__(self, phrase, threshold=-4.4, refractory_frames=50,
                 strong_margin=1.5, strong_ratio=0.5, phones=None,
                 wildcards=None, wildcard_penalty=2.2):
        """threshold: average per-frame deficit tolerated along the keyword
        path (log-prob units, <= 0; closer to 0 = stricter).
        refractory_frames: minimum output frames (20 ms) between fires.
        strong_margin/strong_ratio: evidence requirement -- at least
        strong_ratio of the path's phone frames must have their phone
        within strong_margin log-prob of the frame's best class. Prevents
        noisy/mushy audio (where nothing is clearly heard) from firing on
        average-score alone.
        phones: explicit phoneme sequence (e.g. from voice enrollment);
        overrides the dictionary pronunciation of `phrase`."""
        if phones is None:
            phones = [p for p in text_to_phones(phrase) if p in PHONE_TO_ID]
        else:
            phones = [p for p in phones if p in PHONE_TO_ID]
        if len(phones) < 2:
            raise ValueError(f"phrase too short to spot: {phrase!r}")
        self.phrase = phrase
        self.phones = phones
        self.threshold = threshold
        self.refractory_frames = refractory_frames
        self.strong_margin = strong_margin
        self.strong_ratio = strong_ratio
        # Mismatch tolerance (default OFF): measured on universal clips
        # vs negatives, wildcarding lifted impostor scores as much as
        # genuine ones and REDUCED recall at zero-FA operating points
        # (vucano negmax -4.14 -> -2.89). Kept only for future
        # per-count-lattice experiments.
        if wildcards is None:
            wildcards = 0
        self.max_wild = wildcards * 2   # budget in frames
        self.wild_pen = wildcard_penalty

        # Build states: leading blank, then per phone A,B (+ trailing blank)
        labels = [BLANK]
        self.preds = [[0]]              # predecessor state ids (self-loop
        entry = [True]                  # implied for every state)
        for i, p in enumerate(phones):
            pid = PHONE_TO_ID[p]
            a = len(labels)
            if i == 0:
                labels.append(pid); self.preds.append([0]); entry.append(True)
            else:
                pre = [a - 1]           # blank between phones
                if labels[a - 2] != pid:
                    pre.append(a - 2)   # skip blank (different phones only)
                labels.append(pid); self.preds.append(pre); entry.append(False)
            labels.append(pid)          # B state: only from A
            self.preds.append([a]); entry.append(False)
            labels.append(BLANK)        # blank after phone
            self.preds.append([a + 1]); entry.append(False)
        self.labels = labels
        # allowed emission ids per state (confusable vowels match either)
        self.allowed = []
        id_to_phone = {v: k for k, v in PHONE_TO_ID.items()}
        for lab in labels:
            if lab == BLANK:
                self.allowed.append((BLANK,))
            else:
                ph = id_to_phone[lab]
                self.allowed.append(tuple(
                    PHONE_TO_ID[p] for p in CONFUSABLE.get(ph, (ph,))))
        self.entry = entry
        self.n_states = len(labels)
        self.finals = [self.n_states - 1, self.n_states - 2]
        # min: ~60 ms per phone on average (individual phones may be
        # shorter); max: 400 ms per phone. Anything outside is not a
        # human saying the phrase.
        self.min_dur = 3 * len(phones)
        self.max_dur = MAX_FRAMES_PER_PHONE * len(phones)
        self.reset()

    def reset(self):
        self.rel = np.full(self.n_states, NEG_INF)
        self.start = np.zeros(self.n_states, dtype=np.int64)
        self.pframes = np.zeros(self.n_states, dtype=np.int64)
        self.strong = np.zeros(self.n_states, dtype=np.int64)
        self.wild = np.zeros(self.n_states, dtype=np.int64)
        self.cooldown = 0
        self.t = 0

    def _update(self, log_probs_frame):
        """One Viterbi DP step. Returns the best duration-valid normalized
        score at a final state this frame (or None). No side effects on
        detection state."""
        lp = log_probs_frame
        filler = float(lp.max())
        prev_rel, prev_start = self.rel, self.start
        prev_pf, prev_strong = self.pframes, self.strong
        prev_wild = self.wild
        cur_rel = np.full(self.n_states, NEG_INF)
        cur_start = np.zeros(self.n_states, dtype=np.int64)
        cur_pf = np.zeros(self.n_states, dtype=np.int64)
        cur_strong = np.zeros(self.n_states, dtype=np.int64)
        cur_wild = np.zeros(self.n_states, dtype=np.int64)
        for s in range(self.n_states):
            best = prev_rel[s]
            best_start, best_pf, best_sf, best_w = (
                prev_start[s], prev_pf[s], prev_strong[s], prev_wild[s])
            for q in self.preds[s]:
                if prev_rel[q] > best:
                    best = prev_rel[q]
                    best_start, best_pf, best_sf, best_w = (
                        prev_start[q], prev_pf[q], prev_strong[q],
                        prev_wild[q])
            if self.entry[s] and 0.0 >= best:
                # (re)start the keyword here; >= keeps the start time fresh
                # while idling in silence, so duration stays meaningful
                best, best_start, best_pf, best_sf, best_w = \
                    0.0, self.t, 0, 0, 0
            emit = max(float(lp[i]) for i in self.allowed[s])
            is_phone = self.labels[s] != BLANK
            used_wild = 0
            if is_phone and best_w < self.max_wild:
                # mismatch tolerance: accept the frame's best class at a
                # penalty when the expected phone isn't there (bounded
                # budget converts exact matching into similarity matching)
                soft = filler - self.wild_pen
                if soft > emit:
                    emit = soft
                    used_wild = 1
            cur_rel[s] = best + emit - filler
            cur_start[s] = best_start
            cur_pf[s] = best_pf + (1 if is_phone else 0)
            cur_strong[s] = best_sf + (
                1 if is_phone and emit >= filler - self.strong_margin else 0)
            cur_wild[s] = best_w + used_wild
        self.rel, self.start = cur_rel, cur_start
        self.pframes, self.strong = cur_pf, cur_strong
        self.wild = cur_wild
        self.t += 1
        norm_best = None
        for s in self.finals:
            dur = self.t - cur_start[s]
            pf = cur_pf[s]
            if not (self.min_dur <= dur <= self.max_dur):
                continue
            if pf < 2 * len(self.phones):
                continue
            # evidence requirement: the phrase's phones must actually have
            # been the near-top hypothesis for enough of the match, not
            # merely "not too costly on average" (mushy noisy audio)
            if cur_strong[s] < self.strong_ratio * pf:
                continue
            # normalize by phone frames only: time spent in blank states is
            # free in silence, so counting it would let blank-padded paths
            # dilute their deficit below any threshold
            norm = cur_rel[s] / pf
            if norm_best is None or norm > norm_best[0]:
                norm_best = (float(norm), int(dur))
        return norm_best

    def step(self, log_probs_frame):
        """Consume one frame of log-probs (C,). Returns (score, duration)
        if the keyword fired this frame, else None. duration is in output
        frames (20 ms each), for energy-gating by the caller."""
        norm = self._update(log_probs_frame)
        if self.cooldown > 0:
            self.cooldown -= 1
            return None
        if norm is not None and norm[0] > self.threshold:
            self.cooldown = self.refractory_frames
            self.rel = np.full(self.n_states, NEG_INF)
            return norm
        return None

    def run(self, log_probs):
        """Offline helper: (T, C) log-probs -> list of (frame, score)."""
        hits = []
        for t in range(log_probs.shape[0]):
            s = self.step(log_probs[t])
            if s is not None:
                hits.append((t, s[0]))
        return hits

    def best_score(self, log_probs):
        """Max normalized final-state score over a clip (pure Viterbi, no
        firing/reset side effects) -- used for calibration."""
        self.reset()
        best = -np.inf
        for t in range(log_probs.shape[0]):
            s = self._update(log_probs[t])
            if s is not None and s[0] > best:
                best = s[0]
        return best