"""Independently check the published mask arithmetic with NumPy and Pillow."""

import argparse
import ast
import hashlib
import json
import os
import tempfile
from pathlib import Path
from types import SimpleNamespace

import numpy as np
from PIL import Image

p = argparse.ArgumentParser(description=__doc__)
p.add_argument("assets", type=Path)
p.add_argument(
    "--pinned-save-source",
    type=Path,
    help="Optional pinned vggt/run_warp.py for isolated save boundary check",
)
a = p.parse_args()
root = a.assets
x = np.array(Image.open(root / "specular-brightness-demo.png").convert("RGBA"), dtype=np.float64)
luminance = (0.2126 * x[:, :, 0] + 0.7152 * x[:, :, 1] + 0.0722 * x[:, :, 2]) / 255
results = {}
for name, reverse, curve in [("default", False, 1), ("reversed", True, 1), ("curve-two", False, 2)]:
    value = 32 + x[:, :, 3] / 255 * 192 * (1 - luminance if reverse else luminance) ** curve
    expected = np.floor(value + 0.5).astype(np.uint8)
    actual = np.array(Image.open(root / f"specular-{name}-mask.png").convert("RGBA"))
    if not ((actual[:, :, :3] == expected[:, :, None]).all() and (actual[:, :, 3] == 255).all()):
        raise ValueError("Specular bytes differ " + name)
    results[name] = {"pixels": int(expected.size), "maximumByteDifference": 0}
m = np.array(Image.open(root / "worldforge-authored-mask.png").convert("RGBA"), dtype=np.float64)
g = np.floor(
    (0.2126 * m[:, :, 0] + 0.7152 * m[:, :, 1] + 0.0722 * m[:, :, 2]) * m[:, :, 3] / 255 + 0.5
)
covered = g >= 128
for name, selected in [("center", covered[80:240, 128:384]), ("full", covered)]:
    record = json.loads((root / f"worldforge-{name}-review.json").read_text())
    if not (
        covered.sum() == record["validPixels"]
        and selected.sum() == record["regionValidPixels"]
        and selected.size == record["regionPixels"]
    ):
        raise ValueError("Mask pixel count differs " + name)
    results["mask-" + name] = {
        "fullCovered": int(covered.sum()),
        "regionCovered": int(selected.sum()),
        "regionPixels": int(selected.size),
    }
if a.pinned_save_source:
    source = a.pinned_save_source.read_bytes()
    audit = json.loads((root / "worldforge-source-audit.json").read_text())
    expected_hash = next(d for d in audit["documents"] if d["name"] == "vggt/run_warp.py")["sha256"]
    if hashlib.sha256(source).hexdigest() != expected_hash:
        raise ValueError("Pinned source hash differs")
    fn = next(
        n
        for n in ast.parse(source).body
        if isinstance(n, ast.FunctionDef) and n.name == "save_warped_results"
    )
    captured = {}

    class Writer:
        COLOR_RGB2BGR = 1
        COLOR_GRAY2RGB = 2

        @staticmethod
        def cvtColor(image, code):
            return image

        @staticmethod
        def imwrite(filename, image):
            captured[Path(filename).name] = image.copy()
            return True

    namespace = {"np": np, "os": os, "cv2": Writer}
    exec(
        compile(ast.Module(body=[fn], type_ignores=[]), "<isolated-pinned-save>", "exec"), namespace
    )
    args = SimpleNamespace(
        direction="left", degree=20, frame_single=2, conf_single=1, fill_cracks=True
    )
    white = np.ones((16, 16), dtype=np.uint8)
    checker = np.indices((16, 16)).sum(axis=0) % 2
    with tempfile.TemporaryDirectory() as output:
        namespace["save_warped_results"](
            [np.zeros((16, 16, 3), dtype=np.uint8)] * 2,
            [white, checker.astype(np.float32)],
            [
                {"type": "original"},
                {"type": "single_view_warped", "direction": "left", "angle": 20},
            ],
            output,
            0,
            args,
        )
    if not (captured["mask_cam0_original_00.00_deg.png"] == 255).all():
        raise ValueError("White mask save boundary differs")
    if not np.array_equal(captured["mask_cam0_left_20.00_deg.png"], checker * 255):
        raise ValueError("Binary mask save boundary differs")
    results["pinnedMaskSave"] = {
        "isolatedArrayCases": 2,
        "whiteMeaning": "Pipeline value one becomes byte 255 before image encoding",
        "scope": "Isolated save function with image writer stub not full warping OpenCV encoding or model execution",
    }
results["scope"] = (
    "Independent NumPy byte and pixel count comparison plus optional isolated pinned save boundary no inference or physical calibration"
)
print(json.dumps(results, indent=2))
