#!/usr/bin/env python3
# Copyright 2026 Everpaper. Licensed under the Apache License, Version 2.0 (see ../LICENSE).
"""Check that python/eg2_image.py reproduces the upstream image processor bit for bit on the shipped test images:
the target size, the uint8 antialiased bicubic resize (vs torchvision) and the pixel values / patch layout (vs
transformers' Gemma4ImageProcessor, if transformers is installed).

    python3 check_preprocess.py [--images images] [--model-dir <google/embeddinggemma-2 snapshot>]

Needs torch + torchvision (and transformers for the processor comparison). Exit code 0 = identical everywhere.
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE.parent / "python"))
import eg2_image  # noqa: E402


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--images", default=str(HERE / "images"))
    ap.add_argument("--model-dir", default=None, help="checkpoint folder: also compare with Gemma4ImageProcessor")
    args = ap.parse_args()
    import torch
    from torchvision.transforms.v2 import functional as tvF

    proc = None
    if args.model_dir:
        from transformers import AutoProcessor

        proc = AutoProcessor.from_pretrained(args.model_dir).image_processor
    items = json.loads((Path(args.images) / "manifest.json").read_text())
    bad = 0
    for it in items:
        rgb = eg2_image.load_rgb(Path(args.images) / it["file"])
        th, tw = eg2_image.target_size(rgb.shape[0], rgb.shape[1])
        mine = eg2_image.resize_bicubic_aa_u8(rgb, th, tw)
        t = torch.from_numpy(rgb).permute(2, 0, 1).contiguous()
        ref = tvF.resize(t, [th, tw], interpolation=tvF.InterpolationMode.BICUBIC, antialias=True).permute(1, 2, 0).numpy() if (th, tw) != rgb.shape[:2] else rgb
        same = mine.shape == ref.shape and bool((mine == ref).all())
        msg = f"{it['file']:24s} {rgb.shape[1]}x{rgb.shape[0]} -> {tw}x{th}: resize {'identical' if same else 'DIFFERS'}"
        if proc is not None:
            feats = proc(images=[torch.from_numpy(rgb).permute(2, 0, 1)], return_tensors="pt")
            pv, pos = feats["pixel_values"][0], feats["image_position_ids"][0]
            vi = eg2_image.prepare_vision_inputs(Path(args.images) / it["file"], dtype=np.float32)
            valid = pos[:, 0] >= 0
            key = {(int(x), int(y)): i for i, (x, y) in enumerate(pos[valid].tolist())}
            perm = [key[(int(x), int(y))] for x, y in zip(vi.pos_x[0][: vi.n_soft * 9], vi.pos_y[0][: vi.n_soft * 9])]
            same_pv = bool(np.array_equal(pv[perm].numpy(), vi.patches[0, : vi.n_soft * 9] * eg2_image.RESCALE)) and int(feats["num_soft_tokens_per_image"][0]) == vi.n_soft
            same = same and same_pv
            msg += f", processor pixel values {'identical' if same_pv else 'DIFFER'} (n_soft {vi.n_soft})"
        bad += not same
        print(msg)
    print("CHECK_PREPROCESS", "PASS" if bad == 0 else f"FAIL ({bad} differ)")
    return 0 if bad == 0 else 1


if __name__ == "__main__":
    sys.exit(main())