src/ocr/ocr_cleaner.py

"""
Post-traitement OCR — Module d'auto-correction (ocr_cleaner.py)

Pipeline appliqué sur la sortie de doctr_engine.to_layoutxlm_format() :
    1. Normalisation arabe    → formes Unicode canoniques (Alif, Ya, diacritiques)
    2. Nettoyage numérique   → confusion O/0, l/1, S/5 dans les zones de montants
    3. Correction fuzzy FR   → mots-clés documents commerciaux (rapidfuzz ≥ 85 %)

Les bounding boxes ne sont JAMAIS modifiées.
Implémentation vectorisée (pandas + numpy) pour rester rapide sur CPU 8 Go.

Dépendances :
    pip install rapidfuzz pandas numpy
"""

import re
import numpy as np
import pandas as pd
from typing import List, Dict, Optional, Tuple

from rapidfuzz import process as fuzz_process
from rapidfuzz import fuzz as fuzz_metrics

# ─────────────────────────────────────────────────────────────────────────────
# Dictionnaire bilingue de mots-clés documents commerciaux
# ─────────────────────────────────────────────────────────────────────────────

# Format : "forme canonique" → [variantes OCR fréquentes tolérées]
# La correction remplace toute variante par la forme canonique.

FR_KEYWORDS: Dict[str, List[str]] = {
    "Facture":        ["Factur", "Facturre", "Factuре", "F4cture", "Facfure"],
    "Devis":          ["Dev1s", "Devi5", "Deviss"],
    "Total":          ["T0tal", "Totol", "Tatal", "Tot4l", "Totaï"],
    "Sous-total":     ["Sous-t0tal", "Soustotal", "SousTotal"],
    "TVA":            ["TV4", "TUA", "IVA", "T/A"],
    "HT":             ["H7", "|-|T"],
    "TTC":            ["TIC", "T7C", "TТС"],
    "Montant":        ["Mont4nt", "Montamt", "M0ntant"],
    "Prix":           ["Pr1x", "Pr!x", "Prіx"],
    "Quantité":       ["Quant1té", "Quantité", "Quantitе", "Qte", "Qt€"],
    "Date":           ["D4te", "Dаte", "Dat3"],
    "Client":         ["Cl1ent", "Clіent", "C|ient"],
    "Fournisseur":    ["Fourni55eur", "F0urnisseur", "Fournissuer"],
    "Référence":      ["Ref", "Rеférence", "Ref.", "Réf.", "Réf"],
    "Numéro":         ["Numero", "N°", "Nr", "Num", "Num.", "Numer0"],
    "Paiement":       ["Pa1ement", "Рaiement", "Paiament"],
    "Remise":         ["Rem1se", "Rem!se"],
    "Adresse":        ["Adres5e", "Adresse", "Adr."],
    "Description":    ["Descr1ption", "Descripti0n", "Descript."],
    "Bon":            ["B0n", "8on"],
    "Commande":       ["C0mmande", "Comm4nde"],
    "Reçu":           ["Re_cu", "Recu", "Reçu"],
    "Solde":          ["S0lde", "S01de"],
    "Échéance":       ["Echeance", "Ech.", "Échéance"],
}

AR_KEYWORDS: Dict[str, List[str]] = {
    "فاتورة":   ["فاتوره", "فاتوررة", "ف4تورة"],
    "إجمالي":   ["اجمالي", "إجم4لي", "اجم4لي"],
    "مجموع":    ["مجم0ع", "مجموع"],
    "ضريبة":    ["ضريبه", "ضر1بة", "ضريب"],
    "تاريخ":    ["ت4ريخ", "تار1خ"],
    "عميل":     ["عم1ل", "عمیل"],
    "مورد":     ["م0رد", "مورد"],
    "مرجع":     ["مرج3", "مرجع"],
    "كمية":     ["كم1ة", "كميه"],
    "سعر":      ["س3ر", "سعر"],
    "مبلغ":     ["مبل_غ", "مبلق"],
    "دفع":      ["دف_ع", "دفع"],
    "رقم":      ["رق_م", "رقن"],
    "وصف":      ["وص_ف", "وصف"],
    "خصم":      ["خص_م", "خصم"],
    "ملاحظات":  ["ملاحظ4ت", "ملاحضات"],
    "المجموع":  ["المجم0ع", "المجموع"],
    "الإجمالي": ["الاجمالي", "الإجم4لي"],
    "اسم":      ["4سم", "أسم"],
    "عنوان":    ["عن04ن", "عنوان"],
}

# Liste plate de tous les mots canoniques (pour rapidfuzz)
_ALL_CANONICAL: List[str] = list(FR_KEYWORDS.keys()) + list(AR_KEYWORDS.keys())

# ─────────────────────────────────────────────────────────────────────────────
# Patterns regex — compilés une seule fois au chargement du module
# ─────────────────────────────────────────────────────────────────────────────

# Confusions OCR : caractère erroné → chiffre correct
# Appliqué UNIQUEMENT aux mots détectés comme zones numériques
_CONFUSION_MAP: Dict[str, str] = {
    "O": "0", "o": "0",
    "l": "1", "I": "1", "i": "1",
    "S": "5", "s": "5",
    "B": "8",
    "Z": "2",
    "G": "6",
    "q": "9", "g": "9",
    "T": "7",
}

# Pattern pour un mot « candidat numérique » :
# contient au moins un chiffre ET uniquement des chars digit-like ou ponctuation
_RE_NUMERIC_CANDIDATE = re.compile(
    r'^[0-9OolIiSsBZGqTg,.\s\-\+/\u066B\u066C\u2212]+$'
    # \u066B = virgule décimale arabe, \u066C = séparateur milliers arabe
)

# Montant typique (ex: 1 234,56  ou  1.234,56  ou  ١٢٣٤)
_RE_AMOUNT = re.compile(
    r'^\d[\d\s.,\u0660-\u0669]*$'   # \u0660-\u0669 = chiffres arabes-indiens
)

# ── Normalisation arabe ───────────────────────────────────────────────────────

# Toutes les variantes d'Alif → Alif nu (ا)
_RE_ALIF = re.compile(r'[أإآٱٲٳ]')

# Ya final ambigu (ى) → Ya standard (ي)
_RE_YA   = re.compile(r'ى')

# Kaf persan (ک) → Kaf arabe (ك)
_RE_KAF  = re.compile(r'ک')

# Diacritiques (harakat, shadda, sukun, maddah, hamza haute/basse…)
# Plages Unicode : U+064B–U+065F  et  U+0610–U+061A  et  U+06D6–U+06ED
_RE_DIACRITICS = re.compile(
    r'[\u064B-\u065F\u0610-\u061A\u06D6-\u06DC\u06DF-\u06E4'
    r'\u06E7\u06E8\u06EA-\u06ED\u0640]'
)

# Tatweel (ـ, U+0640) — allongement décoratif, inutile pour l'OCR
_RE_TATWEEL = re.compile(r'\u0640')


# ─────────────────────────────────────────────────────────────────────────────
# Point d'entrée principal
# ─────────────────────────────────────────────────────────────────────────────

def clean_ocr_output(
    words: List[dict],
    fuzzy_threshold:    int   = 85,
    apply_arabic_norm:  bool  = True,
    apply_numeric_fix:  bool  = True,
    apply_fuzzy:        bool  = True,
    remove_diacritics:  bool  = True,
) -> List[dict]:
    """
    Pipeline complet de post-traitement sur la liste de mots docTR.

    Ordre d'application (important) :
        1. Normalisation arabe (simplifie avant le fuzzy)
        2. Nettoyage numérique (O→0, l→1 dans les zones de montants)
        3. Correction fuzzy FR/AR (mots-clés documents commerciaux)

    Les bounding boxes ne sont PAS modifiées.

    Args:
        words            : Sortie de doctr_engine.to_layoutxlm_format()
                           [{"text","box","confidence","is_arabic"}, ...]
        fuzzy_threshold  : Similarité minimale pour la correction (0-100).
                           85 = bon équilibre précision/rappel.
        apply_arabic_norm: Activer la normalisation Unicode arabe.
        apply_numeric_fix: Activer le nettoyage des montants.
        apply_fuzzy      : Activer la correction fuzzy des mots-clés.
        remove_diacritics: Supprimer les diacritiques arabes (harakat…).

    Returns:
        Liste de dicts avec "text" corrigé et "correction" (log des changements).
    """
    if not words:
        return words

    # ── Conversion en DataFrame (opérations vectorisées) ─────────────────────
    df = pd.DataFrame(words).copy()
    df["original_text"] = df["text"]     # conserver l'original pour log
    df["correction"]    = ""              # journal des corrections

    # Masques booléens réutilisés
    mask_ar = df["is_arabic"].astype(bool)
    mask_la = ~mask_ar

    # ── 1. Normalisation arabe ────────────────────────────────────────────────
    if apply_arabic_norm and mask_ar.any():
        df.loc[mask_ar, "text"] = _normalize_arabic_batch(
            df.loc[mask_ar, "text"],
            remove_diacritics=remove_diacritics,
        )
        changed = df.loc[mask_ar, "text"] != df.loc[mask_ar, "original_text"]
        df.loc[mask_ar & changed, "correction"] += "arabic_norm "

    # ── 2. Nettoyage numérique ────────────────────────────────────────────────
    if apply_numeric_fix:
        numeric_mask = df["text"].apply(_is_numeric_candidate)
        if numeric_mask.any():
            df.loc[numeric_mask, "text"] = _sanitize_numeric_batch(
                df.loc[numeric_mask, "text"]
            )
            changed = df.loc[numeric_mask, "text"] != df.loc[numeric_mask, "original_text"]
            df.loc[numeric_mask & changed, "correction"] += "numeric_fix "

    # ── 3. Correction fuzzy ───────────────────────────────────────────────────
    if apply_fuzzy:
        # Latin (français/anglais)
        if mask_la.any():
            fr_keys = list(FR_KEYWORDS.keys())
            corrected_la = _fuzzy_correct_batch(
                df.loc[mask_la, "text"],
                choices=fr_keys,
                threshold=fuzzy_threshold,
            )
            changed = corrected_la != df.loc[mask_la, "text"]
            df.loc[mask_la, "text"] = corrected_la
            df.loc[mask_la & changed, "correction"] += "fuzzy_fr "

        # Arabe
        if mask_ar.any():
            ar_keys = list(AR_KEYWORDS.keys())
            corrected_ar = _fuzzy_correct_batch(
                df.loc[mask_ar, "text"],
                choices=ar_keys,
                threshold=fuzzy_threshold,
            )
            changed = corrected_ar != df.loc[mask_ar, "text"]
            df.loc[mask_ar, "text"] = corrected_ar
            df.loc[mask_ar & changed, "correction"] += "fuzzy_ar "

    df["correction"] = df["correction"].str.strip()
    return df.to_dict(orient="records")


# ─────────────────────────────────────────────────────────────────────────────
# Étape 1 — Normalisation arabe (vectorisée)
# ─────────────────────────────────────────────────────────────────────────────

def _normalize_arabic_batch(
    series: pd.Series,
    remove_diacritics: bool = True,
) -> pd.Series:
    """
    Normalise les caractères arabes pour uniformiser la représentation Unicode.

    Transformations :
        أ إ آ ٱ  →  ا    (variantes Alif → Alif nu)
        ى        →  ي    (Ya final ambigu → Ya standard)
        ک        →  ك    (Kaf persan → Kaf arabe)
        ـ        →  ""   (Tatweel décoratif supprimé)
        ً ٌ ٍ …  →  ""   (diacritiques supprimés si remove_diacritics=True)

    Pourquoi normaliser ?
        - L'OCR peut reconnaître « أحمد » avec différentes formes d'Alif.
        - La recherche/comparaison exige une forme unique.
        - Réduit le vocabulaire effectif → améliore le fuzzy matching.

    Args:
        series           : pd.Series de chaînes arabes brutes.
        remove_diacritics: Si True, supprime harakat, shadda, sukun…

    Returns:
        pd.Series normalisée.
    """
    s = series.copy()

    # Alif
    s = s.str.replace(_RE_ALIF,      "ا",  regex=True)
    # Ya final
    s = s.str.replace(_RE_YA,        "ي",  regex=True)
    # Kaf persan
    s = s.str.replace(_RE_KAF,       "ك",  regex=True)
    # Tatweel
    s = s.str.replace(_RE_TATWEEL,   "",   regex=True)
    # Diacritiques
    if remove_diacritics:
        s = s.str.replace(_RE_DIACRITICS, "", regex=True)

    return s


# ─────────────────────────────────────────────────────────────────────────────
# Étape 2 — Nettoyage numérique (vectorisé)
# ─────────────────────────────────────────────────────────────────────────────

def _is_numeric_candidate(text: str) -> bool:
    """
    Détermine si un mot est dans une zone numérique (montant, code…).

    Critères :
        - Contient au moins un chiffre (ASCII ou arabe-indien U+0660-U+0669)
        - Tous les autres caractères sont des confusions connues ou de la ponctuation

    Exemples :
        "1234,00"  → True   (montant normal)
        "1O0,OO"   → True   (montant avec confusion O/0)
        "Facture"  → False  (mot régulier)
        "TVA"      → False  (sigle mais pas numérique)
    """
    if not text:
        return False
    has_digit = bool(re.search(r'[\d\u0660-\u0669]', text))
    return has_digit and bool(_RE_NUMERIC_CANDIDATE.match(text))


def _sanitize_numeric_batch(series: pd.Series) -> pd.Series:
    """
    Remplace les confusions OCR dans les zones numériques.

    Confusions traitées (uniquement pour mots candidats numériques) :
        O/o → 0   (lettre O confondue avec zéro)
        l/I/i → 1 (lettre L minuscule ou I confondue avec un)
        S/s → 5   (S confondu avec cinq)
        B → 8     (B confondu avec huit)
        Z → 2     (Z confondu avec deux)
        G → 6     (G confondu avec six)
        q/g → 9   (q ou g confondus avec neuf)

    Stratégie vectorisée :
        Un seul str.replace() par caractère de confusion
        (plus rapide qu'une boucle Python sur chaque mot).

    Args:
        series : pd.Series de mots candidats numériques.

    Returns:
        pd.Series avec confusions corrigées.
    """
    s = series.copy()
    for wrong, right in _CONFUSION_MAP.items():
        # On ne remplace que dans le contexte numérique (déjà filtré par le masque)
        s = s.str.replace(wrong, right, regex=False)
    return s


# ─────────────────────────────────────────────────────────────────────────────
# Étape 3 — Correction fuzzy (vectorisée via rapidfuzz)
# ─────────────────────────────────────────────────────────────────────────────

def _fuzzy_correct_batch(
    series: pd.Series,
    choices: List[str],
    threshold: int = 85,
) -> pd.Series:
    """
    Corrige les mots-clés mal reconnus par similarité fuzzy.

    Utilise rapidfuzz.process.cdist() — calcul vectorisé de toutes
    les distances en une seule passe NumPy (évite la boucle Python).

    Algorithme :
        - Pour chaque mot, calcule la similarité avec chaque mot-clé.
        - Si la meilleure similarité > threshold ET la longueur est proche,
          remplace le mot par la forme canonique.
        - Les mots courts (≤ 2 chars) sont ignorés (trop de faux positifs).

    Args:
        series    : pd.Series de textes à corriger.
        choices   : Liste des formes canoniques (mots-clés).
        threshold : Seuil de similarité (0-100), 85 recommandé.

    Returns:
        pd.Series corrigée.
    """
    if series.empty or not choices:
        return series

    texts = series.tolist()

    # Filtrer les mots trop courts (évite les faux positifs)
    eligible_mask = np.array([len(t) > 2 for t in texts])

    if not eligible_mask.any():
        return series

    eligible_texts = [t for t, ok in zip(texts, eligible_mask) if ok]

    # ── Calcul matriciel des similarités (vectorisé rapidfuzz) ───────────────
    # score_matrix shape : (n_eligible_words, n_choices)
    score_matrix = fuzz_process.cdist(
        eligible_texts,
        choices,
        scorer=fuzz_metrics.ratio,    # ratio simple (0-100)
        workers=1,                    # pas de multiprocessing (économie RAM)
        dtype=np.uint8,               # uint8 suffit pour 0-100 → 4× moins de RAM
    )

    # Pour chaque mot, trouver le meilleur match
    best_idx   = score_matrix.argmax(axis=1)
    best_score = score_matrix[np.arange(len(eligible_texts)), best_idx]

    # Appliquer la correction uniquement si ≥ threshold
    corrections: Dict[int, str] = {}
    eli_idx = 0
    for i, ok in enumerate(eligible_mask):
        if not ok:
            continue
        score = int(best_score[eli_idx])
        if score >= threshold:
            candidate     = choices[best_idx[eli_idx]]
            original_text = texts[i]
            # Vérification supplémentaire : longueur similaire (±40 %)
            # évite de remplacer "HT" → "Montant" si la similarity est élevée par hasard
            len_ratio = len(original_text) / max(len(candidate), 1)
            if 0.6 <= len_ratio <= 1.6:
                corrections[i] = candidate
        eli_idx += 1

    # Construire la série de sortie
    result = series.copy()
    for i, corrected in corrections.items():
        result.iloc[i] = corrected

    return result


# ─────────────────────────────────────────────────────────────────────────────
# Utilitaire : résumé des corrections appliquées
# ─────────────────────────────────────────────────────────────────────────────

def correction_report(words_cleaned: List[dict]) -> pd.DataFrame:
    """
    Retourne un DataFrame listant uniquement les mots qui ont été modifiés.

    Utile pour déboguer ou afficher dans l'interface Streamlit.

    Args:
        words_cleaned : Sortie de clean_ocr_output().

    Returns:
        pd.DataFrame avec colonnes : original_text, text, correction, box.
    """
    df = pd.DataFrame(words_cleaned)
    if "original_text" not in df.columns:
        return pd.DataFrame()

    changed = df[df["original_text"] != df["text"]].copy()
    return changed[["original_text", "text", "correction", "box"]].reset_index(drop=True)


# ─────────────────────────────────────────────────────────────────────────────
# Fonctions utilitaires exportées
# ─────────────────────────────────────────────────────────────────────────────

def normalize_arabic(text: str, remove_diacritics: bool = True) -> str:
    """
    Normalise un texte arabe (version non-vectorisée pour usage unitaire).

    Args:
        text             : Chaîne arabe brute.
        remove_diacritics: Si True, supprime harakat et shadda.

    Returns:
        Chaîne normalisée.
    """
    text = _RE_ALIF.sub("ا", text)
    text = _RE_YA.sub("ي", text)
    text = _RE_KAF.sub("ك", text)
    text = _RE_TATWEEL.sub("", text)
    if remove_diacritics:
        text = _RE_DIACRITICS.sub("", text)
    return text


def sanitize_amount(text: str) -> str:
    """
    Nettoie un montant isolé (version non-vectorisée).

    Args:
        text : Chaîne représentant un montant avec éventuelles confusions.

    Returns:
        Chaîne avec confusions corrigées.

    Exemple :
        sanitize_amount("1.2OO,OO") → "1.200,00"
    """
    for wrong, right in _CONFUSION_MAP.items():
        text = text.replace(wrong, right)
    return text