"""Independent pixel verifier for the original palette and camera exercises"""

import argparse
import json
from pathlib import Path

import numpy as np
from PIL import Image

p = argparse.ArgumentParser()
p.add_argument("assets", type=Path)
a = p.parse_args()
r = a.assets
bayer = np.array([[0, 8, 2, 10], [12, 4, 14, 6], [3, 11, 1, 9], [15, 7, 13, 5]], dtype=float)


def convert(source, method, levels):
    h, w, _ = source.shape
    palette = np.floor(np.linspace(0, 255, levels) + 0.5)
    alpha = source[:, :, 3] / 255
    gray = (
        0.2126 * source[:, :, 0] + 0.7152 * source[:, :, 1] + 0.0722 * source[:, :, 2]
    ) * alpha + 255 * (1 - alpha)
    result = np.zeros((h, w), dtype=np.uint8)
    work = gray.copy()
    for y in range(h):
        for x in range(w):
            value = work[y, x] if method == "floyd" else gray[y, x]
            if method == "bayer":
                lower = min(levels - 2, int(np.searchsorted(palette, value, side="right") - 1))
                lower = max(0, lower)
                frac = (value - palette[lower]) / (palette[lower + 1] - palette[lower])
                q = lower + int(frac >= (bayer[y % 4, x % 4] + 0.5) / 16)
            else:
                distances = np.abs(palette - value)
                q = levels - 1 - int(np.argmin(distances[::-1]))
            result[y, x] = int(palette[q])
            if method == "floyd":
                error = value - palette[q]
                for dx, dy, weight in [
                    (1, 0, 7 / 16),
                    (-1, 1, 3 / 16),
                    (0, 1, 5 / 16),
                    (1, 1, 1 / 16),
                ]:
                    if 0 <= x + dx < w and y + dy < h:
                        work[y + dy, x + dx] += error * weight
    return result


source = np.asarray(Image.open(r / "dither-original-scene.png").convert("RGBA")).astype(float)
cases = []
for method in ["threshold", "bayer", "floyd"]:
    for levels in [2, 4, 8]:
        expected = convert(source, method, levels)
        actual = np.asarray(Image.open(r / f"dither-{method}-{levels}.png").convert("RGBA"))
        assert (
            np.all(actual[:, :, 0] == expected)
            and np.all(actual[:, :, 1] == expected)
            and np.all(actual[:, :, 2] == expected)
        )
        assert np.all(actual[:, :, 3] == 255)
        record = json.loads((r / f"dither-{method}-{levels}.json").read_text())
        assert [int(np.sum(expected == v)) for v in record["palette"]] == record["histogram"]
        cases.append(
            {
                "algorithm": method,
                "levels": levels,
                "pixelsCompared": int(expected.size),
                "maximumPixelError": 0,
            }
        )
alpha = np.asarray(Image.open(r / "dither-alpha-source.png").convert("RGBA")).astype(float)
expected = convert(alpha, "bayer", 2)
actual = np.asarray(Image.open(r / "dither-alpha-output.png").convert("RGBA"))
assert (
    np.array_equal(expected, actual[:, :, 0])
    and np.all(actual[:, :, 3] == 255)
    and np.all(actual[:, :21, :3] == 255)
)
print(
    json.dumps(
        {
            "scope": "Independent NumPy full-image error buffer and palette calculations checked every published PNG pixel not AI quality or calibrated color",
            "cases": cases,
            "alphaPixelsCompared": int(expected.size),
        },
        indent=2,
    )
)