"""Hand seeds for the canonical T-Rex 250-point CoTracker layout.

No appearance auto-repair. Tune prompts in ``sam2_prompt_hands.py``.
"""

from __future__ import annotations

import numpy as np

from trex_track.layout import NUM_HEAD_PER_HAND

NUM_HAND_POINTS = NUM_HEAD_PER_HAND

__all__ = [
    "NUM_HAND_POINTS",
    "wrist_hand_50",
    "head_hands_50",
    # Compatibility aliases for callers of the old 25-point implementation.
    "wrist_hand_25",
    "head_hands_25",
]


def _stable_spatial_order(points_xy: np.ndarray) -> np.ndarray:
    """Give sampled slots deterministic top-to-bottom, then left-to-right IDs."""

    points = np.asarray(points_xy, dtype=np.float32).reshape(-1, 2)
    order = np.lexsort((points[:, 0], points[:, 1]))
    return points[order]


def wrist_hand_50(
    predictor,
    rgb_hwc: np.ndarray,
    view: str,
    *,
    n_points: int = NUM_HAND_POINTS,
    seed: int | None = 0,
) -> tuple[np.ndarray, np.ndarray, str]:
    """One SAM2 pass → 50 stably ordered queries on a wrist hand mask.

    Background 5×5 grid is NOT produced here — callers prepend ``make_image_grid``.
    """
    from trex_track.sam2_prompt_hands import sample_points_on_mask, sam2_wrist_mask

    mask, tag, _meta = sam2_wrist_mask(predictor, rgb_hwc, view)
    rng = np.random.default_rng(seed)
    pts = sample_points_on_mask(mask, n_points, rng=rng, distal_bias=True)
    h, w = rgb_hwc.shape[:2]
    pts[:, 0] = np.clip(pts[:, 0], 0, w - 1)
    pts[:, 1] = np.clip(pts[:, 1], 0, h - 1)
    return _stable_spatial_order(pts), mask.astype(bool), tag


def head_hands_50(
    predictor,
    rgb_hwc: np.ndarray,
    *,
    n_points: int = NUM_HAND_POINTS,
    seed: int | None = 0,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, str]:
    """One SAM2 pass per hand → 100 queries (left 50, then right 50)."""
    from trex_track.sam2_prompt_hands import sample_points_on_mask, sam2_head_hand_mask

    rng = np.random.default_rng(seed)
    left, tag_l, _ = sam2_head_hand_mask(predictor, rgb_hwc, "left")
    right, tag_r, _ = sam2_head_hand_mask(predictor, rgb_hwc, "right")
    pts_l = _stable_spatial_order(
        sample_points_on_mask(left, n_points, rng=rng, distal_bias=True)
    )
    pts_r = _stable_spatial_order(
        sample_points_on_mask(right, n_points, rng=rng, distal_bias=True)
    )
    h, w = rgb_hwc.shape[:2]
    queries = np.concatenate([pts_l, pts_r], axis=0).astype(np.float32)
    queries[:, 0] = np.clip(queries[:, 0], 0, w - 1)
    queries[:, 1] = np.clip(queries[:, 1], 0, h - 1)
    tag = f"{tag_l}|{tag_r}"
    return queries, left.astype(bool), right.astype(bool), tag


# Keep old imports working while making their new 50-point behavior explicit.
wrist_hand_25 = wrist_hand_50
head_hands_25 = head_hands_50