# Copyright 2026 AutomatosX
# SPDX-License-Identifier: AGPL-3.0-only
# Modified 2026-10-19 for v1.3; prior component notices are in NOTICE.
"""Public ``rgb-letterbox-1024/1`` geometry. Independent of any detector graph.

This is the signed AXDocLayout-M v1.3 preprocess: HWC uint8 RGB, bilinear resize of the
longer side to 1024, white pad, odd remainder bottom/right, then inverse
mapping of letterbox-space boxes back onto the original raster.

It is not a DocLayout-YOLO drop-in. That incumbent uses a different colour
order, a different pad value, and a different class vocabulary. Those
differences are left unmapped on purpose.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from PIL import Image

from ax_doclayout.contract import PREPROCESS_ID, PageImage
from ax_doclayout.errors import InvalidPageImageError

SIZE = 1024
PAD_RGB = (255, 255, 255)
FILL_SCALE = 1.0 / 255.0


@dataclass(frozen=True, slots=True)
class Letterbox:
    """One 1024-square canvas plus the exact inverse-mapping terms."""

    canvas: np.ndarray
    scale: float
    pad_left: int
    pad_top: int
    pad_right: int
    pad_bottom: int
    size: int = SIZE

    def __post_init__(self) -> None:
        canvas = np.ascontiguousarray(self.canvas)
        if canvas.ndim != 3 or canvas.shape != (self.size, self.size, 3):
            raise InvalidPageImageError("letterbox canvas must be size x size x 3")
        if canvas.dtype != np.uint8:
            raise InvalidPageImageError("letterbox canvas must be uint8")
        object.__setattr__(self, "canvas", canvas)


def letterbox_rgb(rgb: np.ndarray, *, size: int = SIZE) -> Letterbox:
    """Letterbox a decoded HWC uint8 RGB array under ``rgb-letterbox-1024/1``."""

    if not isinstance(rgb, np.ndarray):
        raise InvalidPageImageError("rgb must be a numpy array")
    if rgb.ndim != 3 or rgb.shape[2] != 3:
        raise InvalidPageImageError("rgb must be HWC with three channels")
    if rgb.dtype != np.uint8:
        raise InvalidPageImageError("rgb must be uint8")
    height, width = (int(rgb.shape[0]), int(rgb.shape[1]))
    if height < 1 or width < 1:
        raise InvalidPageImageError("rgb must not be empty")
    scale = size / max(width, height)
    new_width = max(1, round(width * scale))
    new_height = max(1, round(height * scale))
    image = Image.fromarray(np.ascontiguousarray(rgb), mode="RGB")
    resized = image.resize((new_width, new_height), Image.BILINEAR)
    pad_width = size - new_width
    pad_height = size - new_height
    pad_left = pad_width // 2
    pad_right = pad_width - pad_left
    pad_top = pad_height // 2
    pad_bottom = pad_height - pad_top
    if pad_right < 0 or pad_bottom < 0:
        raise InvalidPageImageError("letterbox overflow")
    canvas = Image.new("RGB", (size, size), PAD_RGB)
    canvas.paste(resized, (pad_left, pad_top))
    return Letterbox(
        canvas=np.asarray(canvas, dtype=np.uint8),
        scale=float(scale),
        pad_left=int(pad_left),
        pad_top=int(pad_top),
        pad_right=int(pad_right),
        pad_bottom=int(pad_bottom),
        size=int(size),
    )


def letterbox_page(page: PageImage, *, size: int = SIZE) -> Letterbox:
    """Letterbox a validated ``PageImage``."""

    return letterbox_rgb(page.rgb, size=size)


def letterbox_nchw(letterbox: Letterbox) -> np.ndarray:
    """Return a batch-1 NCHW float32 tensor in ``[0, 1]``."""

    chw = np.transpose(letterbox.canvas, (2, 0, 1)).astype(np.float32)
    chw *= np.float32(FILL_SCALE)
    return np.ascontiguousarray(chw[None, ...])


def inverse_letterbox_box(
    box: tuple[float, float, float, float] | list[float],
    letterbox: Letterbox,
    *,
    width_px: int,
    height_px: int,
) -> tuple[float, float, float, float] | None:
    """Map a letterbox xyxy box onto the original raster, or drop it."""

    if width_px < 1 or height_px < 1:
        return None
    try:
        x_min, y_min, x_max, y_max = (float(value) for value in box)
    except (TypeError, ValueError):
        return None
    if not all(np.isfinite(value) for value in (x_min, y_min, x_max, y_max)):
        return None
    scale = letterbox.scale
    if scale <= 0.0:
        return None
    mapped = (
        (x_min - letterbox.pad_left) / scale,
        (y_min - letterbox.pad_top) / scale,
        (x_max - letterbox.pad_left) / scale,
        (y_max - letterbox.pad_top) / scale,
    )
    clipped_x_min = min(max(mapped[0], 0.0), float(width_px))
    clipped_y_min = min(max(mapped[1], 0.0), float(height_px))
    clipped_x_max = min(max(mapped[2], 0.0), float(width_px))
    clipped_y_max = min(max(mapped[3], 0.0), float(height_px))
    if clipped_x_min >= clipped_x_max or clipped_y_min >= clipped_y_max:
        return None
    return clipped_x_min, clipped_y_min, clipped_x_max, clipped_y_max


def orig_target_sizes(letterbox: Letterbox) -> np.ndarray:
    """Sizes fed to the exported graph. The graph is bound to the letterbox square."""

    side = np.int64(letterbox.size)
    return np.ascontiguousarray([[side, side]])


PREPROCESS = PREPROCESS_ID