# SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
# SPDX-License-Identifier: Apache-2.0
"""Image preprocessing for the HaMeR / HaMeR regressor input.

Mirrors the upstream ``ViTDetDataset`` + ``HAMER.forward`` behaviour:

  1. Around the supplied hand bbox, expand by ``rescale_factor`` (default 2.0)
     and take the larger side → a square bbox of ``bbox_size`` pixels in
     original-image coords centred on the bbox.
  2. Crop that square out of the image, resize to ``256 × 256``.
  3. Centre-slice ``[..., 32:-32]`` along the W axis → ``256 × 192``.  This
     matches HaMeR's ``self.backbone(x[:,:,:,32:-32])``.
  4. ImageNet-normalize.
"""
from __future__ import annotations

import os
import threading
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Optional, Tuple

import numpy as np
import torch
from PIL import Image


IMAGE_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
IMAGE_STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
IMAGE_HEIGHT = 256          # ViT input H after the [:,:,:,32:-32] slice
IMAGE_WIDTH = 192           # ViT input W after the [:,:,:,32:-32] slice
BBOX_CROP_SIZE = 256        # square the bbox is resized to before W-slice
W_SLICE = 32                # pixels removed from each side of W
# round 13: the 192 x 256 resize runs as RESIZE_BANDS horizontal bands of output rows, the first on the calling thread and
# the others on a small thread pool (PIL releases the GIL while resampling).  Bit-identical to the one-call resize: the
# crop side S is an integer, so every band's box edge r * S / 256 and PIL's per-row filter centres box0 + (y + 0.5) * scale
# are exact dyadic floats -> the same filter windows and fixed-point weights as the full resize (checked on 1000 random
# crops x 2 band counts, tests/test_fused_host.py).  NOT the default: a standalone loop gains 0.3-0.5 ms (2 bands), but
# inside the server (300-request A/B on chip 5) the preprocess p50 moved 2.03 / 2.01 -> 1.87 / 2.33 ms (2 bands) and
# 2.41 ms (3 bands) -- the pool threads wake cold on the loaded host.  HAMER_PRE_BANDS=1 (default): one call.
RESIZE_BANDS = max(1, int(os.environ.get("HAMER_PRE_BANDS", "1")))
_POOL: Optional[ThreadPoolExecutor] = None
_POOL_LOCK = threading.Lock()


# round 14 (P1): native crop + resize (hamer/native, bit-identical to PIL, ~3x faster); HAMER_PRE_NATIVE=0 -> PIL path.
# Not used together with HAMER_PRE_BANDS > 1 (that knob measures the PIL band path).
def native_enabled() -> bool:
    return RESIZE_BANDS == 1 and os.environ.get("HAMER_PRE_NATIVE", "1").strip().lower() not in ("0", "false", "no", "off")


def _pool(n: int) -> ThreadPoolExecutor:
    global _POOL
    with _POOL_LOCK:
        if _POOL is None or _POOL._max_workers < n:  # noqa: SLF001
            _POOL = ThreadPoolExecutor(max_workers=n, thread_name_prefix="hamer-resize")
        return _POOL


def resize_crop_hw(crop: Image.Image, is_right: bool = True, bands: Optional[int] = None) -> np.ndarray:
    """The square crop -> uint8 HWC ``[256, 192, 3]``: the centre 192 of the 256 resized columns (BILINEAR), mirrored
    for a left hand.  ``bands`` > 1 splits the output rows over threads (bit-identical, see RESIZE_BANDS)."""
    S = crop.width
    x0, x1 = W_SLICE * S / BBOX_CROP_SIZE, (BBOX_CROP_SIZE - W_SLICE) * S / BBOX_CROP_SIZE
    nb = RESIZE_BANDS if bands is None else max(1, int(bands))
    if nb == 1:
        c = crop.resize((IMAGE_WIDTH, BBOX_CROP_SIZE), Image.BILINEAR, box=(x0, 0, x1, S))
        if not is_right:
            c = c.transpose(Image.FLIP_LEFT_RIGHT)
        return np.asarray(c).copy()
    out = np.empty((BBOX_CROP_SIZE, IMAGE_WIDTH, 3), np.uint8)
    edges = [round(i * BBOX_CROP_SIZE / nb) for i in range(nb + 1)]

    def band(i: int) -> None:
        r0, r1 = edges[i], edges[i + 1]
        b = crop.resize((IMAGE_WIDTH, r1 - r0), Image.BILINEAR,
                        box=(x0, r0 * S / BBOX_CROP_SIZE, x1, r1 * S / BBOX_CROP_SIZE))
        if not is_right:
            b = b.transpose(Image.FLIP_LEFT_RIGHT)
        out[r0:r1] = np.asarray(b)

    futs = [_pool(nb - 1).submit(band, i) for i in range(1, nb)]
    band(0)
    for f in futs:
        f.result()
    return out


def hamer_bbox_size(bbox_xyxy: Tuple[float, float, float, float], rescale_factor: float = 2.0) -> float:
    """Mirrors HaMeR's ``ViTDetDataset`` scale computation:

        scale     = rescale_factor * (bbox_w, bbox_h) / 200
        bbox_size = (scale * 200).max()

    With ``BBOX_SHAPE=None`` (the released checkpoint's default), no
    aspect-ratio expansion is applied — the square bbox just fits the
    longer side.
    """
    x1, y1, x2, y2 = bbox_xyxy
    w, h = x2 - x1, y2 - y1
    return float(rescale_factor * max(w, h))


def crop_bbox(
    image: Image.Image,
    bbox_xyxy: Tuple[float, float, float, float],
    rescale_factor: float = 2.0,
) -> Tuple[Image.Image, Tuple[float, float, float]]:
    """Crop the HaMeR-style square hand bbox out of ``image`` and return the
    256×256 PIL crop plus ``(cx, cy, bbox_size)`` in *original-image coords*.
    """
    crop, meta = _square_crop(image, bbox_xyxy, rescale_factor)
    return crop.resize((BBOX_CROP_SIZE, BBOX_CROP_SIZE), Image.BILINEAR), meta


def _square_crop(
    image: Image.Image,
    bbox_xyxy: Tuple[float, float, float, float],
    rescale_factor: float = 2.0,
) -> Tuple[Image.Image, Tuple[float, float, float]]:
    """The HaMeR-style square crop (grey padding) at source resolution, plus ``(cx, cy, bbox_size)``."""
    x1, y1, x2, y2 = bbox_xyxy
    cx, cy = (x1 + x2) / 2.0, (y1 + y2) / 2.0
    bbox_size = hamer_bbox_size(bbox_xyxy, rescale_factor)
    half = bbox_size / 2.0
    crop_box = (cx - half, cy - half, cx + half, cy + half)
    # Pad with grey if the bbox extends beyond the image edges.
    crop = Image.new("RGB", (int(round(bbox_size)), int(round(bbox_size))), (128, 128, 128))
    src = image.crop((max(0, int(round(crop_box[0]))), max(0, int(round(crop_box[1]))),
                      min(image.width, int(round(crop_box[2]))), min(image.height, int(round(crop_box[3])))))
    paste_x = max(0, int(round(-crop_box[0])))
    paste_y = max(0, int(round(-crop_box[1])))
    crop.paste(src, (paste_x, paste_y))
    return crop, (cx, cy, bbox_size)


def preprocess(
    image_path: str | Path,
    bbox_xyxy: Optional[Tuple[float, float, float, float]] = None,
    rescale_factor: float = 2.0,
    is_right: bool = True,
    raw: bool = False,
) -> Tuple[torch.Tensor, Optional[Tuple[float, float, float]]]:
    """Load image, crop HaMeR-style around ``bbox_xyxy``, return
    ``(input_tensor, bbox_meta)``.

    ``bbox_xyxy`` is in original-image pixel coords.  When omitted, the
    image is centre-cropped to a square — useful for smoke-tests, but
    geometrically meaningless for the camera prediction.

    ``is_right=False`` mirrors the crop horizontally before normalisation —
    HaMeR is trained on right hands only, so a left hand is fed in
    flipped form (matches ``ViTDetDataset``'s ``flip = (right == 0)``).

    ``bbox_meta = (cx, cy, bbox_size)`` lets the visualiser reverse the
    crop to overlay the mesh back onto the original frame.
    """
    img = Image.open(image_path).convert("RGB")
    return preprocess_pil(img, bbox_xyxy, rescale_factor=rescale_factor, is_right=is_right, raw=raw)


def center_square_bbox(img: Image.Image) -> Tuple[float, float, float, float]:
    """The centre square of ``img`` as an xyxy bbox (the no-bbox fallback)."""
    side = min(img.width, img.height)
    return (
        (img.width - side) / 2.0,
        (img.height - side) / 2.0,
        (img.width + side) / 2.0,
        (img.height + side) / 2.0,
    )


def preprocess_pil(
    img: Image.Image,
    bbox_xyxy: Optional[Tuple[float, float, float, float]] = None,
    rescale_factor: float = 2.0,
    is_right: bool = True,
    raw: bool = False,
    contiguous: bool = True,
) -> Tuple[torch.Tensor, Tuple[float, float, float]]:
    """In-memory variant of :func:`preprocess` for an already-decoded PIL
    image (what a server holds).  Same crop / flip / slice / normalise
    sequence, byte-for-byte.  ``raw=True``: the uint8 ``[1, 3, 256, 192]`` crop before the
    normalisation (``normalize_raw`` of it is bit-identical to ``raw=False``); the
    ``TT_FUSED_INPUT=u8`` device graph takes it directly.  ``contiguous=False`` (raw only): return the
    ``[1, 3, 256, 192]`` view of the HWC array instead of a CHW copy (the server; same values).
    """
    if img.mode != "RGB":
        img = img.convert("RGB")
    if bbox_xyxy is None:
        bbox_xyxy = center_square_bbox(img)
    res = None
    if native_enabled():
        # round 14 (P1): crop + grey pad + box resize + mirror in one native call, bit-identical to the PIL path below
        # (hamer/native/resize.c; self-checked at load, tests/test_fused_host.py::test_native_crop_resize_bit_identical)
        from hamer import native

        # planar CHW straight from the native pass: the raw crop is contiguous [1, 3, 256, 192] at no extra cost (the
        # model's centring then reads it contiguously: 0.02 vs 0.07 ms for the strided HWC view)
        # round 15: the frame's pixel buffer is read in place (HAMER_PRE_ZEROCOPY=1; no np.asarray copy)
        res = native.crop_resize_image(img, bbox_xyxy, rescale_factor, is_right, chw=True)
    if res is not None:
        u8_chw, bbox_meta = res
        x = torch.from_numpy(u8_chw).unsqueeze(0)
        if raw:
            return x, bbox_meta
        return (x.float() / 255.0 - IMAGE_MEAN) / IMAGE_STD, bbox_meta
    crop, bbox_meta = _square_crop(img, bbox_xyxy, rescale_factor)
    # Only the centre 192 of the 256 resized columns are used: resize just that box (round 6: bit-identical to
    # resize-to-256 + slice -- PIL's filter windows depend only on the output pixel centre and the source image,
    # checked on 400 random crops -- and ~25% less host time).  The flip commutes with the symmetric slice.
    return _finish(torch.from_numpy(resize_crop_hw(crop, is_right)), bbox_meta, raw, contiguous)


def _finish(u8: torch.Tensor, bbox_meta, raw: bool, contiguous: bool):
    if raw:
        x = u8.permute(2, 0, 1).unsqueeze(0)
        # round 13 (server): contiguous=False keeps the view (no CHW copy); the model's input centring reads it strided
        return (x.contiguous() if contiguous else x), bbox_meta
    arr = u8.float() / 255.0
    arr = arr.permute(2, 0, 1).unsqueeze(0)                 # (1, 3, 256, 192)
    return (arr - IMAGE_MEAN) / IMAGE_STD, bbox_meta


def normalize_raw(u8: torch.Tensor) -> torch.Tensor:
    """uint8 ``[B, 3, 256, 192]`` crop (``preprocess_pil(raw=True)``) -> the normalized fp32 input,
    bit-identical to ``preprocess_pil(raw=False)``."""
    return (u8.float() / 255.0 - IMAGE_MEAN) / IMAGE_STD


def denormalize(x: torch.Tensor) -> np.ndarray:
    img = x * IMAGE_STD + IMAGE_MEAN
    img = img.squeeze(0).clamp(0, 1).permute(1, 2, 0).numpy()
    return (img * 255).astype(np.uint8)


def map_network_to_original(
    pts_192x256: np.ndarray,
    bbox_meta: Tuple[float, float, float],
) -> np.ndarray:
    """Map 2-D points in network-input pixel coords (192 × 256) back to
    original-image pixel coords.

    Inverse of the preprocess: shift x by ``W_SLICE`` to undo the slice,
    then scale by ``bbox_size / BBOX_CROP_SIZE`` and translate by the bbox
    top-left corner.
    """
    cx, cy, bbox_size = bbox_meta
    pts = pts_192x256.astype(np.float64).copy()
    pts[:, 0] += W_SLICE                                    # 192-wide → 256-wide
    pts *= bbox_size / BBOX_CROP_SIZE                       # 256×256 → bbox_size×bbox_size
    pts[:, 0] += cx - bbox_size / 2.0
    pts[:, 1] += cy - bbox_size / 2.0
    return pts