src/ocr/doctr_engine.py

"""
Moteur OCR docTR — Optimisé CPU 8 Go RAM, Multilingue (FR / EN / AR).

Modèles choisis pour CPU sans GPU :
    Détection   : db_mobilenet_v3_large  (~12 Mo,  ~1-2 s/page)
    Reconnaissance : crnn_mobilenet_v3_large  (~8 Mo, ~2-3 s/page)

    → Consommation mémoire totale : ≈ 350-450 Mo RAM
    → Temps traitement page A4 scannée : 3-6 s sur CPU standard

    Alternatives si encore trop lent :
        reco_arch = "crnn_mobilenet_v3_small"  → ≈ 180 Mo, ~2 s/page

Sortie : [{"text": str, "box": [x1,y1,x2,y2], "confidence": float, "is_arabic": bool}]
         Coordonnées normalisées 0-1000 (format LayoutLM / LayoutXLM).
"""

import gc
import logging
import os
from typing import Union, Optional

import numpy as np
from PIL import Image, ImageDraw, ImageFont

logger = logging.getLogger(__name__)

# ─────────────────────────────────────────────────────────────────────────────
# Constantes d'optimisation CPU
# ─────────────────────────────────────────────────────────────────────────────

# Largeur max envoyée à docTR pour l'inférence.
# Au-delà, l'image est redimensionnée → gain de vitesse important.
# 1024 px = bon compromis vitesse/précision sur CPU.
MAX_INFERENCE_WIDTH = 512

# Architectures recommandées pour CPU 8 Go
_DEFAULT_DET  = "db_mobilenet_v3_large"    # léger, rapide
_DEFAULT_RECO = "crnn_mobilenet_v3_large"  # bon équilibre précision/vitesse

# Singleton : le prédicteur est chargé une seule fois par session Streamlit
_predictor = None


# ─────────────────────────────────────────────────────────────────────────────
# Initialisation du prédicteur (singleton)
# ─────────────────────────────────────────────────────────────────────────────

def get_predictor(
    det_arch:  str = _DEFAULT_DET,
    reco_arch: str = _DEFAULT_RECO,
    assume_straight_pages: bool = True,
):
    """
    Charge et met en cache le prédicteur docTR (une seule fois par session).

    Optimisations CPU appliquées automatiquement :
        - Tous les threads disponibles (torch.set_num_threads)
        - Mode inférence uniquement (pas de gradient calculé)
        - Modèles légers MobileNetV3

    Args:
        det_arch  : Architecture de détection.
                    'db_mobilenet_v3_large'  ← recommandé CPU  (12 Mo)
                    'db_resnet50'            ← précis mais lent sur CPU
        reco_arch : Architecture de reconnaissance.
                    'crnn_mobilenet_v3_large' ← recommandé CPU  (8 Mo)
                    'crnn_mobilenet_v3_small' ← ultra-léger     (3 Mo)
                    'crnn_vgg16_bn'           ← moyen            (60 Mo)
                    'parseq'                  ← précis mais lent (90 Mo)
        assume_straight_pages : True pour pages droites (plus rapide).

    Returns:
        doctr.models.OCRPredictor prêt à l'emploi.
    """
    global _predictor

    if _predictor is not None:
        return _predictor

    # ── Import conditionnel (évite d'importer torch au démarrage) ─────────────
    try:
        import torch
        from doctr.models import ocr_predictor
    except ImportError:
        raise ImportError(
            "docTR n'est pas installé.\n"
            "Installez avec : pip install python-doctr[torch]\n"
            "Documentation  : https://mindee.github.io/doctr/"
        )

    # ── Libération mémoire avant chargement ──────────────────────────────────
    gc.collect()

    # ── Optimisations CPU ─────────────────────────────────────────────────────
    n_threads = min(os.cpu_count() or 2, 4)   # plafonné à 4 pour limiter la RAM
    torch.set_num_threads(n_threads)
    torch.set_num_interop_threads(max(1, n_threads // 2))

    logger.info(
        "Chargement docTR — det=%s  reco=%s  threads=%d",
        det_arch, reco_arch, n_threads,
    )

    _predictor = ocr_predictor(
        det_arch=det_arch,
        reco_arch=reco_arch,
        pretrained=True,
        assume_straight_pages=assume_straight_pages,
    )

    # Passer en mode évaluation (désactive dropout, batchnorm en mode train, etc.)
    _predictor.det_predictor.model.eval()
    _predictor.reco_predictor.model.eval()

    logger.info("Modèles docTR prêts (RAM ≈ 350-450 Mo).")
    return _predictor


def reset_predictor():
    """Libère les modèles de la mémoire (utile si on change d'architecture)."""
    global _predictor
    _predictor = None
    gc.collect()


def trim_memory():
    """
    Demande au système d'exploitation de récupérer les pages mémoire inutilisées
    par PyTorch après une inférence CPU. Sans effet sur GPU.

    Sur Windows : SetProcessWorkingSetSize(-1, -1) vide le working set.
    Sur Linux   : malloc_trim(0) via libc.
    """
    gc.collect()
    gc.collect()
    try:
        import sys
        if sys.platform == "win32":
            import ctypes
            ctypes.windll.kernel32.SetProcessWorkingSetSize(
                ctypes.windll.kernel32.GetCurrentProcess(), -1, -1
            )
        else:
            ctypes.CDLL("libc.so.6").malloc_trim(0)
    except Exception:
        pass


# ─────────────────────────────────────────────────────────────────────────────
# Fonction principale d'OCR
# ─────────────────────────────────────────────────────────────────────────────

def run_ocr(
    image_input: Union[np.ndarray, Image.Image],
    det_arch:  str   = _DEFAULT_DET,
    reco_arch: str   = _DEFAULT_RECO,
    assume_straight_pages: bool  = True,
    min_confidence:        float = 0.30,
) -> dict:
    """
    Exécute l'OCR docTR sur une image prétraitée.

    L'image est redimensionnée à MAX_INFERENCE_WIDTH avant inférence
    pour réduire la charge mémoire et accélérer le traitement sur CPU.

    Args:
        image_input     : numpy array (BGR ou gris) OU PIL.Image.
        det_arch        : Architecture de détection.
        reco_arch       : Architecture de reconnaissance.
        assume_straight_pages : Optimisation pages droites.
        min_confidence  : Mots en dessous de ce seuil sont ignorés.

    Returns:
        dict :
            "words"      → liste LayoutXLM (voir to_layoutxlm_format)
            "full_text"  → texte complet (str)
            "page_w/h"   → dimensions originales de la page
            "raw_result" → résultat brut docTR
    """
    import io as _io

    try:
        import torch
        ctx = torch.inference_mode()   # pas de calcul de gradient → économie RAM
    except ImportError:
        ctx = _NullContext()

    # ── Préparation image ─────────────────────────────────────────────────────
    pil_rgb = _to_rgb_pil(image_input)
    page_w, page_h = pil_rgb.size

    # Redimensionner pour l'inférence (gain de vitesse × 2-4 selon la taille)
    pil_inf = _resize_for_inference(pil_rgb, max_width=MAX_INFERENCE_WIDTH)

    # ── Inférence ─────────────────────────────────────────────────────────────
    predictor = get_predictor(det_arch, reco_arch, assume_straight_pages)

    try:
        from doctr.io import DocumentFile
        buf = _io.BytesIO()
        pil_inf.save(buf, format="PNG")
        buf.seek(0)
        doc = DocumentFile.from_images([buf.read()])
    except Exception:
        doc = [np.array(pil_inf)]

    with ctx:
        raw_result = predictor(doc)

    # Libérer immédiatement le document (tenseurs d'entrée)
    del doc
    gc.collect()

    # ── Extraction ────────────────────────────────────────────────────────────
    words     = to_layoutxlm_format(raw_result, min_confidence=min_confidence)
    full_text = _reconstruct_text(raw_result)

    # Ne pas faire del raw_result ici : Python 3.12 lève UnboundLocalError si
    # predictor() a échoué auparavant. raw_result est libéré automatiquement
    # par le GC au retour de la fonction (plus aucune référence externe).
    del pil_inf
    gc.collect()
    gc.collect()   # double appel : libère les cycles du premier passage

    return {
        "words":     words,
        "full_text": full_text,
        "page_w":    page_w,
        "page_h":    page_h,
    }


# ─────────────────────────────────────────────────────────────────────────────
# Conversion vers format LayoutXLM (0-1000)
# ─────────────────────────────────────────────────────────────────────────────

def to_layoutxlm_format(
    doctr_result,
    min_confidence: float = 0.30,
) -> list:
    """
    Convertit la sortie docTR en liste de dicts compatibles LayoutXLM.

    Format de sortie :
        [
            {
                "text":       "Facture",
                "box":        [42, 80, 198, 112],  # 0-1000
                "confidence": 0.97,
                "is_arabic":  False,
            },
            ...
        ]

    Normalisation 0-1000 (standard LayoutLM v1/v2/XLM) :
        Les coordonnées relatives [0,1] de docTR sont multipliées par 1000.

    Args:
        doctr_result   : Objet Document retourné par le prédicteur docTR.
        min_confidence : Seuil de confiance minimal.

    Returns:
        Liste de dicts triés de haut en bas puis gauche à droite.
    """
    words_out = []

    for page in doctr_result.pages:
        for block in page.blocks:
            for line in block.lines:
                for word in line.words:

                    if word.confidence < min_confidence:
                        continue

                    geo = word.geometry

                    # Extraire les coins (format bbox 2 pts ou polygone 4 pts)
                    if len(geo) == 2:
                        (rx1, ry1), (rx2, ry2) = geo
                    else:
                        xs = [p[0] for p in geo]
                        ys = [p[1] for p in geo]
                        rx1, ry1, rx2, ry2 = min(xs), min(ys), max(xs), max(ys)

                    # Normalisation 0-1000 avec clamp de sécurité
                    x1 = max(0, min(1000, int(rx1 * 1000)))
                    y1 = max(0, min(1000, int(ry1 * 1000)))
                    x2 = max(0, min(1000, int(rx2 * 1000)))
                    y2 = max(0, min(1000, int(ry2 * 1000)))

                    text = word.value.strip()
                    if not text:
                        continue

                    words_out.append({
                        "text":       text,
                        "box":        [x1, y1, x2, y2],
                        "confidence": round(word.confidence, 4),
                        "is_arabic":  _is_arabic(text),
                    })

    # Tri : lecture naturelle haut → bas, gauche → droite
    words_out.sort(key=lambda w: (w["box"][1], w["box"][0]))
    return words_out


# ─────────────────────────────────────────────────────────────────────────────
# Fusion de mots en blocs de lignes (Grouping pour LayoutXLM)
# ─────────────────────────────────────────────────────────────────────────────

def merge_words_into_lines(
    words: list,
    y_thresh: int = 15,
    x_gap_thresh: int = 40,
) -> list:
    """
    Regroupe les mots qui appartiennent à la même ligne visuelle en blocs fusionnés.

    Pourquoi ?
        docTR renvoie les mots un par un (ex : une désignation → 15 tokens isolés).
        LayoutXLM fonctionne mieux sur des blocs de ligne cohérents car :
          - Chaque bloc consomme 1 slot sur les 512 disponibles (au lieu de 15).
          - Le contexte spatial est plus lisible pour le modèle.

    Algorithme (3 étapes) :
        1. Tri par Y_centre puis X  →  ordre de lecture naturel.
        2. Regroupement vertical   :  deux mots forment la même ligne si
                                      |Y_centre_A − Y_centre_B| ≤ y_thresh.
        3. Fusion horizontale      :  deux mots adjacents sur la même ligne
                                      sont fusionnés si le gap horizontal
                                      (x1_droit − x2_gauche) ≤ x_gap_thresh.
                                      Au-delà du seuil, une nouvelle colonne commence.

    Paramètres :
        y_thresh      (int, 0-1000) : Tolérance verticale en unités normalisées.
                                      Recommandé : 10-20  (1-2 % de la hauteur).
                                      • 15 → absorbe la légère inclinaison / skew.
                                      • Trop grand (>30) → deux lignes adjacentes
                                        pourraient être fusionnées par erreur.
        x_gap_thresh  (int, 0-1000) : Distance horizontale max entre deux mots
                                      pour les fusionner.
                                      Recommandé : 30-60  (3-6 % de la largeur).
                                      • 40 → fusionne les mots d'une même phrase.
                                      • Garde les gaps inter-colonnes (>60) séparés.

    Args:
        words : Liste [{"text": str, "box": [x1,y1,x2,y2], "confidence": float,
                        "is_arabic": bool}] — sortie de to_layoutxlm_format().

    Returns:
        Liste de blocs fusionnés, même format que words, triée en ordre de lecture.
    """
    if not words:
        return words

    # ── Étape 1 : tri initial par Y_centre puis X (ordre lecture naturel) ─────
    sorted_words = sorted(
        words,
        key=lambda w: ((w["box"][1] + w["box"][3]) / 2, w["box"][0]),
    )

    # ── Étape 2 : regroupement vertical en lignes ─────────────────────────────
    # Chaque élément de `lines` est une liste de mots sur la même ligne.
    lines: list[list[dict]] = []

    for word in sorted_words:
        y_c = (word["box"][1] + word["box"][3]) / 2

        placed = False
        # Chercher une ligne existante compatible (parcours en sens inverse
        # pour comparer à la ligne la plus récente en priorité)
        for line in reversed(lines):
            line_y_c = sum((w["box"][1] + w["box"][3]) / 2 for w in line) / len(line)
            if abs(y_c - line_y_c) <= y_thresh:
                line.append(word)
                placed = True
                break

        if not placed:
            lines.append([word])

    # ── Étape 3 : tri gauche→droite dans chaque ligne ─────────────────────────
    for line in lines:
        line.sort(key=lambda w: w["box"][0])

    # ── Étape 4 : fusion horizontale par ligne ────────────────────────────────
    merged: list[dict] = []

    for line in lines:
        group: dict | None = None

        for word in line:
            if group is None:
                group = {
                    "text":       word["text"],
                    "box":        list(word["box"]),      # copie mutable
                    "confidence": word["confidence"],
                    "is_arabic":  word.get("is_arabic", False),
                }
            else:
                # Gap horizontal : distance entre la fin du groupe et le début du mot
                gap = word["box"][0] - group["box"][2]

                if gap <= x_gap_thresh:
                    # ── Fusion : concaténer le texte et étendre la bbox ───────
                    group["text"] += " " + word["text"]
                    # Bbox englobante : min des coins haut-gauche, max des coins bas-droit
                    group["box"][0] = min(group["box"][0], word["box"][0])
                    group["box"][1] = min(group["box"][1], word["box"][1])
                    group["box"][2] = max(group["box"][2], word["box"][2])
                    group["box"][3] = max(group["box"][3], word["box"][3])
                    # Confiance = pire confiance du groupe (prudence)
                    group["confidence"] = min(group["confidence"], word["confidence"])
                    group["is_arabic"]  = group["is_arabic"] or word.get("is_arabic", False)
                else:
                    # ── Nouveau bloc (colonne différente ou espacement trop grand)
                    merged.append(group)
                    group = {
                        "text":       word["text"],
                        "box":        list(word["box"]),
                        "confidence": word["confidence"],
                        "is_arabic":  word.get("is_arabic", False),
                    }

        if group is not None:
            merged.append(group)

    # ── Étape 5 : tri final en ordre de lecture (haut→bas, gauche→droite) ─────
    merged.sort(key=lambda b: (b["box"][1], b["box"][0]))

    logger.debug(
        "merge_words_into_lines : %d mots → %d blocs (y_thresh=%d, x_gap=%d)",
        len(words), len(merged), y_thresh, x_gap_thresh,
    )
    return merged


# ─────────────────────────────────────────────────────────────────────────────
# Visualisation — bounding boxes annotées
# ─────────────────────────────────────────────────────────────────────────────

def draw_bounding_boxes(
    image_input: Union[np.ndarray, Image.Image],
    words: list,
    color_latin:     str  = "#0088DD",
    color_arabic:    str  = "#E05A00",
    line_width:      int  = 2,
    show_text:       bool = True,
    show_confidence: bool = False,
) -> Image.Image:
    """
    Dessine les bounding boxes sur l'image pour vérification visuelle.

    Couleurs :
        Bleu   (#0088DD) → texte latin (français / anglais)
        Orange (#E05A00) → texte arabe (détecté automatiquement)

    Args:
        image_input      : Image source (PIL ou numpy BGR/gris).
        words            : Sortie de to_layoutxlm_format().
        color_latin      : Couleur bbox texte latin.
        color_arabic     : Couleur bbox texte arabe.
        line_width       : Épaisseur contour (pixels).
        show_text        : Afficher le texte reconnu au-dessus de la bbox.
        show_confidence  : Ajouter le score de confiance dans le label.

    Returns:
        PIL.Image RGB annotée.
    """
    pil   = _to_rgb_pil(image_input).copy()
    draw  = ImageDraw.Draw(pil)
    w_img, h_img = pil.size

    font_size = max(11, w_img // 90)
    try:
        font = ImageFont.truetype("arial.ttf", size=font_size)
    except IOError:
        font = ImageFont.load_default()

    for word in words:
        x1, y1, x2, y2 = word["box"]

        # Dénormalisation 0-1000 → pixels
        px1 = int(x1 * w_img / 1000)
        py1 = int(y1 * h_img / 1000)
        px2 = int(x2 * w_img / 1000)
        py2 = int(y2 * h_img / 1000)

        color = color_arabic if word.get("is_arabic") else color_latin
        draw.rectangle([px1, py1, px2, py2], outline=color, width=line_width)

        if show_text:
            label = word["text"]
            if show_confidence:
                label += f" {word['confidence']:.2f}"
            try:
                bbox_t = draw.textbbox((px1, py1), label, font=font)
                draw.rectangle(
                    [bbox_t[0]-2, bbox_t[1]-2, bbox_t[2]+2, bbox_t[3]+2],
                    fill=color,
                )
                draw.text((px1, py1), label, fill="white", font=font, anchor="lb")
            except Exception:
                pass   # Si la police ne supporte pas le glyphe arabe, on ignore

    return pil


# ─────────────────────────────────────────────────────────────────────────────
# Pipeline complet (prétraitement → OCR → LayoutXLM → visualisation)
# ─────────────────────────────────────────────────────────────────────────────

def ocr_pipeline(
    file_input,
    language_hint:  Optional[str] = None,
    min_confidence: float = 0.30,
    apply_cleaning: bool  = True,
) -> dict:
    """
    Pipeline tout-en-un optimisé CPU 8 Go :
        1. Prétraitement (image_preprocessor)
        2. OCR léger docTR (MobileNetV3)
        3. Formatage LayoutXLM
        4. Visualisation bounding boxes
        5. Auto-correction OCR (ocr_cleaner — optionnel)

    Exemple d'appel depuis Streamlit :

        from ocr.doctr_engine import ocr_pipeline

        result = ocr_pipeline(uploaded_file, language_hint="ar")
        words  = result["words"]          # liste LayoutXLM (corrigée)
        text   = result["full_text"]      # texte brut
        viz    = result["visualization"]  # PIL.Image annotée
        report = result["correction_report"]  # DataFrame des corrections

    Args:
        file_input      : UploadedFile Streamlit, chemin, bytes ou PIL.Image.
        language_hint   : Code ISO optionnel ('fr', 'en', 'ar').
        min_confidence  : Seuil de confiance minimal.
        apply_cleaning  : Active l'auto-correction post-OCR (défaut : True).

    Returns:
        dict : words, full_text, page_w, page_h, visualization, preprocessed,
               correction_report (DataFrame ou None).
    """
    from preprocessing.image_preprocessor import preprocess_for_ocr

    # ── 1. Prétraitement ──────────────────────────────────────────────────────
    preprocessed_pil = preprocess_for_ocr(
        file_input,
        dpi=200,          # 200 DPI suffisant (300 trop lourd sur CPU)
        denoise=True,
        deskew=True,
        binarize=True,
        return_pil=True,
    )

    # ── 2. OCR ───────────────────────────────────────────────────────────────
    ocr_result = run_ocr(
        preprocessed_pil,
        det_arch=_DEFAULT_DET,
        reco_arch=_DEFAULT_RECO,
        min_confidence=min_confidence,
    )

    # ── 3. Visualisation ─────────────────────────────────────────────────────
    viz = draw_bounding_boxes(
        preprocessed_pil,
        ocr_result["words"],
        show_text=True,
        show_confidence=False,
    )

    # ── 4. Auto-correction post-OCR ──────────────────────────────────────────
    words_final    = ocr_result["words"]
    corr_report    = None
    cleaning_error = None

    if apply_cleaning and words_final:
        try:
            import importlib
            _cleaner = importlib.import_module("ocr.ocr_cleaner")
            words_final = _cleaner.clean_ocr_output(words_final)
            corr_report = _cleaner.correction_report(words_final)
            logger.info(
                "Auto-correction : %d mot(s) corrigé(s) sur %d.",
                len(corr_report), len(words_final),
            )
        except Exception as exc:
            cleaning_error = str(exc)
            logger.warning("Erreur auto-correction : %s", exc)

    # ── 5. Post-traitement métier (vocabulaire factures, mots fusionnés…) ────
    post_text              = None
    post_corr_report       = None
    post_processing_error  = None

    if words_final:
        try:
            import importlib
            _pp      = importlib.import_module("ocr.post_processor")
            words_final    = _pp.clean_invoice_ocr(words_final)
            lines          = _pp.reconstruct_lines_table_aware(words_final)
            col_bounds     = _pp.detect_column_boundaries(words_final)
            post_text      = _pp.lines_to_text(
                lines,
                col_bounds if len(col_bounds) >= 1 else None,
            )
            post_corr_report = _pp.correction_report_invoice(words_final)
            logger.info(
                "Post-traitement : %d correction(s), %d ligne(s), %d colonne(s).",
                len(post_corr_report), len(lines), len(col_bounds) + 1,
            )
        except Exception as exc:
            post_processing_error = str(exc)
            logger.warning("Erreur post-traitement : %s", exc)

    # Nettoyage mémoire
    gc.collect()

    return {
        "words":                words_final,
        "full_text":            ocr_result["full_text"],
        "page_w":               ocr_result["page_w"],
        "page_h":               ocr_result["page_h"],
        "visualization":        viz,
        "preprocessed":         preprocessed_pil,
        "correction_report":    corr_report,
        "cleaning_error":       cleaning_error,
        "post_text":            post_text,
        "post_corr_report":     post_corr_report,
        "post_processing_error": post_processing_error,
    }


# ─────────────────────────────────────────────────────────────────────────────
# Utilitaires internes
# ─────────────────────────────────────────────────────────────────────────────

def _resize_for_inference(pil: Image.Image, max_width: int) -> Image.Image:
    """
    Réduit l'image à max_width si elle est plus large.
    Conserve le ratio. Utilise LANCZOS (meilleur pour le texte).
    """
    w, h = pil.size
    if w <= max_width:
        return pil
    new_h = int(h * max_width / w)
    return pil.resize((max_width, new_h), Image.LANCZOS)


def _to_rgb_pil(image_input) -> Image.Image:
    """Convertit numpy (BGR/gris) ou PIL en PIL.Image RGB."""
    if isinstance(image_input, Image.Image):
        return image_input.convert("RGB")

    import cv2
    arr = np.asarray(image_input)
    if arr.ndim == 2:
        arr_rgb = cv2.cvtColor(arr, cv2.COLOR_GRAY2RGB)
    elif arr.shape[2] == 4:
        arr_rgb = cv2.cvtColor(arr, cv2.COLOR_BGRA2RGB)
    else:
        arr_rgb = cv2.cvtColor(arr, cv2.COLOR_BGR2RGB)
    return Image.fromarray(arr_rgb)


def _is_arabic(text: str) -> bool:
    """Détecte si un mot contient majoritairement des caractères arabes (U+0600-U+06FF)."""
    if not text:
        return False
    arabic = sum(1 for ch in text if "\u0600" <= ch <= "\u06FF")
    return arabic / len(text) > 0.4


def _reconstruct_text(doctr_result) -> str:
    """Reconstruit le texte complet : mots → lignes → blocs."""
    lines_out = []
    for page in doctr_result.pages:
        for block in page.blocks:
            for line in block.lines:
                words = [w.value for w in line.words if w.value.strip()]
                if words:
                    lines_out.append(" ".join(words))
    return "\n".join(lines_out)


class _NullContext:
    """Context manager vide si torch n'est pas disponible."""
    def __enter__(self): return self
    def __exit__(self, *_): pass