"""Image contracts: uint8 HWC BGR input; float32 RGB/CHW model input."""
from __future__ import annotations

from typing import cast

import cv2
import numpy as np
from numpy.typing import NDArray

ByteImage = NDArray[np.uint8]
FloatImage = NDArray[np.float32]

# region pixels
def pixel_demo() -> tuple[ByteImage, ByteImage, ByteImage]:
    """Return a (2,2,3) BGR image, its RGB copy, and a copied one-pixel crop."""
    bgr: ByteImage = np.array(
        [[[0, 0, 255], [0, 255, 0]], [[255, 0, 0], [30, 60, 90]]],
        dtype=np.uint8,
    )
    rgb: ByteImage = cast(ByteImage, cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB))
    crop: ByteImage = rgb[1:2, 1:2, :].copy()
    crop[0, 0, 0] = 0
    return bgr, rgb, crop
# endregion pixels

# region normalize
def to_unit_rgb(bgr: ByteImage) -> FloatImage:
    """Nonempty uint8 (H,W,3) BGR -> float32 (H,W,3) RGB in [0,1]."""
    if bgr.dtype != np.uint8 or bgr.ndim != 3 or bgr.shape[2] != 3:
        raise ValueError("Expected uint8 HWC with three BGR channels")
    if min(bgr.shape[:2]) == 0:
        raise ValueError("Image must have pixels")
    rgb: ByteImage = cast(ByteImage, cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB))
    unit_rgb: FloatImage = rgb.astype(np.float32) / np.float32(255.0)
    return unit_rgb
# endregion normalize

# region layout
def model_input(unit_rgb: FloatImage) -> FloatImage:
    """Finite float32 HWC RGB [0,1] -> contiguous NCHW batch of size one."""
    if unit_rgb.dtype != np.float32 or unit_rgb.ndim != 3 or unit_rgb.shape[2] != 3:
        raise ValueError("Expected float32 HWC with three RGB channels")
    if unit_rgb.size == 0 or not np.isfinite(unit_rgb).all():
        raise ValueError("Expected nonempty finite image")
    if np.any(unit_rgb < 0) or np.any(unit_rgb > 1):
        raise ValueError("Expected values in [0,1]")
    chw: FloatImage = unit_rgb.transpose(2, 0, 1)
    batch: FloatImage = np.ascontiguousarray(chw[None, ...])
    return batch
# endregion layout