src/preprocessing/image_preprocessor.py

"""
Pipeline de prétraitement d'images pour OCR multi-langue.

Étapes : PDF/Image → Normalisation → Niveaux de gris →
          Débruitage → Deskewing → Binarisation adaptative

Optimisé pour documents contenant du texte arabe cursif,
français et anglais — photos de qualité variable acceptées.
"""

import io
import logging
from pathlib import Path
from typing import Union, Optional

import cv2
import numpy as np
from PIL import Image

logger = logging.getLogger(__name__)


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

def preprocess_for_ocr(
    file_input,
    dpi: int = 300,
    page_index: int = 0,
    denoise: bool = True,
    deskew: bool = True,
    binarize: bool = True,
    return_pil: bool = True,
) -> Union[Image.Image, np.ndarray]:
    """
    Pipeline complet de prétraitement : charge le fichier puis applique
    niveaux de gris → débruitage → correction d'inclinaison → binarisation.

    Args:
        file_input  : UploadedFile Streamlit, chemin (str/Path), bytes ou PIL.Image.
        dpi         : Résolution pour la conversion PDF (300 recommandé pour l'OCR).
        page_index  : Numéro de page à traiter (PDF multi-pages).
        denoise     : Activer la réduction de bruit.
        deskew      : Activer la correction automatique d'inclinaison.
        binarize    : Activer la binarisation adaptative.
        return_pil  : True → retourne PIL.Image ; False → numpy uint8.

    Returns:
        Image prétraitée prête pour docTR / PaddleOCR / EasyOCR.
    """
    # ── Chargement ────────────────────────────────────────────────────────────
    pages = _load_document(file_input, dpi=dpi)
    img   = pages[min(page_index, len(pages) - 1)].copy()

    # ── 1. Normalisation de la résolution (upscale si trop petite) ───────────
    img = _normalize_resolution(img, min_width=1200)

    # ── 2. Niveaux de gris ───────────────────────────────────────────────────
    gray = _to_grayscale(img)

    # ── 3. Débruitage (paramètres conservateurs pour l'arabe cursif) ─────────
    if denoise:
        gray = _denoise(gray)

    # ── 4. Correction d'inclinaison ──────────────────────────────────────────
    if deskew:
        gray = _deskew(gray)

    # ── 5. Binarisation adaptative ───────────────────────────────────────────
    result = _binarize(gray) if binarize else gray

    return _cv2_to_pil(result) if return_pil else result


def get_preview_pair(
    file_input,
    dpi: int = 300,
    page_index: int = 0,
) -> tuple:
    """
    Retourne (original_pil, preprocessed_pil) pour affichage côte à côte.

    Args:
        file_input  : Source du fichier.
        dpi         : Résolution DPI pour PDF.
        page_index  : Page à utiliser.

    Returns:
        (PIL.Image original, PIL.Image prétraitée)
    """
    pages = _load_document(file_input, dpi=dpi)
    raw   = pages[min(page_index, len(pages) - 1)]

    original_pil     = _cv2_to_pil(raw)
    preprocessed_pil = preprocess_for_ocr(
        file_input, dpi=dpi, page_index=page_index, return_pil=True
    )
    return original_pil, preprocessed_pil


# ─────────────────────────────────────────────────────────────────────────────
# Chargement multi-format
# ─────────────────────────────────────────────────────────────────────────────

def _load_document(file_input, dpi: int = 300) -> list:
    """
    Charge un document et retourne une liste de pages (numpy BGR).

    Formats supportés : JPEG, PNG, TIFF, BMP, PDF.
    """
    # PIL.Image direct
    if isinstance(file_input, Image.Image):
        return [_pil_to_cv2(file_input)]

    # numpy array direct
    if isinstance(file_input, np.ndarray):
        return [file_input]

    # Lire les bytes selon la source
    if hasattr(file_input, "read"):
        # UploadedFile Streamlit ou file-like object
        raw_bytes = file_input.read()
        filename  = getattr(file_input, "name", "file.bin")
        try:
            file_input.seek(0)   # Rembobiner pour usage ultérieur dans l'app
        except Exception:
            pass
    elif isinstance(file_input, (str, Path)):
        filename  = str(file_input)
        with open(file_input, "rb") as f:
            raw_bytes = f.read()
    elif isinstance(file_input, bytes):
        raw_bytes = file_input
        filename  = "file.bin"
    else:
        raise TypeError(f"Type d'entrée non supporté : {type(file_input)}")

    ext = Path(filename).suffix.lower()

    # ── PDF → images haute résolution ─────────────────────────────────────────
    if ext == ".pdf":
        try:
            from pdf2image import convert_from_bytes
            pil_pages = convert_from_bytes(raw_bytes, dpi=dpi)
            return [_pil_to_cv2(p) for p in pil_pages]
        except ImportError:
            raise ImportError(
                "pdf2image est requis pour les PDF.\n"
                "Installez-le avec : pip install pdf2image\n"
                "Sous Windows, téléchargez aussi Poppler : "
                "https://github.com/oschwartz10612/poppler-windows/releases"
            )

    # ── Image standard (JPEG, PNG, TIFF, BMP…) ────────────────────────────────
    arr = np.frombuffer(raw_bytes, dtype=np.uint8)
    img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
    if img is None:
        raise ValueError(
            f"Impossible de décoder l'image '{filename}'. "
            "Vérifiez que le fichier n'est pas corrompu."
        )
    return [img]


# ─────────────────────────────────────────────────────────────────────────────
# Étapes du pipeline
# ─────────────────────────────────────────────────────────────────────────────

def _normalize_resolution(
    img: np.ndarray,
    min_width: int = 1400,
    max_width: int = 2400,
) -> np.ndarray:
    """
    Maintient la largeur de l'image dans [min_width, max_width].

    - Upscale si trop petite (< min_width) : améliore la précision OCR.
      · LANCZOS4 si scale > 1.5 (meilleure conservation des arêtes de chars)
      · CUBIC sinon (plus rapide, suffisant pour facteur faible)
    - Downscale si trop grande (> max_width) : réduit la RAM et le temps
      de traitement sur CPU — INTER_AREA est optimal pour réduire.

    Seuils relevés à 1400/2400 px (vs 1000/1800) pour réduire le CER
    sur les factures avec petits caractères (numéros de ligne, montants).
    Consommation mémoire : ≈ 280 Mo max sur CPU 8 Go.
    """
    h, w = img.shape[:2]

    if w < min_width:
        scale  = min_width / w
        interp = cv2.INTER_LANCZOS4 if scale > 1.5 else cv2.INTER_CUBIC
        img    = cv2.resize(
            img,
            (int(w * scale), int(h * scale)),
            interpolation=interp,
        )
    elif w > max_width:
        scale = max_width / w
        img   = cv2.resize(
            img,
            (int(w * scale), int(h * scale)),
            interpolation=cv2.INTER_AREA,
        )
    return img


def _to_grayscale(img: np.ndarray) -> np.ndarray:
    """Convertit BGR → niveaux de gris (no-op si déjà en gris)."""
    if len(img.shape) == 2:
        return img
    return cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)


def _denoise(gray: np.ndarray) -> np.ndarray:
    """
    Réduction du bruit par filtre Non-Local Means.

    Paramètres conservateurs (h=7) pour préserver :
    - Les connexions entre lettres arabes (cursif)
    - Les diacritiques (harakat, shadda, tanwin…)
    - Les petits détails des caractères français accentués

    Un h trop élevé (>15) estompe ces détails et dégrade l'OCR arabe.
    """
    return cv2.fastNlMeansDenoising(
        gray,
        h=7,                  # Force du filtre — conservateur pour l'arabe
        templateWindowSize=7,
        searchWindowSize=21,
    )


def _deskew(gray: np.ndarray) -> np.ndarray:
    """
    Correction automatique de l'inclinaison du document.

    Algorithme :
    1. Binarisation Otsu temporaire (rapide)
    2. Détection des lignes de texte via Hough probabiliste
    3. Angle médian des lignes quasi-horizontales
    4. Rotation de correction si l'angle est > 0.3° (seuil de tolérance)

    Fonctionne pour l'arabe car les lignes de base sont horizontales
    quelle que soit la direction d'écriture (RTL).
    """
    angle = _estimate_skew_angle(gray)

    if abs(angle) < 0.3:
        return gray   # Inclinaison négligeable → pas de traitement

    logger.debug("Deskew : correction de %.2f°", angle)
    return _rotate_image(gray, angle)


def _estimate_skew_angle(gray: np.ndarray) -> float:
    """
    Estime l'angle d'inclinaison dominant (en degrés).

    Retourne 0.0 si aucune ligne exploitable n'est détectée.
    """
    # Binarisation grossière pour la détection de contours
    _, binary = cv2.threshold(
        gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU
    )

    edges = cv2.Canny(binary, 50, 150, apertureSize=3)

    # Longueur minimale de ligne = 25 % de la largeur (ignore le bruit)
    min_len = gray.shape[1] // 4

    lines = cv2.HoughLinesP(
        edges,
        rho=1,
        theta=np.pi / 180,
        threshold=80,
        minLineLength=min_len,
        maxLineGap=20,
    )

    if lines is None or len(lines) == 0:
        return 0.0

    angles = []
    for line in lines:
        x1, y1, x2, y2 = line[0]
        dx = x2 - x1
        if dx == 0:
            continue
        angle = np.degrees(np.arctan2(y2 - y1, dx))
        # Garder uniquement les lignes proches de l'horizontale (±15°)
        if -15.0 < angle < 15.0:
            angles.append(angle)

    if not angles:
        return 0.0

    # Médiane pour être robuste face aux lignes aberrantes
    return float(np.median(angles))


def _rotate_image(gray: np.ndarray, angle: float) -> np.ndarray:
    """
    Rotation autour du centre de l'image, fond blanc.
    Les nouvelles dimensions sont recalculées pour éviter le rognage.
    """
    h, w  = gray.shape[:2]
    cx, cy = w // 2, h // 2

    M     = cv2.getRotationMatrix2D((cx, cy), angle, scale=1.0)
    cos_a = abs(M[0, 0])
    sin_a = abs(M[0, 1])

    new_w = int(h * sin_a + w * cos_a)
    new_h = int(h * cos_a + w * sin_a)

    # Ajuster la translation pour centrer l'image dans les nouvelles dimensions
    M[0, 2] += (new_w - w) / 2
    M[1, 2] += (new_h - h) / 2

    return cv2.warpAffine(
        gray, M, (new_w, new_h),
        flags=cv2.INTER_CUBIC,
        borderMode=cv2.BORDER_CONSTANT,
        borderValue=255,   # Fond blanc
    )


def _binarize(gray: np.ndarray) -> np.ndarray:
    """
    Binarisation adaptative gaussienne.

    Avantages sur Otsu global :
    - Gère les ombres portées et l'éclairage non uniforme (photos de documents)
    - Préserve les zones de texte foncé sur fond clair ET clair sur fond foncé
    - blockSize=31 : fenêtre suffisamment grande pour les ombres diffuses
    - C=11 : soustraction de constante locale pour affiner le seuil

    Pour l'arabe :
    - La binarisation adaptative respecte mieux les diacritiques (points,
      shadda, tanwin) que la binarisation globale qui les efface parfois.
    - Pré-lissage léger (GaussianBlur 3×3) pour éviter l'amplification
      du bruit dans les zones uniformes avant seuillage.
    """
    # Lissage préalable pour réduire le bruit haute fréquence
    blurred = cv2.GaussianBlur(gray, (3, 3), 0)

    binary = cv2.adaptiveThreshold(
        blurred,
        maxValue=255,
        adaptiveMethod=cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        thresholdType=cv2.THRESH_BINARY,
        blockSize=31,
        C=11,
    )
    return binary


# ─────────────────────────────────────────────────────────────────────────────
# Utilitaires de conversion PIL ↔ OpenCV
# ─────────────────────────────────────────────────────────────────────────────

def _pil_to_cv2(pil_img: Image.Image) -> np.ndarray:
    """PIL.Image (RGB) → numpy BGR pour OpenCV."""
    rgb = np.array(pil_img.convert("RGB"))
    return cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)


def _cv2_to_pil(img: np.ndarray) -> Image.Image:
    """numpy BGR ou niveaux de gris → PIL.Image."""
    if len(img.shape) == 2:
        return Image.fromarray(img, mode="L")
    rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    return Image.fromarray(rgb)