"""Verify the converted ONNX model against Google's own MediaPipe runtime (TASK 51).

Feeds the SAME 256x144 RGB frame to (a) MediaPipe's ImageSegmenter running the original
.tflite (with its custom op) and (b) onnxruntime running our converted .onnx, then compares
the person-probability masks. Also writes side-by-side PNGs for an eyeball check.

Usage:  python verify_against_mediapipe.py model.tflite model.onnx photo.jpg [out_prefix]
Requires: mediapipe, onnxruntime, numpy, pillow.
"""
import sys

import mediapipe as mp
import numpy as np
import onnxruntime as ort
from mediapipe.tasks import python as mp_python
from mediapipe.tasks.python import vision
from PIL import Image

W, H = 256, 144


def main(tfl, onx, photo, prefix="verify"):
    img = Image.open(photo).convert("RGB")
    # Cover-crop to 16:9, then resize to the model's exact input so neither runtime resamples.
    sw, sh = img.size
    target = W / H
    if sw / sh > target:
        nw = int(sh * target); x = (sw - nw) // 2; img = img.crop((x, 0, x + nw, sh))
    else:
        nh = int(sw / target); y = (sh - nh) // 2; img = img.crop((0, y, sw, y + nh))
    img = img.resize((W, H), Image.BILINEAR)
    rgb = np.asarray(img, dtype=np.uint8)

    seg = vision.ImageSegmenter.create_from_options(vision.ImageSegmenterOptions(
        base_options=mp_python.BaseOptions(model_asset_path=tfl), output_confidence_masks=True))
    res = seg.segment(mp.Image(image_format=mp.ImageFormat.SRGB, data=rgb))
    ref = res.confidence_masks[0].numpy_view().astype(np.float32).reshape(H, W)

    sess = ort.InferenceSession(onx, providers=["CPUExecutionProvider"])
    x = (rgb.astype(np.float32) / 255.0)[None]
    got = sess.run(None, {"input": x})[0].reshape(H, W)

    diff = np.abs(ref - got)
    agree = np.mean((ref > 0.5) == (got > 0.5))
    print(f"ref range {ref.min():.3f}..{ref.max():.3f} person%={np.mean(ref > 0.5) * 100:.1f}")
    print(f"onnx range {got.min():.3f}..{got.max():.3f} person%={np.mean(got > 0.5) * 100:.1f}")
    print(f"max|diff|={diff.max():.5f} mean|diff|={diff.mean():.6f} binary-agreement={agree * 100:.2f}%")

    def to_img(m):
        return Image.fromarray((np.clip(m, 0, 1) * 255).astype(np.uint8)).convert("RGB")
    sheet = Image.new("RGB", (W * 3, H))
    sheet.paste(img, (0, 0)); sheet.paste(to_img(ref), (W, 0)); sheet.paste(to_img(got), (W * 2, 0))
    sheet.resize((W * 6, H * 2), Image.NEAREST).save(f"{prefix}_sheet.png")
    return diff.max()


if __name__ == "__main__":
    md = main(*sys.argv[1:])
    sys.exit(0 if md < 1e-3 else 1)