ocr/preprocessing.py

"""
InvoiceForge AI — ocr/preprocessing.py
Advanced multi-pipeline image preprocessing for OCR accuracy maximisation.

Pipelines:
  A) Printed Documents  — CLAHE + NLMeans + Adaptive Threshold
  B) Handwritten Bills  — Bilateral Filter + Otsu + Stroke Dilation
  C) Mobile/Low-Res     — Super-scale + Shadow Removal + Perspective Correction
  D) Thermal Receipts   — Inversion + Gamma + Contrast Stretch

All functions accept BGR numpy arrays and return single-channel or BGR arrays
ready for PaddleOCR / EasyOCR inference.
"""

from __future__ import annotations

import logging
import math
from typing import Tuple

import cv2
import numpy as np
from PIL import Image

logger = logging.getLogger(__name__)

# ─────────────────────────────────────────────────────────────────────────────
# CONSTANTS
# ─────────────────────────────────────────────────────────────────────────────
MIN_PRINTED_DIM: int = 1200   # upscale target (px on longest side)
MIN_HW_DIM: int = 1600        # handwritten needs higher resolution
MIN_MOBILE_DIM: int = 1400
BLUR_LAPLACIAN_THRESHOLD: float = 80.0  # below → blurry image


# ─────────────────────────────────────────────────────────────────────────────
# UTILITY HELPERS
# ─────────────────────────────────────────────────────────────────────────────

def _to_gray(img: np.ndarray) -> np.ndarray:
    """Convert BGR or already-gray image to grayscale."""
    if len(img.shape) == 2:
        return img
    return cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)


def _upscale_if_small(gray: np.ndarray, min_dim: int) -> np.ndarray:
    """Upscale image if its longest side is below min_dim using INTER_CUBIC."""
    h, w = gray.shape[:2]
    if max(h, w) < min_dim:
        scale = min_dim / max(h, w)
        new_w, new_h = int(w * scale), int(h * scale)
        gray = cv2.resize(gray, (new_w, new_h), interpolation=cv2.INTER_CUBIC)
        logger.debug("Upscaled image %.1fx → (%d, %d)", scale, new_w, new_h)
    return gray


def _apply_clahe(gray: np.ndarray, clip: float = 2.0,
                 tile: Tuple[int, int] = (8, 8)) -> np.ndarray:
    """Apply Contrast-Limited Adaptive Histogram Equalisation."""
    clahe = cv2.createCLAHE(clipLimit=clip, tileGridSize=tile)
    return clahe.apply(gray)


def _gamma_correction(gray: np.ndarray, gamma: float = 1.5) -> np.ndarray:
    """Brighten/darken image via power-law gamma transform."""
    lut = np.array(
        [((i / 255.0) ** (1.0 / gamma)) * 255 for i in range(256)],
        dtype=np.uint8,
    )
    return cv2.LUT(gray, lut)


# ─────────────────────────────────────────────────────────────────────────────
# DESKEW
# ─────────────────────────────────────────────────────────────────────────────

def auto_deskew(gray: np.ndarray) -> np.ndarray:
    """
    Deskew image using Hough line detection.
    Corrects rotations between -45° and +45°.
    Returns the deskewed grayscale array.
    """
    edges = cv2.Canny(gray, 50, 150, apertureSize=3)
    lines = cv2.HoughLines(edges, 1, np.pi / 180, 200)
    if lines is None or len(lines) == 0:
        return gray

    angles: list[float] = []
    for rho, theta in lines[:30, 0]:
        angle = (theta - np.pi / 2) * 180 / np.pi
        if abs(angle) < 45:
            angles.append(angle)

    if not angles:
        return gray

    median_angle = float(np.median(angles))
    if abs(median_angle) <= 0.3:
        return gray   # tiny skew, not worth rotating

    h, w = gray.shape[:2]
    M = cv2.getRotationMatrix2D((w / 2, h / 2), median_angle, 1.0)
    deskewed = cv2.warpAffine(
        gray, M, (w, h),
        flags=cv2.INTER_CUBIC,
        borderMode=cv2.BORDER_REPLICATE,
    )
    logger.debug("Deskewed by %.2f°", median_angle)
    return deskewed


# ─────────────────────────────────────────────────────────────────────────────
# PERSPECTIVE CORRECTION
# ─────────────────────────────────────────────────────────────────────────────

def correct_perspective(img_bgr: np.ndarray) -> np.ndarray:
    """
    Detect the document quadrilateral in a camera-captured image and apply a
    four-point perspective warp to produce a flat, top-down view.

    If no clear quadrilateral is found the original image is returned unchanged.
    """
    gray = _to_gray(img_bgr)
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    edged = cv2.Canny(blurred, 75, 200)

    # Morphological closing to connect broken edges
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
    closed = cv2.morphologyEx(edged, cv2.MORPH_CLOSE, kernel)

    contours, _ = cv2.findContours(
        closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
    )
    if not contours:
        return img_bgr

    # Sort by area descending and try the largest contours
    contours = sorted(contours, key=cv2.contourArea, reverse=True)[:5]
    doc_contour = None
    for cnt in contours:
        peri = cv2.arcLength(cnt, True)
        approx = cv2.approxPolyDP(cnt, 0.02 * peri, True)
        if len(approx) == 4:
            doc_contour = approx
            break

    if doc_contour is None:
        return img_bgr

    pts = doc_contour.reshape(4, 2).astype(np.float32)
    # Order: top-left, top-right, bottom-right, bottom-left
    rect = _order_points(pts)
    (tl, tr, br, bl) = rect

    width_a = np.linalg.norm(br - bl)
    width_b = np.linalg.norm(tr - tl)
    max_w = max(int(width_a), int(width_b))

    height_a = np.linalg.norm(tr - br)
    height_b = np.linalg.norm(tl - bl)
    max_h = max(int(height_a), int(height_b))

    if max_w < 100 or max_h < 100:
        return img_bgr

    dst = np.array(
        [[0, 0], [max_w - 1, 0], [max_w - 1, max_h - 1], [0, max_h - 1]],
        dtype=np.float32,
    )
    M = cv2.getPerspectiveTransform(rect, dst)
    warped = cv2.warpPerspective(img_bgr, M, (max_w, max_h))
    logger.debug("Perspective corrected to %dx%d", max_w, max_h)
    return warped


def _order_points(pts: np.ndarray) -> np.ndarray:
    """Return points in (top-left, top-right, bottom-right, bottom-left) order."""
    rect = np.zeros((4, 2), dtype=np.float32)
    s = pts.sum(axis=1)
    rect[0] = pts[np.argmin(s)]   # top-left  → smallest sum
    rect[2] = pts[np.argmax(s)]   # bottom-right → largest sum
    diff = np.diff(pts, axis=1)
    rect[1] = pts[np.argmin(diff)]  # top-right → smallest diff
    rect[3] = pts[np.argmax(diff)]  # bottom-left → largest diff
    return rect


# ─────────────────────────────────────────────────────────────────────────────
# SHADOW REMOVAL
# ─────────────────────────────────────────────────────────────────────────────

def remove_shadow(img_bgr: np.ndarray) -> np.ndarray:
    """
    Remove uneven illumination / shadows using morphological background
    estimation and normalisation. Works on color images.
    """
    channels = cv2.split(img_bgr)
    result_channels = []
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (7, 7))

    for ch in channels:
        # Dilate then blur to estimate background illumination
        dilated = cv2.dilate(ch, kernel)
        bg = cv2.medianBlur(dilated, 21)
        diff = 255 - cv2.absdiff(ch, bg)
        normalized = cv2.normalize(
            diff, None, alpha=0, beta=255,
            norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8U,
        )
        result_channels.append(normalized)

    shadow_free = cv2.merge(result_channels)
    return shadow_free


# ─────────────────────────────────────────────────────────────────────────────
# IMAGE QUALITY METRICS
# ─────────────────────────────────────────────────────────────────────────────

def detect_blur(img: np.ndarray) -> float:
    """
    Measure image blurriness using the variance of the Laplacian.
    Returns a float; values below BLUR_LAPLACIAN_THRESHOLD indicate blur.
    """
    gray = _to_gray(img)
    lap_var = float(cv2.Laplacian(gray, cv2.CV_64F).var())
    return lap_var


def score_image_quality(img: np.ndarray) -> dict:
    """
    Return a composite quality dict with:
      blur_score    — higher is sharper  (>80 is acceptable)
      brightness    — mean pixel intensity (0-255)
      contrast      — std dev of pixel intensities
      resolution    — (w, h) tuple
      quality_score — aggregate 0.0-1.0 score
    """
    gray = _to_gray(img)
    blur = detect_blur(img)
    brightness = float(gray.mean())
    contrast = float(gray.std())
    h, w = gray.shape[:2]

    # Normalise each component to 0-1
    blur_norm = min(blur / 200.0, 1.0)
    bright_norm = 1.0 - abs(brightness - 128.0) / 128.0
    contrast_norm = min(contrast / 80.0, 1.0)
    res_norm = min((w * h) / (1920 * 1080), 1.0)

    quality = (
        0.40 * blur_norm
        + 0.25 * bright_norm
        + 0.20 * contrast_norm
        + 0.15 * res_norm
    )

    return {
        "blur_score": round(blur, 2),
        "brightness": round(brightness, 2),
        "contrast": round(contrast, 2),
        "resolution": (w, h),
        "quality_score": round(quality, 4),
        "is_blurry": blur < BLUR_LAPLACIAN_THRESHOLD,
    }


# ─────────────────────────────────────────────────────────────────────────────
# PIPELINE A — PRINTED DOCUMENTS
# ─────────────────────────────────────────────────────────────────────────────

def preprocess_printed(img_bgr: np.ndarray) -> np.ndarray:
    """
    Full preprocessing pipeline for printed / scanned invoices.

    Steps:
      1. Perspective correction (camera shots)
      2. Shadow removal
      3. Convert to gray
      4. Upscale to minimum resolution
      5. Auto-deskew
      6. CLAHE contrast enhancement
      7. Fast non-local means denoising
      8. Adaptive Gaussian threshold
      9. Morphological opening (noise spots)

    Returns binarised single-channel image.
    """
    # Step 1 — Perspective
    img_bgr = correct_perspective(img_bgr)

    # Step 2 — Shadow removal
    img_bgr = remove_shadow(img_bgr)

    # Step 3 — Grayscale
    gray = _to_gray(img_bgr)

    # Step 4 — Resolution
    gray = _upscale_if_small(gray, MIN_PRINTED_DIM)

    # Step 5 — Deskew
    gray = auto_deskew(gray)

    # Step 6 — CLAHE
    gray = _apply_clahe(gray, clip=2.0, tile=(8, 8))

    # Step 7 — Denoise
    denoised = cv2.fastNlMeansDenoising(gray, h=10, templateWindowSize=7,
                                         searchWindowSize=21)

    # Step 8 — Adaptive threshold
    thresh = cv2.adaptiveThreshold(
        denoised, 255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY,
        blockSize=15,
        C=8,
    )

    # Step 9 — Clean tiny specks
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 1))
    cleaned = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel)

    return cleaned


# ─────────────────────────────────────────────────────────────────────────────
# PIPELINE B — HANDWRITTEN DOCUMENTS
# ─────────────────────────────────────────────────────────────────────────────

def preprocess_handwritten(img_bgr: np.ndarray) -> np.ndarray:
    """
    Preprocessing pipeline optimised for handwritten bills and receipts.

    Steps:
      1. Convert to gray + upscale (higher target — 1600 px)
      2. Auto-deskew
      3. Bilateral filter (edge-preserving smoothing)
      4. CLAHE (higher clip limit for low-contrast handwriting)
      5. Otsu's binarisation (bimodal ink-on-paper histogram)
      6. Morphological dilation (thicken thin strokes)

    Returns binarised single-channel image.
    """
    # Step 1
    gray = _to_gray(img_bgr)
    gray = _upscale_if_small(gray, MIN_HW_DIM)

    # Step 2
    gray = auto_deskew(gray)

    # Step 3 — Bilateral preserves stroke edges
    filtered = cv2.bilateralFilter(gray, d=9, sigmaColor=75, sigmaSpace=75)

    # Step 4
    enhanced = _apply_clahe(filtered, clip=3.0, tile=(8, 8))

    # Step 5 — Otsu
    _, thresh = cv2.threshold(
        enhanced, 0, 255,
        cv2.THRESH_BINARY + cv2.THRESH_OTSU,
    )

    # Step 6 — Dilate strokes
    kernel = np.ones((2, 2), np.uint8)
    dilated = cv2.dilate(thresh, kernel, iterations=1)

    return dilated


# ─────────────────────────────────────────────────────────────────────────────
# PIPELINE C — MOBILE / LOW-RESOLUTION IMAGES
# ─────────────────────────────────────────────────────────────────────────────

def preprocess_mobile(img_bgr: np.ndarray) -> np.ndarray:
    """
    Pipeline for smartphone-captured invoice photos.
    Handles motion blur, uneven lighting, and perspective distortion.

    Steps:
      1. Perspective correction
      2. Shadow removal (heavy uneven lighting)
      3. Grayscale + aggressive upscale
      4. Wiener-like sharpening (unsharp mask)
      5. CLAHE
      6. Bilateral filter (preserves text edges)
      7. Adaptive threshold with larger block size
    """
    # Step 1 & 2
    img_bgr = correct_perspective(img_bgr)
    img_bgr = remove_shadow(img_bgr)

    gray = _to_gray(img_bgr)
    gray = _upscale_if_small(gray, MIN_MOBILE_DIM)

    # Step 4 — Unsharp mask for deblurring
    blurred = cv2.GaussianBlur(gray, (0, 0), 3)
    sharpened = cv2.addWeighted(gray, 1.5, blurred, -0.5, 0)

    # Step 5
    enhanced = _apply_clahe(sharpened, clip=3.5, tile=(8, 8))

    # Step 6
    filtered = cv2.bilateralFilter(enhanced, d=9, sigmaColor=75, sigmaSpace=75)

    # Step 7
    thresh = cv2.adaptiveThreshold(
        filtered, 255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY,
        blockSize=21,
        C=10,
    )

    return thresh


# ─────────────────────────────────────────────────────────────────────────────
# PIPELINE D — THERMAL RECEIPTS
# ─────────────────────────────────────────────────────────────────────────────

def preprocess_thermal(img_bgr: np.ndarray) -> np.ndarray:
    """
    Pipeline for thermal paper receipts (POS terminals).
    Thermal paper fades over time — needs aggressive contrast stretch.

    Steps:
      1. Grayscale
      2. Auto-invert if paper is dark (thermal negative)
      3. Gamma correction (brighten)
      4. CLAHE (high clip for low-contrast paper)
      5. Adaptive threshold
      6. Morphological dilation (thin thermal text)
    """
    gray = _to_gray(img_bgr)

    # Auto-invert detection: if mean > 128 the paper is likely light
    if gray.mean() < 128:
        gray = cv2.bitwise_not(gray)
        logger.debug("Thermal image auto-inverted.")

    # Gamma brighten
    gray = _gamma_correction(gray, gamma=1.8)

    # CLAHE high clip
    enhanced = _apply_clahe(gray, clip=4.0, tile=(4, 4))

    # Adaptive threshold
    thresh = cv2.adaptiveThreshold(
        enhanced, 255,
        cv2.ADAPTIVE_THRESH_MEAN_C,
        cv2.THRESH_BINARY,
        blockSize=11,
        C=6,
    )

    # Dilate to recover thin thermal characters
    kernel = np.ones((1, 1), np.uint8)
    result = cv2.dilate(thresh, kernel, iterations=1)

    return result


# ─────────────────────────────────────────────────────────────────────────────
# AUTO PIPELINE SELECTOR
# ─────────────────────────────────────────────────────────────────────────────

def auto_select_pipeline(img_bgr: np.ndarray, doc_type: str = "unknown") -> np.ndarray:
    """
    Automatically select and apply the appropriate preprocessing pipeline
    based on the document type string from the classifier.

    Args:
        img_bgr:  Input BGR image array.
        doc_type: One of the classifier output labels.

    Returns:
        Preprocessed image array.
    """
    doc_type = doc_type.lower()

    handwritten_types = {"handwritten_bill"}
    thermal_types = {"retail_receipt"}
    mobile_indicators = {"mobile"}

    if doc_type in handwritten_types:
        logger.info("Using HANDWRITTEN preprocessing pipeline.")
        return preprocess_handwritten(img_bgr)
    elif doc_type in thermal_types:
        logger.info("Using THERMAL preprocessing pipeline.")
        return preprocess_thermal(img_bgr)
    else:
        # Check blur — might be a mobile capture
        quality = score_image_quality(img_bgr)
        if quality["is_blurry"] or max(img_bgr.shape[:2]) < 800:
            logger.info("Using MOBILE preprocessing pipeline (low quality image).")
            return preprocess_mobile(img_bgr)

        logger.info("Using PRINTED preprocessing pipeline.")
        return preprocess_printed(img_bgr)