"""
skew_corrector.py
-----------------
Detects and corrects rotational skew in scanned/photographed floor plans.

Strategy:
  1. Detect dominant lines using Probabilistic Hough Transform.
  2. Compute the median angle of near-horizontal and near-vertical lines.
  3. Rotate the image to align the dominant axis.

Floor plans are axis-aligned by design, so the dominant line angle
should be close to 0° or 90°.
"""

import cv2
import numpy as np
from typing import Optional


def detect_skew_angle(binary: np.ndarray, angle_threshold: float = 45.0) -> float:
    """
    Estimate the skew angle of the image using the Hough Line Transform.

    Args:
        binary:           Binary image (uint8, 0/255).
        angle_threshold:  Only consider lines within this many degrees of
                          horizontal (0°) or vertical (90°).

    Returns:
        Estimated skew angle in degrees. Positive = clockwise skew.
        Returns 0.0 if no dominant angle is found.
    """
    # Hough works on edges — use Canny on the binary image
    edges = cv2.Canny(binary, threshold1=50, threshold2=150, apertureSize=3)

    # Probabilistic Hough: faster and more robust for long wall lines
    lines = cv2.HoughLinesP(
        edges,
        rho=1,
        theta=np.pi / 180,
        threshold=80,
        minLineLength=binary.shape[1] // 8,  # at least 1/8 of image width
        maxLineGap=20,
    )

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

    angles = []
    for line in lines:
        x1, y1, x2, y2 = line[0]
        if x2 == x1:
            continue  # vertical line → 90°, handle separately
        angle = np.degrees(np.arctan2(y2 - y1, x2 - x1))

        # Normalise to [-90, 90]
        if angle > 90:
            angle -= 180
        elif angle < -90:
            angle += 180

        # Keep only near-horizontal lines (close to 0°)
        if abs(angle) <= angle_threshold:
            angles.append(angle)

    if not angles:
        return 0.0

    # Use the median to be robust against outliers
    return float(np.median(angles))


def correct_skew(
    img: np.ndarray,
    angle: Optional[float] = None,
    binary: Optional[np.ndarray] = None,
    background_color: int = 255,
) -> np.ndarray:
    """
    Rotate an image to correct for skew.

    Args:
        img:              Grayscale image to correct.
        angle:            Skew angle in degrees (provide this OR binary).
        binary:           Binary image used to auto-detect angle if angle=None.
        background_color: Fill value for borders created by rotation (0=black, 255=white).

    Returns:
        Rotated (deskewed) grayscale image, same size as input.
    """
    if img is None:
        raise ValueError("img must not be None.")

    if angle is None:
        if binary is None:
            raise ValueError("Provide either 'angle' or 'binary' for auto-detection.")
        angle = detect_skew_angle(binary)

    # If skew is negligible, skip rotation to avoid resampling artifacts
    if abs(angle) < 0.3:
        return img.copy()

    h, w = img.shape[:2]
    center = (w / 2.0, h / 2.0)

    # Rotation matrix — negate angle because OpenCV y-axis is flipped
    M = cv2.getRotationMatrix2D(center, -angle, scale=1.0)

    rotated = cv2.warpAffine(
        img,
        M,
        (w, h),
        flags=cv2.INTER_LINEAR,
        borderMode=cv2.BORDER_CONSTANT,
        borderValue=background_color,
    )

    return rotated


def deskew(img: np.ndarray, binary: np.ndarray) -> tuple[np.ndarray, float]:
    """
    Convenience wrapper: detect angle and correct in one call.

    Args:
        img:    Original grayscale image.
        binary: Binary version used for angle detection.

    Returns:
        (corrected_image, detected_angle_degrees)
    """
    angle = detect_skew_angle(binary)
    corrected = correct_skew(img, angle=angle)
    return corrected, angle