# Copyright 2026 Everpaper. Licensed under the Apache License, Version 2.0 (see ../LICENSE).
#
# Ports the logic of the Gemma 4 image processor (image_processing_gemma4.py in Hugging Face transformers 5.19,
# Copyright 2026 The HuggingFace Team, Apache-2.0), and the resize arithmetic of ATen's uint8
# `_upsample_bicubic2d_aa` kernel (PyTorch, BSD-3-Clause; itself modelled on Pillow's Resample.c, MIT-CMU/HPND).
# See ../NOTICE for the attributions and the PyTorch and Pillow license notices.
"""Exact host-side image preprocessing for the EmbeddingGemma 2 vision encoder (numpy + Pillow, no torch).

Reproduces what `Gemma4ImageProcessor` (transformers 5.19, torchvision backend) does to an image, bit for bit:

  1. decode to 8-bit RGB: PIL `convert("RGB")` (alpha dropped, not composited; palette expanded; EXIF orientation
     NOT applied; no colour management) - identical to torchvision `decode_image(..., mode=RGB)` for 8-bit files.
  2. target size: keep the aspect ratio, each side a multiple of 48 px, at most 2520 patches of 16x16 (budget of
     280 soft tokens); tiny images are upscaled (1x1 -> 768x768), extreme aspect ratios go up to 48x13440.
  3. resize: antialiased bicubic (a = -0.5) on uint8, the ATen `_upsample_bicubic2d_aa` uint8 kernel: separable,
     WIDTH pass first then HEIGHT, weights computed in float64 then rounded to int16 with a dynamic precision, each
     pass accumulated in int32, rounded (+half, arithmetic shift) and clamped to 0..255. Verified bit-identical to
     torchvision 0.29.1 on 239 images of 1x1 .. 8000x2000 px (tests/check_preprocess.py).
  4. 16x16x3 patches flattened in (py, px, c) order. Upstream then multiplies by 1/255 (no mean/std) and the
     encoder computes 2p - 1; here the patches stay INTEGER pixel values 0..255 (exact in float16) and the graph
     computes P - 127.5 with the 2/255 scale folded into its first projection (same maths, no input rounding).
     Upstream's pixel_values == patches * RESCALE.

The vision graph wants the patches in CELL-MAJOR order (the 9 patches of each 3x3 pooling cell contiguous, cells
row-major), padded to 2520, plus per-patch (x, y) positions, the axial-RoPE cos/sin table and a validity mask; see
`VisionInputs` / `prepare_vision_inputs`. The model is otherwise unchanged (see ../NOTICE).
"""
from __future__ import annotations

import math
from dataclasses import dataclass
from pathlib import Path

import numpy as np

PATCH = 16
POOL = 3
CELL_PX = PATCH * POOL  # 48
MAX_SOFT_TOKENS = 280
MAX_PATCHES = MAX_SOFT_TOKENS * POOL * POOL  # 2520
MAX_SIDE_PATCHES = MAX_SOFT_TOKENS * POOL  # 840: an image can be at most 840 patches wide or high (48 x 13440 px)
HEAD_DIM = 64
ROPE_THETA = 100.0
ROPE_FREQS = HEAD_DIM // 4  # 16 frequencies per axis


# ---------------------------------------------------------------------------------------------- decode
def load_rgb(image, apply_exif_orientation: bool = False) -> np.ndarray:
    """`image`: a path or a PIL image. Returns HWC uint8 RGB exactly like the upstream processor (8-bit files).

    apply_exif_orientation=True rotates camera photos upright first (NOT what upstream does; the embedding of a
    rotated photo differs from the upright one by cosine ~0.97, see README)."""
    from PIL import Image, ImageOps

    im = image if isinstance(image, Image.Image) else Image.open(Path(image))
    if apply_exif_orientation:
        im = ImageOps.exif_transpose(im)
    if im.mode in ("I;16", "I;16B", "I;16L", "I;16N", "I"):
        # 16-bit images: upstream divides the uint16 samples by 255 and produces garbage, PIL's convert clips at
        # 255; both runners take the high byte instead (documented divergence, same in Swift).
        a = np.asarray(im)
        a = (a >> 8).astype(np.uint8) if a.dtype == np.uint16 else np.clip(a // 256, 0, 255).astype(np.uint8)
        return np.repeat(a[..., None], 3, axis=-1)
    return np.asarray(im.convert("RGB"), dtype=np.uint8)


# ---------------------------------------------------------------------------------------------- size
def target_size(height: int, width: int, max_patches: int = MAX_PATCHES) -> tuple[int, int]:
    """Gemma4 `get_aspect_ratio_preserving_size`, verbatim logic. Returns (height, width)."""
    target_px = max_patches * PATCH * PATCH
    factor = math.sqrt(target_px / (height * width))
    th = int(math.floor(factor * height / CELL_PX)) * CELL_PX
    tw = int(math.floor(factor * width / CELL_PX)) * CELL_PX
    max_side = (max_patches // (POOL * POOL)) * CELL_PX
    if th == 0 and tw == 0:
        raise ValueError("image would resize to 0x0")
    if th == 0:
        th, tw = CELL_PX, min(int(math.floor(width / height)) * CELL_PX, max_side)
    elif tw == 0:
        tw, th = CELL_PX, min(int(math.floor(height / width)) * CELL_PX, max_side)
    if th * tw > target_px:
        raise ValueError("size exceeds the patch budget")
    return th, tw


# ---------------------------------------------------------------------------------------------- resize
def _cubic(x: float, a: float = -0.5) -> float:
    x = abs(x)
    if x < 1.0:
        return ((a + 2.0) * x - (a + 3.0)) * x * x + 1.0
    if x < 2.0:
        return (((x - 5.0) * x + 8.0) * x - 4.0) * a
    return 0.0


def aa_int16_weights(in_size: int, out_size: int) -> tuple[np.ndarray, np.ndarray, int]:
    """Antialiased bicubic taps for one axis, as ATen computes them for uint8 input.
    Returns (xmin [out] int, weights [out, taps] int32 (0-padded), precision)."""
    scale = in_size / out_size  # float64, align_corners=False, no explicit scale factor
    support = 2.0 * scale if scale >= 1.0 else 2.0
    invscale = 1.0 / scale if scale >= 1.0 else 1.0
    taps = int(math.ceil(support)) * 2 + 1
    xmin = np.zeros(out_size, dtype=np.int64)
    wts = np.zeros((out_size, taps), dtype=np.float64)
    for i in range(out_size):
        center = scale * (i + 0.5)
        lo = max(int(center - support + 0.5), 0)
        size = min(int(center + support + 0.5), in_size) - lo
        w = [_cubic((j + lo - center + 0.5) * invscale) for j in range(size)]
        tot = sum(w)
        if tot != 0.0:
            w = [v / tot for v in w]
        xmin[i] = lo
        wts[i, :size] = w
    wt_max = float(wts.max())
    precision = 0
    while precision < 22:
        if int(0.5 + wt_max * (1 << (precision + 1))) >= (1 << 15):
            break
        precision += 1
    s = float(1 << precision)
    # round half away from zero (what both std::round and Pillow do; all three rounding conventions agree on
    # every weight we have seen, this is the documented one)
    wi = np.where(wts >= 0, np.floor(wts * s + 0.5), np.ceil(wts * s - 0.5)).astype(np.int32)
    return xmin, wi, precision


def _resize_axis_u8(img: np.ndarray, axis: int, out_size: int) -> np.ndarray:
    in_size = img.shape[axis]
    xmin, wts, precision = aa_int16_weights(in_size, out_size)
    src = np.moveaxis(img, axis, 0).astype(np.int32)  # [in, ...]
    acc = np.full((out_size,) + src.shape[1:], 1 << (precision - 1), dtype=np.int32)
    taps = wts.shape[1]
    for t in range(taps):
        w = wts[:, t]
        if not w.any():
            continue
        idx = np.minimum(xmin + t, in_size - 1)  # taps beyond the range have weight 0
        acc += src[idx] * w.reshape((-1,) + (1,) * (src.ndim - 1))
    out = np.clip(acc >> precision, 0, 255).astype(np.uint8)
    return np.moveaxis(out, 0, axis)


def resize_bicubic_aa_u8(img: np.ndarray, height: int, width: int) -> np.ndarray:
    """img HWC uint8 -> [height, width, C] uint8, bit-identical to torchvision's uint8 antialiased bicubic."""
    if img.shape[0] == height and img.shape[1] == width:
        return img
    out = img
    if img.shape[1] != width:
        out = _resize_axis_u8(out, 1, width)  # width first
    if img.shape[0] != height:
        out = _resize_axis_u8(out, 0, height)  # then height
    return out


# ---------------------------------------------------------------------------------------------- patches
@dataclass
class VisionInputs:
    """Everything the `vision_2520` function wants, plus what the text side needs."""
    patches: np.ndarray  # [1, 2520, 768] float16, integer pixel values 0..255, cell-major, zero padded
    pos_x: np.ndarray  # [1, 2520] int32, padding = 0
    pos_y: np.ndarray  # [1, 2520] int32, padding = 0
    cos: np.ndarray  # [1, 1, 2520, 64] float16 axial RoPE
    sin: np.ndarray  # [1, 1, 2520, 64] float16
    patch_mask: np.ndarray  # [1, 2520] float16, 1 = real patch, 0 = padding
    n_soft: int  # number of valid soft tokens (cells)
    grid: tuple[int, int]  # (patch rows, patch cols)
    resized: tuple[int, int]  # (height, width) px


def rope_tables(pos_x: np.ndarray, pos_y: np.ndarray, dtype=np.float16) -> tuple[np.ndarray, np.ndarray]:
    """Axial RoPE: channels 0..31 rotate by x (16 freqs, duplicated), 32..63 by y. float32 maths, cast to dtype."""
    inv = 1.0 / (ROPE_THETA ** (np.arange(0, ROPE_FREQS * 2, 2, dtype=np.float32) / (ROPE_FREQS * 2)))  # [16]
    fx = pos_x.astype(np.float32)[:, None] * inv[None]
    fy = pos_y.astype(np.float32)[:, None] * inv[None]
    f = np.concatenate([fx, fx, fy, fy], -1)  # [N, 64]
    return np.cos(f).astype(dtype), np.sin(f).astype(dtype)


def patchify_cell_major(rgb: np.ndarray, max_patches: int = MAX_PATCHES, rope_dtype=np.float16,
                        patch_dtype=np.float16) -> VisionInputs:
    """rgb: [H, W, 3] uint8 (pixel values 0..255) with H, W multiples of 48. The patches keep the integer values
    (the graph takes 0..255). dtypes: float16 for the Core AI functions, float32 for the fp32 correctness gates."""
    if rgb.dtype != np.uint8:
        raise ValueError("patchify_cell_major takes uint8 pixels (0..255)")
    H, W = rgb.shape[:2]
    ph, pw = H // PATCH, W // PATCH
    kh, kw = ph // POOL, pw // POOL
    n = kh * kw
    p = rgb.reshape(kh, POOL, PATCH, kw, POOL, PATCH, 3)  # ky dy py kx dx px c
    p = p.transpose(0, 3, 1, 4, 2, 5, 6).reshape(n * POOL * POOL, PATCH * PATCH * 3)  # (ky kx dy dx) x (py px c)
    ky, kx, dy, dx = np.meshgrid(np.arange(kh), np.arange(kw), np.arange(POOL), np.arange(POOL), indexing="ij")
    x = (kx * POOL + dx).ravel()
    y = (ky * POOL + dy).ravel()
    patches = np.zeros((1, max_patches, PATCH * PATCH * 3), patch_dtype)
    patches[0, : n * 9] = p
    pos_x = np.zeros((1, max_patches), np.int32)
    pos_y = np.zeros((1, max_patches), np.int32)
    pos_x[0, : n * 9] = x
    pos_y[0, : n * 9] = y
    mask = np.zeros((1, max_patches), np.float16)
    mask[0, : n * 9] = 1
    cos, sin = rope_tables(pos_x[0], pos_y[0], rope_dtype)
    return VisionInputs(patches, pos_x, pos_y, cos[None, None], sin[None, None], mask, n, (ph, pw), (H, W))


RESCALE = np.float32(1.0 / 255.0)  # upstream pixel_values = patches * RESCALE (the graph folds this scale in)


def prepare_vision_inputs(image, apply_exif_orientation: bool = False, dtype=np.float16) -> VisionInputs:
    """path or PIL image -> VisionInputs (decode, size, exact resize, cell-major patches of pixel values 0..255)."""
    rgb = load_rgb(image, apply_exif_orientation)
    th, tw = target_size(rgb.shape[0], rgb.shape[1])
    rgb = resize_bicubic_aa_u8(rgb, th, tw)
    return patchify_cell_major(rgb, rope_dtype=dtype, patch_dtype=dtype)


def image_sequence(n_soft: int, bos: int, boi: int, image: int, eoi: int, eos: int) -> list[int]:
    """Token ids of a single image input, as the upstream processor builds them (no prompt)."""
    return [bos, boi] + [image] * n_soft + [eoi, eos]