src/sphinx_corrector.py

#!/usr/bin/env python3
"""
NLP pre-transformer correction layer for SphinxEyes.

Role in the pipeline
    YOLOv11 (top-3 per slot)
    BigramModel       (Markov chain built from BBAW corpus)
        ViterbiSegmenter  (dynamic programming over SphinxTrie)
        ConfidenceGate    (Naive Bayes: respect high-conf YOLO predictions)
        UnknownResolver   (enumerate trie children at Unknown slots)
        correct()         (single entry point for FastAPI)
        CorrectionResult  (JSON-ready dict for GPT prompt)

Algorithms used
    1. Bigram language model  (Markov chain, order 1)
       Built from BBAW gardiner_seq column.
       P(code_B | code_A) = count(A->B) / count(A)
       Stored as log-probabilities to avoid float underflow.

    2. Viterbi segmenter
       Dynamic programming over the flat YOLO sequence.
       At each position i, tries all segments codes[i:j] (j-i <= MAX_WORD_LEN).
       Scores each candidate segment with:
           freq_score   = log(total_record_freq + 1)
           trans_score  = log P(segment[0] | prev_last_code)   [bigram]
           edit_penalty = 0.0 (exact) | -1.0 (fuzzy dist=1) | -2.5 (dist=2)
       Backtracks from dp[N] to recover the optimal word sequence.
       Fallback: if dp[N] == -inf, segments sign-by-sign (graceful degradation).

    3. Beam search over top-3 YOLO candidates
       At each slot, YOLO supplies up to 3 (code, confidence) pairs.
       The beam keeps the top-B partial sequences, expanding each with all
       3 candidates. At the end, the highest-scoring complete path wins.
       This replaces the single-best assumption and lets the corrector
       recover from YOLO's first-choice errors.

    4. Naive Bayes confidence gate
       Prevents the corrector from overriding a high-confidence YOLO prediction.
       P(correction_correct | yolo_conf, edit_dist) ~ (1-yolo_conf)/(edit_dist+1)
       If this probability is below a threshold, the original YOLO code is kept.

    5. Unknown-slot resolution
       When YOLO outputs 'Unknown' for a slot, the trie enumerates all child
       edges at that depth in surviving beam nodes and proposes the most
       frequent real Gardiner code as a replacement. gardiner_orig on the
       SourceRecord provides the ground-truth code the Unknown was substituted for.

Usage
    from sphinx_corrector import build_bigram_model, correct, CorrectionResult
    from sphinx_trie import SphinxTrie

    trie      = SphinxTrie.from_pickle('sphinx_trie_v4.pkl')
    log_prob, unigrams = build_bigram_model('bbaw_clean.parquet')

    # yolo_topk: list of slots, each slot is list of (code, confidence)
    yolo_topk = [
        [('G17', 0.97), ('G18', 0.02), ('Unknown', 0.01)],
        [('N35', 0.91), ('N36', 0.06), ('Z7',  0.03)],
        [('D21', 0.88), ('D22', 0.08), ('D19', 0.04)],
    ]
    result = correct(yolo_topk, trie, log_prob, unigrams)
    print(result['flat_translit'])
"""
from __future__ import annotations

import json
import math
import pickle
from collections import defaultdict, Counter
from dataclasses import dataclass, field, asdict
from pathlib import Path
from typing import Optional


# Public re-export so callers only need to import this module

try:
    from sphinx_trie import SphinxTrie, SourceRecord, UNKNOWN_TOKEN
except ImportError:
    # Allow the module to be imported standalone for unit tests
    SphinxTrie     = None       # type: ignore
    SourceRecord   = None       # type: ignore
    UNKNOWN_TOKEN  = 'Unknown'

# Constants

MAX_WORD_LEN    = 12    # max Gardiner codes in a single dictionary entry
BEAM_WIDTH      = 8     # number of partial paths kept during beam search
FUZZY_MAX_DIST  = 2     # max Levenshtein distance for trie fuzzy lookup
NEG_INF         = float('-inf')
YOLO_TRUST_THRESHOLD = 0.75
# Edit penalty table indexed by Levenshtein distance.
# Used when sub_cost_matrix is unavailable, or when the trie's fuzzy hit
# changed segment length (insertions/deletions — no per-pair lookup possible).
EDIT_PENALTY = {0: 0.0, 1: -1.0, 2: -2.5}


def _compute_edit_penalty(
    orig_seg : list[str],
    eff_seg  : list[str],
    edit_dist: int,
    sub_cost : Optional[dict],
) -> float:
    """
    Score contribution (negative, i.e. a penalty) for an edited segment.

    When `sub_cost` is None, or the original and effective lengths differ
    (insert/delete), fall back to the coarse `EDIT_PENALTY` table.

    When lengths match and `sub_cost` is provided, sum per-pair
    `-(-log P(truth|pred))` = `log P(truth|pred)` over the substituted
    positions. Frequently-confused pairs (high P) → small penalty;
    rarely-confused pairs (low P) → large penalty.
    """
    if edit_dist == 0:
        return 0.0
    
    if sub_cost is None or len(orig_seg) != len(eff_seg):
        return EDIT_PENALTY.get(edit_dist, -5.0)
    total = 0.0
    
    for o, e in zip(orig_seg, eff_seg):
        if o == e:
            continue
        
        c = sub_cost.get((o, e))
        
        if c is None:
            total += EDIT_PENALTY.get(1, -1.0)   # per-pair fallback
        else:
            total += -c     
            # cost ≥ 0; penalty ≤ 0
    return total

# Soft linguistic priors added to the Viterbi segment score.
# Tuned conservatively so corpus frequency still dominates; raise these
# if the segmenter under-uses determinatives / particles in practice.
DETERMINATIVE_BONUS = 0.5    # last code of segment is a determinative
PARTICLE_BONUS      = 1.0    # segment exactly matches a phrase-initial particle
LAYOUT_BONUS        = 1.5    # segment ends at an image-derived row/col break

# Lazy-loaded prior artifacts. Populated from determinatives.json and
# initial_particles.json on first use. Set to None to disable.
_DETERMINATIVES_CACHE: Optional[set] = None
_PARTICLES_CACHE: Optional[list] = None
_PRIORS_DIR = Path(__file__).parent

# Data Structures

@dataclass
class SegmentedWord:
    codes       : list[str]
    translit    : str
    translation : str
    freq        : int
    edit_dist   : int               # 0 = exact match, 1-2 = fuzzy corrected
    confidence  : float             # mean YOLO confidence over the slot range
    source      : str               # "dickson" | "bbaw" | "fallback"


@dataclass
class ResolvedUnknown:
    slot        : int               # index in the flat yolo sequence
    proposed    : str               # the Gardiner code we propose
    reason      : str               # human-readable justification
    freq        : int               # corpus frequency of the proposal


@dataclass
class CorrectionResult:
    segmented_words    : list[SegmentedWord]
    unknowns_resolved  : list[ResolvedUnknown]
    flat_corrected_seq : list[str]          # for GPT prompt
    flat_translit      : str                # for GPT prompt
    flat_translation   : str                # for GPT prompt (Dickson English)
    score              : float              # total Viterbi score
    had_fallback       : bool               # True if graceful degradation used

    def to_dict(self) -> dict:
        return {
            'segmented_words'   : [asdict(w) for w in self.segmented_words],
            'unknowns_resolved' : [asdict(u) for u in self.unknowns_resolved],
            'flat_corrected_seq': self.flat_corrected_seq,
            'flat_translit'     : self.flat_translit,
            'flat_translation'  : self.flat_translation,
            'score'             : self.score,
            'had_fallback'      : self.had_fallback,
        }

# Linguistic prior loaders (determinatives, initial particles)


def load_determinatives(path: Optional[Path] = None) -> set:
    """
    Load the set of Gardiner codes flagged as determinatives. Cached.
    Falls back to an empty set if determinatives.json is absent.
    """
    global _DETERMINATIVES_CACHE
    
    if _DETERMINATIVES_CACHE is not None:
        
        return _DETERMINATIVES_CACHE
    
    p = Path(path) if path else _PRIORS_DIR / 'determinatives.json'
    
    if not p.exists():
        
        _DETERMINATIVES_CACHE = set()
        
        return _DETERMINATIVES_CACHE
    
    with open(p, encoding='utf-8') as f:
        
        data = json.load(f)
        
    _DETERMINATIVES_CACHE = set(data.get('codes', {}).keys())
    
    return _DETERMINATIVES_CACHE


def load_initial_particles(path: Optional[Path] = None) -> list:
    """
    Load phrase-initial particles as a list of (gardiner_spelling, meta)
    tuples, sorted by spelling length descending so longer matches are
    tested first. Cached. Returns [] if initial_particles.json is absent.
    """
    global _PARTICLES_CACHE
    
    if _PARTICLES_CACHE is not None:
        
        return _PARTICLES_CACHE
    
    p = Path(path) if path else _PRIORS_DIR / 'initial_particles.json'
    
    if not p.exists():
        
        _PARTICLES_CACHE = []
        
        return _PARTICLES_CACHE
    
    with open(p, encoding='utf-8') as f:
        
        data = json.load(f)
        
    out: list = []
    
    for name, meta in data.get('particles', {}).items():
        
        spellings = [meta.get('gardiner', [])]
        
        spellings.extend(meta.get('alt_spellings', []))
        
        for spelling in spellings:
            
            if spelling:
                
                out.append((tuple(spelling),
                            
                            {'name': name, 'meaning': meta.get('meaning', '')}))
                
    out.sort(key=lambda kv: -len(kv[0]))
    
    _PARTICLES_CACHE = out
    
    return _PARTICLES_CACHE


# Bigram language model  (Markov chain, order 1)


def build_bigram_model(
    bbaw_source,   # str | Path to parquet, OR pandas DataFrame
    gardiner_col: str = 'gardiner_seq',
    smoothing: float = 1e-6,
) -> tuple[dict, Counter]:
    """
    Build a bigram language model from the BBAW corpus.

    Parameters:
    bbaw_source   : path to bbaw_clean.parquet  OR  a pandas DataFrame
    gardiner_col  : column containing space-separated Gardiner codes
    smoothing     : additive (Laplace) smoothing for unseen transitions.
                    Small positive float avoids log(0).

    Returns
    log_prob : dict[str, dict[str, float]]
        log_prob[A][B] = log P(B | A)
    unigrams : Counter
        Raw unigram counts, used for vocabulary size in the Viterbi.
    """
    try:
        import pandas as pd
    except ImportError:
        raise ImportError("pandas is required: pip install pandas")

    if isinstance(bbaw_source, (str, Path)):
        df = pd.read_parquet(bbaw_source)
    else:
        df = bbaw_source

    bigrams  = defaultdict(Counter)
    unigrams = Counter()

    for seq_str in df[gardiner_col].dropna():
        
        codes = seq_str.strip().split()
        
        if not codes:
            continue
        for c in codes:
            
            unigrams[c] += 1
            
        for a, b in zip(codes, codes[1:]):
            
            bigrams[a][b] += 1

    # Convert to log-probabilities with smoothing
    vocab_size = len(unigrams)
    log_prob: dict[str, dict[str, float]] = {}

    for a, followers in bigrams.items():
        
        total = sum(followers.values()) + smoothing * vocab_size
        
        log_prob[a] = {
            
            b: math.log((cnt + smoothing) / total)
            
            for b, cnt in followers.items()
        }

    return log_prob, unigrams


def transition_log_prob(
    prev_code : str,
    next_code : str,
    log_prob  : dict,
    unigrams  : Counter,
    smoothing : float = 1e-6,
) -> float:
    """
    Return log P(next_code | prev_code) with fallback to unigram smoothing.
    """
    if prev_code in log_prob and next_code in log_prob[prev_code]:
        
        return log_prob[prev_code][next_code]
    
    # Unigram fallback with smoothing
    vocab_size  = len(unigrams)
    total       = sum(unigrams.values()) + smoothing * vocab_size
    count       = unigrams.get(next_code, 0) + smoothing
    
    return math.log(count / total)



# Viterbi segmenter


def viterbi_segment(
    codes       : list[str],
    trie        : 'SphinxTrie',
    log_prob    : dict,
    unigrams    : Counter,
    sub_cost_matrix : Optional[dict] = None,
    max_word_len    : int = MAX_WORD_LEN,
    fuzzy_max_dist  : int = FUZZY_MAX_DIST,
    boundary_hints  : Optional[list] = None,
    determinatives  : Optional[set] = None,
    particles       : Optional[list] = None,
) -> tuple[list[tuple[list[str], list, int]], bool]:
    """
    Segment a flat sequence of Gardiner codes into the most probable
    sequence of words using dynamic programming over the trie.

    Parameters
    ----------
    codes           : flat list of Gardiner IDs from YOLO
    trie            : SphinxTrie with Dickson + BBAW entries
    log_prob        : bigram log-probability table
    unigrams        : unigram counts (for smoothing)
    sub_cost_matrix : optional dict[(a,b)] -> float substitution cost.
                      If None, uses uniform cost 1.0 for all substitutions.
    max_word_len    : maximum number of signs in a single dictionary word
    fuzzy_max_dist  : maximum Levenshtein distance for fuzzy trie lookup
    boundary_hints  : optional list[int] of slot indices where the image
                      breaks (end of row/column). Ending a segment at one
                      of these positions earns LAYOUT_BONUS.
    determinatives  : optional set of Gardiner codes treated as
                      determinatives. Defaults to lazy-load from
                      determinatives.json.
    particles       : optional list of (gardiner_spelling_tuple, meta)
                      pairs for phrase-initial particles. Defaults to
                      lazy-load from initial_particles.json.

    Returns
    -------
    path        : list of (segment_codes, records, edit_dist)
                  segment_codes is the post-fuzzy-correction segment when
                  a fuzzy hit fired; cursor accounting uses the original
                  span length internally so backtrack stays correct.
    had_fallback: True if graceful sign-by-sign fallback was used
    """
    N = len(codes)
    if N == 0:
        return [], False

    if determinatives is None:
        
        determinatives = load_determinatives()
        
    if particles is None:
        
        particles = load_initial_particles()
        
    boundary_set = set(boundary_hints) if boundary_hints else set()

    # dp[i]   = best total score for the prefix codes[0:i]
    # back[i] = (effective_segment, records, edit_dist, span)
    #   effective_segment : codes used for translit/transition (post-fuzzy)
    #   span              : original-input length consumed (= j-i at insert)
    #   The cursor walks `span`, never `len(effective_segment)`. This is the
    #   fix for the off-by-one when fuzzy returns a different-length match.
    dp   = [NEG_INF] * (N + 1)
    back = [None]    * (N + 1)
    dp[0] = 0.0

    for i in range(N):
        if dp[i] == NEG_INF:
            continue
    # dp[i]   = best total score for the prefix codes[0:i]
    # back[i] = (effective_segment, records, edit_dist, span)
    #   effective_segment : codes used for translit/transition (post-fuzzy)
        for j in range(i + 1, min(i + max_word_len + 1, N + 1)):
            
            original_segment   = codes[i:j]
            span               = j - i
            edit_dist          = 0
            effective_segment  = original_segment

            #  Exact trie lookup  --
            records = trie.search(original_segment)

            # Fuzzy trie lookup (Levenshtein) 
            if records is None and fuzzy_max_dist > 0:
                
                candidates = trie.fuzzy_search(
                    original_segment,
                    max_distance=fuzzy_max_dist,
                    max_results=1
                )
                if candidates:
                    best              = candidates[0]
                    edit_dist         = best['distance']
                    records           = best['records']
                    effective_segment = best['gardiner_seq'].split()

            if records is None:
                continue

            # --- Score components ---
            total_freq  = sum(r.freq for r in records)
            freq_score  = math.log(total_freq + 1)
            edit_pen    = _compute_edit_penalty(
                original_segment, effective_segment, edit_dist, sub_cost_matrix
            )

            # Bigram transition from previous word's last code
            if back[i] is not None and effective_segment:
                
                prev_eff_seg = back[i][0]
                
                if prev_eff_seg:
                    
                    trans_score = transition_log_prob(
                        prev_eff_seg[-1], effective_segment[0],
                        log_prob, unigrams
                    )
                else:
                    trans_score = 0.0
            else:
                trans_score = 0.0

            # --- Linguistic priors ---
            # Determinative bonus: last sign of segment is a determinative
            det_bonus = (DETERMINATIVE_BONUS
                         if (determinatives
                             and effective_segment
                             and effective_segment[-1] in determinatives)
                         else 0.0)

            # Particle bonus: this exact span matches a phrase-initial
            # particle's Gardiner spelling
            part_bonus = 0.0
            if particles:
                
                seg_tuple = tuple(original_segment)
                
                for spelling, _meta in particles:
                    
                    if len(spelling) == span and seg_tuple == spelling:
                        
                        part_bonus = PARTICLE_BONUS
                        break

            # Layout bonus: segment ends at an image-derived break
            layout_bonus = LAYOUT_BONUS if j in boundary_set else 0.0

            score = (dp[i] + freq_score + trans_score + edit_pen
                     + det_bonus + part_bonus + layout_bonus)

            if score > dp[j]:
                dp[j]   = score
                back[j] = (effective_segment, records, edit_dist, span)

    # --- Reconstruct path ---
    if dp[N] == NEG_INF:
        # Graceful fallback: one sign per segment, no trie records
        fallback = [([c], [], 0) for c in codes]
        return fallback, True

    path, pos = [], N
    while pos > 0 and back[pos] is not None:
        _eff_seg, recs, ed, span = back[pos]
        # Emit ORIGINAL codes consumed (the YOLO span), so downstream
        # cursor accounting in correct() stays simple. The trie's corrected
        # reading is conveyed via records (record.translit, record.gardiner_orig).
        original_seg = codes[pos - span:pos]
        path.append((original_seg, recs, ed))
        pos -= span

    # Handle any remaining prefix not covered (partial match)
    if pos > 0:
        for c in reversed(codes[:pos]):
            path.append(([c], [], 0))

    path.reverse()
    return path, False



#  Beam search over top-3 YOLO candidates


@dataclass
class BeamState:
    codes       : list[str]     = field(default_factory=list)
    confs       : list[float]   = field(default_factory=list)
    score       : float         = 0.0

    def extend(self, code: str, conf: float, trans_score: float) -> 'BeamState':
        return BeamState(
            codes  = self.codes  + [code],
            confs  = self.confs  + [conf],
            score  = self.score + math.log(conf + 1e-9) + trans_score,
        )


def beam_decode_topk(
    slots     : list[list[tuple[str, float]]],
    log_prob  : dict,
    unigrams  : Counter,
    beam_width: int = BEAM_WIDTH,
    trust_threshold:  float = YOLO_TRUST_THRESHOLD,
) ->  tuple[list[str], list[float]]:
    """
    Given a list of slots where each slot is [(code, conf), ...],
    return the single best flat code sequence using beam search
    with bigram language model scoring.

    Parameters
    ----------
    slots      : list of slots, each slot is a list of (code, confidence)
                 sorted by confidence descending (as YOLO produces them)
    log_prob   : bigram log-prob table
    unigrams   : unigram counts
    beam_width : number of partial hypotheses to keep at each step

    Returns
    -------
    best_codes : list[str]  the single best code sequence
    """
    if not slots:
        return [], []

    # Initialize beam with empty state
    beam: list[BeamState] = [BeamState()]

    for slot_candidates in slots:
        new_beam: list[BeamState] = []

        for state in beam:
            for code, conf in slot_candidates:
                
                trans = (transition_log_prob(state.codes[-1], code, log_prob, unigrams) if state.codes else 0.0)
                
                new_beam.append( state.extend(code, conf, trans) )
        
        new_beam.sort(key=lambda s: s.score, reverse=True)

        beam = new_beam[:beam_width]

    best = beam[0] if beam else BeamState()

    # Compare best codes to YOLO's top-1 per slot and revert if YOLO was more confident
    overridden = []
    
    for i, chosen in enumerate(best.codes):
        
        top1_code, top1_conf = slots[i][0]
        
        deviated = (chosen != top1_code)
        
        if deviated and top1_conf >= trust_threshold:
            
            best.codes[i] = top1_code     # revertir: YOLO manda
            overridden.append(False)
        else:
            overridden.append(deviated)

    return best.codes, overridden

# 4. Naive Bayes confidence gate


def should_correct(
    yolo_conf   : float,
    edit_dist   : int,
    threshold   : float = 0.30,
) -> bool:
    """
    Decide whether to apply the trie correction or keep the YOLO prediction.

    Uses a simple Bayesian estimate:
        P(correction_correct | yolo_conf, edit_dist)
            ~ (1 - yolo_conf) / (edit_dist + 1)

    If this probability is below `threshold`, the original YOLO code is kept
    (YOLO is probably right and the trie correction is risky).

    Parameters
    ----------
    yolo_conf  : YOLO's confidence for this prediction (0.0 to 1.0)
    edit_dist  : Levenshtein distance of the trie correction (0 = exact)
    threshold  : minimum probability required to apply correction

    Returns
    -------
    True  = apply the trie correction
    False = keep the original YOLO prediction
    """
    if edit_dist == 0:
        return True     # exact match: always accept
    p_correction = (1.0 - yolo_conf) / (edit_dist + 1)
    return p_correction >= threshold


# 5. Unknown-slot resolution

def resolve_unknowns(
    path       : list[tuple[list[str], list, int]],
    trie       : 'SphinxTrie',
    yolo_topk  : list[ list[ tuple[ str,  float]]],
    slot_offset: int = 0,
    trust_threshold: float = YOLO_TRUST_THRESHOLD,
) -> list[ResolvedUnknown]:
    """
    For each segment that contains the UNKNOWN_TOKEN edge, enumerate the
    trie's terminal records to find the most frequent real Gardiner code
    that was originally substituted.

    Uses the gardiner_orig field stored on SourceRecord during ingestion
    with oov_strategy='unknown'.

    Parameters

    path        : output of viterbi_segment
    trie        : SphinxTrie (needed for context; records are already in path)
    slot_offset : global offset of this path in the full YOLO sequence

    Returns

    list of ResolvedUnknown, one per Unknown slot that could be resolved
    """
    resolved = []
    current_slot = slot_offset

    for seg_codes, records, _ in path:

        for idx, code in enumerate(seg_codes):

            if code != UNKNOWN_TOKEN or not records:
                continue

            global_idx = current_slot + idx
            slot_candidates = (yolo_topk[global_idx]
                               if global_idx < len(yolo_topk) else [])
            yolo_alt_codes = {c for c, _ in slot_candidates
                              if c != UNKNOWN_TOKEN}

            # Tally the ground-truth codes this Unknown stood in for.
            orig_counter: Counter = Counter()
            for rec in records:
                if rec.gardiner_orig:
                    orig_codes = rec.gardiner_orig.split()
                    if idx < len(orig_codes):
                        orig_counter[orig_codes[idx]] += rec.freq

            if not orig_counter:
                continue

            overlap = [c for c in orig_counter if c in yolo_alt_codes]

            if overlap:
                # Best case: corpus proposal is also one of YOLO's top-3.
                best_code = max(overlap, key=lambda c: orig_counter[c])
                best_freq = orig_counter[best_code]
                reason = (f"corpus proposal '{best_code}' confirmed by a "
                          f"YOLO top-3 candidate at this slot")
            else:
                # No overlap: prefer a decent YOLO alternative, else corpus.
                decent_yolo_alt = next(
                    (c for c, conf in slot_candidates
                     if c != UNKNOWN_TOKEN and conf >= 0.2),
                    None,
                )
                if decent_yolo_alt is not None:
                    best_code = decent_yolo_alt
                    best_freq = 0
                    reason = (f"no corpus/YOLO overlap; falling back to "
                              f"YOLO's own alternative '{decent_yolo_alt}'")
                else:
                    best_code, best_freq = orig_counter.most_common(1)[0]
                    reason = (f"no usable YOLO alternative; corpus frequency "
                              f"proposal '{best_code}' (freq={best_freq})")

            resolved.append(ResolvedUnknown(
                slot=global_idx, proposed=best_code,
                reason=reason, freq=best_freq,
            ))

        current_slot += len(seg_codes)

    return resolved
                

# 6. correct() — single entry point for FastAPI

def correct(
    yolo_topk       : list[list[tuple[str, float]]],
    trie            : 'SphinxTrie',
    log_prob        : dict,
    unigrams        : Counter,
    sub_cost_matrix : Optional[dict] = None,
    conf_threshold  : float = 0.30,
    beam_width      : int   = BEAM_WIDTH,
    max_word_len    : int   = MAX_WORD_LEN,
    fuzzy_max_dist  : int   = FUZZY_MAX_DIST,
    boundary_hints  : Optional[list] = None,
) -> CorrectionResult:
    """
    Main entry point. Takes YOLO's top-K predictions per slot and returns
    a fully corrected, segmented, transliterated result ready for GPT.

    Parameters
    ----------
    yolo_topk       : list of slots. Each slot is a list of (code, confidence)
                      tuples sorted by confidence descending.
                      Minimum: 1 tuple per slot. Maximum recommended: 3.

    trie            : SphinxTrie loaded from sphinx_trie_v4.pkl
    log_prob        : bigram log-prob table from build_bigram_model()
    unigrams        : unigram counter from build_bigram_model()
    sub_cost_matrix : optional dict[(a,b)] -> float.
                      If provided, overrides uniform substitution cost.
                      Load from confusion_matrix.csv after V2 training.
    conf_threshold  : Naive Bayes threshold (default 0.30).
                      Lower = more aggressive correction.
                      Higher = more conservative, trust YOLO more.
    beam_width      : beam search width (default 8)
    max_word_len    : max signs per dictionary entry (default 12)
    fuzzy_max_dist  : max Levenshtein distance for trie fuzzy lookup (default 2)

    Returns
    -------
    CorrectionResult dataclass. Call .to_dict() for JSON serialization.
    """
    if not yolo_topk:
        return CorrectionResult(
            segmented_words    = [],
            unknowns_resolved  = [],
            flat_corrected_seq = [],
            flat_translit      = '',
            flat_translation   = '',
            score              = NEG_INF,
            had_fallback       = False,
        )

    #   Beam search over top-3 candidates to get best flat sequence
 
    best_codes, beam_overrides = beam_decode_topk(
        yolo_topk, log_prob, unigrams, beam_width=beam_width
    )

    # Mean confidence per slot (for the confidence gate)
    slot_confs = [
        candidates[0][1] if candidates else 0.5
        for candidates in yolo_topk
    ]
    #  Viterbi segmentation over best_codes
 
    path, had_fallback = viterbi_segment(
        best_codes, trie, log_prob, unigrams,
        sub_cost_matrix = sub_cost_matrix,
        max_word_len    = max_word_len,
        fuzzy_max_dist  = fuzzy_max_dist,
        boundary_hints  = boundary_hints,
    )

 
    # Apply confidence gate — revert corrections where YOLO was
    # more reliable than the trie correction
 
    gated_path = []
    cursor     = 0
    for seg_codes, records, edit_dist in path:
        seg_len    = len(seg_codes)
        mean_conf  = (
            sum(slot_confs[cursor:cursor + seg_len]) / seg_len
            if seg_len > 0 else 0.5
        )

        if edit_dist > 0 and not should_correct(mean_conf, edit_dist, conf_threshold):
            # Revert: keep original YOLO codes for this segment
            orig_codes = best_codes[cursor:cursor + seg_len]
            gated_path.append((orig_codes, [], 0))
        else:
            gated_path.append((seg_codes, records, edit_dist))

        cursor += seg_len
 
    #  Resolve Unknown slots
 
    unknowns_resolved = resolve_unknowns(
    gated_path, trie, yolo_topk, slot_offset=0
    )

    # Apply Unknown resolutions to the corrected sequence
    flat_corrected = []
    for seg_codes, records, _ in gated_path:
        flat_corrected.extend(seg_codes)

    for resolution in unknowns_resolved:
        if resolution.slot < len(flat_corrected):
            flat_corrected[resolution.slot] = resolution.proposed

 
    #  Build SegmentedWord list and flat strings
 
    segmented_words : list[SegmentedWord] = []
    total_score      = 0.0
    cursor           = 0

    for seg_codes, records, edit_dist in gated_path:
        seg_len   = len(seg_codes)
        mean_conf = (
            sum(slot_confs[cursor:cursor + seg_len]) / seg_len
            if seg_len > 0 else 0.5
        )

        # Pick best record by frequency
        if records:
            best_rec = max(records, key=lambda r: r.freq)
            translit    = best_rec.translit
            translation = best_rec.translation
            source      = best_rec.source
            freq        = best_rec.freq
        else:
            translit    = ' '.join(seg_codes)
            translation = ''
            source      = 'fallback'
            freq        = 0

        total_score += math.log(freq + 1) + EDIT_PENALTY.get(edit_dist, -5.0)

        segmented_words.append(SegmentedWord(
            codes       = seg_codes,
            translit    = translit,
            translation = translation,
            freq        = freq,
            edit_dist   = edit_dist,
            confidence  = round(mean_conf, 4),
            source      = source,
        ))

        cursor += seg_len

    flat_translit   = ' '.join(w.translit    for w in segmented_words if w.translit)
    flat_translation = ' | '.join(
        w.translation for w in segmented_words
        if w.translation and w.source != 'fallback'
    )

    return CorrectionResult(
        segmented_words    = segmented_words,
        unknowns_resolved  = unknowns_resolved,
        flat_corrected_seq = flat_corrected,
        flat_translit      = flat_translit,
        flat_translation   = flat_translation,
        score              = total_score,
        had_fallback       = had_fallback,
    )


# ---------------------------------------------------------------------------
# Convenience: load artifacts from disk
# ---------------------------------------------------------------------------

def load_corrector(
    trie_pkl     : str | Path,
    bbaw_parquet : str | Path,
    confusion_csv: Optional[str | Path] = None,
    gardiner_col : str = 'gardiner_seq',
) -> tuple['SphinxTrie', dict, Counter, Optional[dict]]:
    """
    Load every artifact needed by correct() in one call.

    Parameters
    ----------
    trie_pkl      : path to sphinx_trie_v4.pkl
    bbaw_parquet  : path to bbaw_clean.parquet (for the bigram LM)
    confusion_csv : optional path to a YOLO confusion-matrix CSV
                    (e.g. confusion_matrix_v2_on_v3val_normalized.csv).
                    When provided, a per-pair substitution-cost dict is
                    built and returned; otherwise the 4th return is None
                    and the corrector falls back to coarse EDIT_PENALTY.
    gardiner_col  : column name inside bbaw_parquet

    Returns
    -------
    trie     : SphinxTrie
    log_prob : bigram log-probability table
    unigrams : unigram Counter
    sub_cost : dict[(pred, truth) -> -log P(truth|pred)] or None
    """
    if SphinxTrie is None:
        raise ImportError("sphinx_trie.py must be on the Python path")

    trie = SphinxTrie.from_pickle(trie_pkl)
    log_prob, unigrams = build_bigram_model(bbaw_parquet, gardiner_col)
    sub_cost = load_sub_cost_matrix(confusion_csv) if confusion_csv else None
    return trie, log_prob, unigrams, sub_cost


# ---------------------------------------------------------------------------
# Cost matrix loader (stub — activate after V2 training)
# ---------------------------------------------------------------------------

def load_sub_cost_matrix(
    confusion_csv: str | Path,
    classes      : Optional[list[str]] = None,
    epsilon      : float = 1e-4,
) -> dict:
    """
    Build a per-pair substitution cost dict from a YOLO confusion matrix CSV.

    CSV orientation (Ultralytics `ConfusionMatrix.matrix` convention, after
    the V3 notebook's `axis=1` row-normalization):
        rows    = predicted class
        columns = truth class
        cell[pred, truth] = P(truth | pred)

    The extra `background` row/column that Ultralytics adds (index nc) is
    dropped — substituting against background is never a valid Gardiner edit.

    Cost convention
    ---------------
        cost(pred, truth) = -log( P(truth | pred) + epsilon )    # ≥ 0
            ~0  for frequently-confused pairs   (cheap to substitute)
            ~9.2 for never-observed pairs       (expensive to substitute,
                                                 epsilon=1e-4 caps -log here)

    The caller's score adds `-cost` so the penalty contribution is ≤ 0,
    matching the sign convention of `EDIT_PENALTY`. See `_compute_edit_penalty`.

    Parameters
    ----------
    confusion_csv : path to e.g. `confusion_matrix_v2_on_v3val_normalized.csv`
    classes       : optional whitelist of Gardiner codes to keep. If None,
                    keeps everything in the CSV except `background`.
    epsilon       : Laplace floor (1e-4 → -log caps near 9.2 instead of inf)

    Returns
    -------
    dict[(pred_code, truth_code)] -> float cost (non-negative).
    Pairs absent from the matrix are silently absent from the dict;
    `_compute_edit_penalty` then falls back to `EDIT_PENALTY[1]`.

    Usage
    -----
        sub_cost = load_sub_cost_matrix(
            'confusion_matrix_v2_on_v3val_normalized.csv'
        )
        result = correct(yolo_topk, trie, log_prob, unigrams,
                         sub_cost_matrix=sub_cost)
    """
    try:
        import pandas as pd
    except ImportError:
        raise ImportError("pandas required: pip install pandas")

    df = pd.read_csv(confusion_csv, index_col=0)

    # Drop Ultralytics' nc+1 "background" axis on both ends if present.
    for label in ('background', 'Background'):
        if label in df.index:
            df = df.drop(index=label)
        if label in df.columns:
            df = df.drop(columns=label)

    # Optional whitelist
    if classes is not None:
        keep_idx = [c for c in classes if c in df.index]
        keep_col = [c for c in classes if c in df.columns]
        df = df.loc[keep_idx, keep_col]

    # The CSV is already row-normalized (per V3 notebook cell 121:
    # `cm_norm = cm / cm.sum(axis=1)`). Read rows directly as P(truth|pred).
    cost = {}
    rows, cols = list(df.index), list(df.columns)
    mat = df.values.astype(float)
    for i, pred in enumerate(rows):
        for j, truth in enumerate(cols):
            p = float(mat[i, j])
            cost[(pred, truth)] = -math.log(p + epsilon)
            
    return cost