utils/confidence.py

"""
InvoiceForge AI — utils/confidence.py
Multi-engine OCR confidence fusion.

Implements weighted voting at the token level:
  - PaddleOCR weight: 0.50
  - EasyOCR weight:   0.30
  - TrOCR weight:     0.20  (only for handwritten documents)

Tokens from multiple engines are spatially matched using IoU overlap of
bounding boxes, then the highest weighted-confidence text is selected.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field

import numpy as np

logger = logging.getLogger(__name__)

# ─────────────────────────────────────────────────────────────────────────────
# CONSTANTS
# ─────────────────────────────────────────────────────────────────────────────
ENGINE_WEIGHTS: dict[str, float] = {
    "paddleocr": 0.50,
    "easyocr":   0.30,
    "trocr":     0.20,
}

IOU_MATCH_THRESHOLD: float = 0.40  # Minimum IoU to consider tokens overlapping
X_SNAP_PX: float = 30.0            # Fallback positional match if no bbox
Y_SNAP_PX: float = 15.0


@dataclass
class FusedToken:
    """Token produced after multi-engine confidence fusion."""
    text: str
    confidence: float
    x: float
    y: float
    bbox: list = field(default_factory=list)
    engine: str = "fused"
    candidates: list[dict] = field(default_factory=list)


# ─────────────────────────────────────────────────────────────────────────────
# IOu HELPERS
# ─────────────────────────────────────────────────────────────────────────────

def _bbox_to_xyxy(bbox: list) -> tuple[float, float, float, float]:
    """Convert [[x1,y1],[x2,y1],[x2,y2],[x1,y2]] to (x1,y1,x2,y2)."""
    if not bbox:
        return (0.0, 0.0, 0.0, 0.0)
    xs = [p[0] for p in bbox]
    ys = [p[1] for p in bbox]
    return (min(xs), min(ys), max(xs), max(ys))


def _iou(a: list, b: list) -> float:
    """Compute intersection-over-union of two bboxes."""
    ax1, ay1, ax2, ay2 = _bbox_to_xyxy(a)
    bx1, by1, bx2, by2 = _bbox_to_xyxy(b)

    ix1 = max(ax1, bx1)
    iy1 = max(ay1, by1)
    ix2 = min(ax2, bx2)
    iy2 = min(ay2, by2)

    if ix2 <= ix1 or iy2 <= iy1:
        return 0.0

    inter = (ix2 - ix1) * (iy2 - iy1)
    area_a = max((ax2 - ax1) * (ay2 - ay1), 1e-6)
    area_b = max((bx2 - bx1) * (by2 - by1), 1e-6)
    return inter / (area_a + area_b - inter)


def _pos_match(a: dict, b: dict) -> bool:
    """Positional proximity match when bounding boxes are unavailable."""
    return abs(a["x"] - b["x"]) < X_SNAP_PX and abs(a["y"] - b["y"]) < Y_SNAP_PX


# ─────────────────────────────────────────────────────────────────────────────
# CONFIDENCE FUSION
# ─────────────────────────────────────────────────────────────────────────────

class ConfidenceFusion:
    """
    Fuses OCR outputs from multiple engines using confidence-weighted voting.

    Usage:
        fusion = ConfidenceFusion()
        fusion.add_engine_result("paddleocr", paddle_tokens)
        fusion.add_engine_result("easyocr", easy_tokens)
        fused = fusion.fuse()
    """

    def __init__(self) -> None:
        self._engine_results: dict[str, list[dict]] = {}

    def add_engine_result(self, engine_name: str, tokens: list) -> None:
        """
        Add OCR results from one engine.

        Args:
            engine_name: One of "paddleocr", "easyocr", "trocr".
            tokens:      List of token dicts or OCRToken objects.
        """
        normalised: list[dict] = []
        for tok in tokens:
            if hasattr(tok, "__dict__"):
                # OCRToken dataclass
                normalised.append(
                    {
                        "text": tok.text,
                        "confidence": tok.confidence,
                        "bbox": tok.bbox,
                        "x": tok.x,
                        "y": tok.y,
                        "engine": engine_name,
                    }
                )
            else:
                d = dict(tok)
                d["engine"] = engine_name
                normalised.append(d)

        self._engine_results[engine_name] = normalised
        logger.debug(
            "Fusion: added %d tokens from %s.", len(normalised), engine_name
        )

    def fuse(self) -> list[FusedToken]:
        """
        Perform confidence-weighted fusion.

        Algorithm:
          1. Build a list of all tokens across engines.
          2. For each unmatched token, find all overlapping tokens from
             other engines using IoU or positional proximity.
          3. Compute weighted confidence for each text candidate.
          4. Select the text with the highest weighted score.

        Returns:
            Sorted list of FusedToken objects.
        """
        all_tokens: list[dict] = []
        for engine, tokens in self._engine_results.items():
            for tok in tokens:
                tok = dict(tok)
                tok["engine"] = engine
                tok["_used"] = False
                all_tokens.append(tok)

        fused_tokens: list[FusedToken] = []
        used_indices: set[int] = set()

        for i, anchor in enumerate(all_tokens):
            if i in used_indices:
                continue

            group: list[dict] = [anchor]
            used_indices.add(i)

            # Find matching tokens from other engines
            for j, candidate in enumerate(all_tokens):
                if j in used_indices:
                    continue
                if candidate["engine"] == anchor["engine"]:
                    continue

                matched = False
                if anchor["bbox"] and candidate["bbox"]:
                    matched = _iou(anchor["bbox"], candidate["bbox"]) >= IOU_MATCH_THRESHOLD
                else:
                    matched = _pos_match(anchor, candidate)

                if matched:
                    group.append(candidate)
                    used_indices.add(j)

            # Weighted voting across group
            fused = _vote(group)
            fused_tokens.append(fused)

        # Sort by reading order (top to bottom, left to right)
        fused_tokens.sort(key=lambda t: (t.y, t.x))
        logger.debug("Fusion: %d fused tokens produced.", len(fused_tokens))
        return fused_tokens

    def reset(self) -> None:
        """Clear all engine results for reuse."""
        self._engine_results.clear()


def _vote(group: list[dict]) -> FusedToken:
    """
    Select the best text from a group of candidate tokens using
    confidence-weighted voting.
    """
    text_scores: dict[str, float] = {}
    best_bbox: list = []
    best_x: float = 0.0
    best_y: float = 0.0

    for tok in group:
        engine = tok.get("engine", "paddleocr")
        weight = ENGINE_WEIGHTS.get(engine, 0.30)
        text = tok.get("text", "").strip()
        conf = float(tok.get("confidence", 0.5))
        weighted = conf * weight

        text_scores[text] = text_scores.get(text, 0.0) + weighted

        if tok.get("bbox"):
            best_bbox = tok["bbox"]
        best_x = tok.get("x", 0.0)
        best_y = tok.get("y", 0.0)

    # Pick text with highest aggregate weighted score
    if not text_scores:
        return FusedToken(text="", confidence=0.0, x=best_x, y=best_y)

    best_text = max(text_scores, key=lambda t: text_scores[t])
    total_weight = sum(ENGINE_WEIGHTS.values())
    normalised_conf = min(text_scores[best_text] / total_weight, 1.0)

    return FusedToken(
        text=best_text,
        confidence=round(normalised_conf, 4),
        x=best_x,
        y=best_y,
        bbox=best_bbox,
        engine="fused",
        candidates=group,
    )


# ─────────────────────────────────────────────────────────────────────────────
# DOCUMENT-LEVEL CONFIDENCE SCORING
# ─────────────────────────────────────────────────────────────────────────────

def compute_document_confidence(
    fused_tokens: list[FusedToken],
    validation_errors: list[str],
) -> float:
    """
    Compute an aggregate document-level confidence score.

    Factors:
      - Mean per-token confidence from fusion
      - Penalty for each validation error (−0.05 per error, max −0.30)

    Returns:
        Float 0.0 – 1.0.
    """
    if not fused_tokens:
        return 0.0

    mean_conf = float(np.mean([t.confidence for t in fused_tokens]))
    error_penalty = min(len(validation_errors) * 0.05, 0.30)
    score = max(0.0, mean_conf - error_penalty)
    return round(score, 4)