"""
preprocess.py — REAL OpenCV preprocessing (Knowledge Module 05 §2).

OCR/extraction accuracy is dominated by image quality, so a few cheap,
deterministic OpenCV operations often beat swapping models. These run offline on
any image (a numpy array). The highest-value step for documents is DESKEW.
"""
from __future__ import annotations

import cv2
import numpy as np


def to_grayscale(img: np.ndarray) -> np.ndarray:
    """Color (BGR) -> single intensity channel. Most document tasks don't need
    color and it's 3x less data."""
    if img.ndim == 3:
        return cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    return img


def denoise(gray: np.ndarray) -> np.ndarray:
    """Remove scanner/camera speckle that confuses OCR."""
    return cv2.fastNlMeansDenoising(gray, h=10)


def binarize(gray: np.ndarray) -> np.ndarray:
    """Adaptive threshold -> black text on white, robust to uneven lighting
    (phone photos). Returns a 0/255 image."""
    return cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                 cv2.THRESH_BINARY, 31, 10)


def estimate_skew_angle(gray: np.ndarray) -> float:
    """Estimate the document's skew in degrees via the minimum-area rectangle of
    the foreground (text) pixels. Positive = rotated counter-clockwise."""
    # Foreground = dark text on light bg -> invert so text is the non-zero mass.
    inv = cv2.bitwise_not(gray)
    _, bw = cv2.threshold(inv, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
    coords = np.column_stack(np.where(bw > 0))
    if len(coords) < 10:
        return 0.0
    angle = cv2.minAreaRect(coords.astype(np.float32))[-1]
    # OpenCV returns the angle in [-90, 0); normalize to a small skew around 0.
    if angle < -45:
        angle = 90 + angle
    return float(angle)


def _rotate(img: np.ndarray, degrees: float) -> np.ndarray:
    (h, w) = img.shape[:2]
    M = cv2.getRotationMatrix2D((w / 2, h / 2), degrees, 1.0)
    border = (255, 255, 255) if img.ndim == 3 else 255
    return cv2.warpAffine(img, M, (w, h), flags=cv2.INTER_CUBIC,
                          borderMode=cv2.BORDER_CONSTANT, borderValue=border)


def deskew(img: np.ndarray) -> tuple[np.ndarray, float]:
    """Detect skew and rotate the image to straighten it. Returns
    (deskewed_image, angle_applied). Crooked scans wreck OCR; this is the single
    highest-value preprocessing step. We try both sign conventions of the
    estimated angle and keep whichever leaves the smaller residual skew — robust
    to OpenCV's minAreaRect angle ambiguity."""
    gray = to_grayscale(img)
    angle = estimate_skew_angle(gray)
    best_img, best_corr, best_resid = img, 0.0, abs(angle)
    for corr in (-angle, angle):
        cand = _rotate(img, corr)
        resid = abs(estimate_skew_angle(to_grayscale(cand)))
        if resid < best_resid:
            best_img, best_corr, best_resid = cand, corr, resid
    return best_img, best_corr


def crop_to_document(img: np.ndarray) -> np.ndarray:
    """Find the largest 4-corner contour (the document edges in a photo) and
    perspective-correct it to a flat, cropped scan. Falls back to the original
    if no quadrilateral is found."""
    gray = to_grayscale(img)
    edges = cv2.Canny(gray, 50, 150)
    contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return img
    biggest = max(contours, key=cv2.contourArea)
    peri = cv2.arcLength(biggest, True)
    approx = cv2.approxPolyDP(biggest, 0.02 * peri, True)
    if len(approx) != 4:
        return img
    pts = approx.reshape(4, 2).astype(np.float32)
    rect = _order_corners(pts)
    (tl, tr, br, bl) = rect
    width = int(max(np.linalg.norm(br - bl), np.linalg.norm(tr - tl)))
    height = int(max(np.linalg.norm(tr - br), np.linalg.norm(tl - bl)))
    if width < 10 or height < 10:
        return img
    dst = np.array([[0, 0], [width - 1, 0], [width - 1, height - 1], [0, height - 1]],
                   dtype=np.float32)
    M = cv2.getPerspectiveTransform(rect, dst)
    return cv2.warpPerspective(img, M, (width, height))


def _order_corners(pts: np.ndarray) -> np.ndarray:
    """Order 4 points as top-left, top-right, bottom-right, bottom-left."""
    s = pts.sum(axis=1)
    diff = np.diff(pts, axis=1)
    return np.array([pts[np.argmin(s)], pts[np.argmin(diff)],
                     pts[np.argmax(s)], pts[np.argmax(diff)]], dtype=np.float32)


def preprocess_pipeline(img: np.ndarray) -> np.ndarray:
    """The standard chain: deskew -> grayscale -> denoise -> binarize. Returns a
    clean binary image ready for OCR."""
    deskewed, _ = deskew(img)
    gray = to_grayscale(deskewed)
    return binarize(denoise(gray))