"""Coordinate and prompt transforms in pure numpy."""

import numpy as np

from sam2_onnx.types import Float32Array, HW, Int64Array, PromptCoords


def transform_coords(
    coords: Float32Array,
    orig_hw: HW,
    resolution: int,
) -> Float32Array:
    """Transform pixel coordinates to model coordinate space.

    Args:
        coords: (N, 2) array of (x, y) pixel coordinates.
        orig_hw: Original image (height, width).
        resolution: Model input resolution.

    Returns:
        Transformed (N, 2) coordinates scaled to model resolution.
    """
    h, w = orig_hw
    if h <= 0 or w <= 0:
        raise ValueError(f"orig_hw dimensions must be positive, got {orig_hw}")
    scale = np.array([resolution / w, resolution / h], dtype=np.float32)
    return coords.astype(np.float32) * scale


def transform_boxes(
    boxes: Float32Array,
    orig_hw: HW,
    resolution: int,
) -> Float32Array:
    """Transform XYXY boxes to model coordinate space.

    Args:
        boxes: (B, 4) array of [x0, y0, x1, y1] pixel coordinates.
        orig_hw: Original image (height, width).
        resolution: Model input resolution.

    Returns:
        Transformed (B, 4) boxes scaled to model resolution.
    """
    h, w = orig_hw
    if h <= 0 or w <= 0:
        raise ValueError(f"orig_hw dimensions must be positive, got {orig_hw}")
    scale = np.array([resolution / w, resolution / h, resolution / w, resolution / h], dtype=np.float32)
    return boxes.astype(np.float32) * scale


def concat_points(
    point_coords: Float32Array | None,
    point_labels: Int64Array | None,
    box: Float32Array | None,
    orig_hw: HW,
    resolution: int,
) -> PromptCoords:
    """Combine points and box into a single prompt, transformed to model space.

    Matches SAM2ImagePredictor._predict: box corners (labels [2, 3]) are prepended
    before point prompts. No padding point is added — the PromptEncoder inside the
    ONNX decoder handles that.

    Args:
        point_coords: (N, 2) point coordinates or None.
        point_labels: (N,) point labels or None.
        box: (4,) XYXY box or None.
        orig_hw: Original image (height, width).
        resolution: Model input resolution.

    Returns:
        Tuple of (coords (1, P, 2), labels (1, P)) ready for the decoder.
    """
    all_coords: list[Float32Array] = []
    all_labels: list[Float32Array] = []

    # Box corners come first (matching SAM2ImagePredictor._predict)
    if box is not None:
        box_t = transform_boxes(box.reshape(1, 4), orig_hw, resolution).reshape(2, 2)
        all_coords.append(box_t)
        all_labels.append(np.array([2.0, 3.0], dtype=np.float32))

    if point_coords is not None and point_labels is not None:
        transformed = transform_coords(point_coords, orig_hw, resolution)
        all_coords.append(transformed)
        all_labels.append(point_labels.astype(np.float32))

    if not all_coords:
        raise ValueError("At least one of point_coords or box must be provided")

    coords = np.concatenate(all_coords, axis=0)
    labels = np.concatenate(all_labels, axis=0)  # pragma: no mutate

    # Add batch dimension
    return coords[np.newaxis], labels[np.newaxis]
