#!/usr/bin/env python3

# Based on the visualizer for BAPC 2023, Problem H ("Hidden Art").

import random
import struct
import zlib
from typing import Iterable


# Inspired by https://gist.github.com/darka/061cfac5e95b80b078b769eaae7adf84
def save_png(w: int, h: int, img: Iterable[Iterable[tuple[int, int, int]]], filename: str) -> None:
    assert w > 0 and h > 0  # assume we were not given empty image data
    with open(filename, "wb") as out:
        out.write(b"\x89PNG\r\n\x1a\n")  # start by writing the header
        for chunk_type, data in [
            # width, height, bits per pixel, color_type, pixel is RGB triple (0, 0, 0)
            (b"IHDR", struct.pack(">2I5B", w, h, 8, 2, 0, 0, 0)),
            (b"IDAT", zlib.compress(bytearray(x for row in img for x in [0, *[v for pixel in row for v in pixel]]))),
            (b"IEND", b""),
        ]:
            out.write(struct.pack(">I", len(data)))
            out.write(chunk_type)
            out.write(data)
            out.write(struct.pack(">I", zlib.crc32(data, zlib.crc32(chunk_type))))  # checksum


random.seed(42)
w, h = 300, 200
save_png(w, h, ((tuple(random.randint(0, 255) for _ in range(3)) for _ in range(w)) for _ in range(h)), "garbage.png")
