"""Independent PNG chunk, sample, quantization and upgrade check using stdlib only."""

import csv
import json
import struct
import sys
import zlib
from collections import Counter
from pathlib import Path

root = Path(sys.argv[1]) if len(sys.argv) > 1 else Path(__file__).parent


def decode(name):
    blob = (root / name).read_bytes()
    assert blob[:8] == b"\x89PNG\r\n\x1a\n"
    p, compressed, header = 8, b"", None
    while p < len(blob):
        n = struct.unpack(">I", blob[p : p + 4])[0]
        body = blob[p + 4 : p + 8 + n]
        assert zlib.crc32(body) == struct.unpack(">I", blob[p + 8 + n : p + 12 + n])[0]
        if body[:4] == b"IHDR":
            header = struct.unpack(">IIBBBBB", body[4:])
        if body[:4] == b"IDAT":
            compressed += body[4:]
        p += 12 + n
    assert p == len(blob) and header is not None
    w, h, depth, color, compression, filtering, interlace = header
    assert (w, h, color, compression, filtering, interlace) == (1024, 160, 0, 0, 0, 0)
    raw = zlib.decompress(compressed)
    stride = 1 + w * depth // 8
    assert len(raw) == h * stride
    values = []
    for y in range(h):
        row = raw[y * stride : (y + 1) * stride]
        assert row[0] == 0
        values.extend(row[1:] if depth == 8 else struct.unpack(">" + "H" * w, row[1:]))
    return values, depth


def rnd(value):
    return int(value + 0.5)


records = []
for mode in ["source16", "nearest", "noise"]:
    actual, depth = decode(f"banding-{mode}.png")
    expected, state, error = [], 17, 0

    def random():
        global state
        state = (state ^ (state << 13)) & 0xFFFFFFFF
        state = (state ^ (state >> 17)) & 0xFFFFFFFF
        state = (state ^ (state << 5)) & 0xFFFFFFFF
        return state / 4294967296

    for y in range(160):
        for x in range(1024):
            target = 16 + 16 * x / 1023
            noise = random() - random() if mode == "noise" else 0
            value = rnd(target * 257) if depth == 16 else max(0, min(255, rnd(target + noise)))
            expected.append(value)
            error += abs(value / (257 if depth == 16 else 1) - target)
    assert actual == expected, mode
    counts = Counter(actual)
    report = json.loads((root / f"banding-{mode}.json").read_text())
    assert report["distinctStoredSamples"] == len(counts)
    assert abs(report["meanAbsoluteByteError"] - error / len(actual)) < 1e-12
    assert report["counts"] == [{"sample": s, "pixels": n} for s, n in sorted(counts.items())]
    with (root / f"banding-{mode}.csv").open() as f:
        assert [(int(r["sample"]), int(r["pixels"])) for r in csv.DictReader(f)] == sorted(
            counts.items()
        )
    inspection, idepth = decode(f"banding-{mode}-inspection.png")
    assert idepth == 8
    assert inspection == [
        max(0, min(255, rnd((v / (257 if depth == 16 else 1) - 16) * 255 / 16))) for v in actual
    ]
    records.append(
        {"mode": mode, "depth": depth, "distinct": len(counts), "pixelsChecked": len(actual)}
    )
nearest, _ = decode("banding-nearest.png")
up, d = decode("banding-upconverted16.png")
assert d == 16 and up == [v * 257 for v in nearest] and len(set(up)) == 17
print(
    json.dumps(
        {
            "ramps": records,
            "upconvertDistinct": len(set(up)),
            "scope": "Exact saved samples not a display or AI quality benchmark",
        },
        indent=2,
    )
)