#!/usr/bin/env python3
"""Generate the megademo's vaporwave picture, and its six frames of video.

The picture is composed directly in akbasic's 16-colour palette space at
160x120 -- one cell is a 5x5 pixel block on the 800x600 window -- then
run-length encoded into string literals the demo decodes with INSTR and
draws with three WIDTH-2 lines per run.

The video is delta frames. Six phases of a perfect loop -- the floor
grid advancing on a geometric progression parameterised so phase 6 lands
exactly on phase 0, and the sun's slice pattern crawling on a period-six
cycle -- and each frame encodes only the rows that differ from the one
before it: an 'R' record carrying the row in two base-36 digits, then
ordinary runs, with unreachable colour 'Q' as a skip. The rays are not
in the encoding at all; the demo overdraws them live, which is what
turns a thirty-run grid row back into a five-run one.

Why strings and not DATA: the interpreter's DATA pool is 512 items and
the stroke font already holds ~350 of them. A string literal carries one
RLE run in two characters. TODO.md section 4 records the gap.

Why the dithering is by row and not by pixel: per-pixel ordered dither
shatters every gradient into one-cell runs, which is death for RLE.
Row-phase dither keeps the runs long, and horizontal banding is what a
CRT did to a gradient anyway.

Usage:
    python3 vaporwave.py --preview OUT.png     write a x5 preview of
                                               phase 0, rays included
    python3 vaporwave.py --splice MEGADEMO.BAS rewrite the generated
                                               block between the
                                               PICTURE-BEGIN/END markers
"""

import argparse
import struct
import sys
import zlib

W, H = 160, 120
PHASES = 6

# src/graphics_tables.c, Pepto's PAL measurement. Index 0 unused.
PALETTE = [
    (0x00, 0x00, 0x00), (0x00, 0x00, 0x00), (0xff, 0xff, 0xff),
    (0x88, 0x39, 0x32), (0x67, 0xb6, 0xbd), (0x8b, 0x3f, 0x96),
    (0x55, 0xa0, 0x49), (0x40, 0x31, 0x8d), (0xbf, 0xce, 0x72),
    (0x8b, 0x54, 0x29), (0x57, 0x42, 0x00), (0xb8, 0x69, 0x62),
    (0x50, 0x50, 0x50), (0x78, 0x78, 0x78), (0x94, 0xe0, 0x89),
    (0x78, 0x69, 0xc4), (0x9f, 0x9f, 0x9f),
]

# Run encoding alphabets. Position decides meaning, so overlap is fine.
# 'Q' is a skip (INSTR misses the colour table and the decoder's guard
# draws nothing) and 'R' opens a row record.
COLORCH = "ABCDEFGHIJKLMNOP"
SKIP = "Q"
ROWREC = "R"
LENCH = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
MAXRUN = len(LENCH)
PAYLOAD = 240

SUN_CX, SUN_CY, SUN_R = 80, 50, 30
HORIZON = 74
GRID_R = 1.55
RAY_SPREAD = 26


def band_mix(y, x, frac, a, b):
    threshold = (((y * 5) + ((x // 8) * 3)) % 8) / 8.0
    return b if frac > threshold else a


def sky_color(y, x):
    bands = [(7, 5), (5, 11), (11, 9)]
    seg = HORIZON / len(bands)
    i = min(int(y / seg), len(bands) - 1)
    frac = (y - (i * seg)) / seg
    a, b = bands[i]
    return band_mix(y, x, frac, a, b)


def sun_color(y):
    frac = (y - (SUN_CY - SUN_R)) / (2.0 * SUN_R)
    if frac < 0.4:
        return 8
    if frac < 0.7:
        return 9
    return 3


def sun_sliced(y, phase):
    """Period-six cuts below the sun's midline, crawling with the phase."""
    if y < SUN_CY:
        return False
    return ((y + phase) % 6) < (1 + min(1, (y - SUN_CY) // 14))


def grid_rows(t):
    """Horizontal grid lines at loop parameter t in [0,1): row k sits at
    HORIZON+1 + B*(r^(k+t)-1), so t=1 reproduces t=0 shifted one line."""
    rows, k = [], 0
    while True:
        y = HORIZON + 1 + 2.2 * (GRID_R ** (k + t) - 1.0)
        if y >= H:
            return rows
        rows.append(int(y))
        k += 1


def compose(phase, rays=False):
    img = [[1] * W for _ in range(H)]
    for y in range(HORIZON):
        for x in range(W):
            img[y][x] = sky_color(y, x)
    for i, (sx, sy) in enumerate([(9, 5), (31, 11), (52, 3), (74, 8),
                                  (99, 14), (126, 6), (147, 12), (18, 21),
                                  (139, 24), (61, 17)]):
        img[sy][sx] = 2 if i % 3 else 16
    for y in range(SUN_CY - SUN_R, SUN_CY + SUN_R + 1):
        if y < 0 or y >= HORIZON or sun_sliced(y, phase):
            continue
        dy = y - SUN_CY
        half = int((SUN_R * SUN_R - dy * dy) ** 0.5)
        for x in range(SUN_CX - half, SUN_CX + half + 1):
            img[y][x] = sun_color(y)
    for r in grid_rows(phase / float(PHASES)):
        for x in range(W):
            img[r][x] = 5
    if rays:
        for i in range(-9, 10):
            x0, y0 = SUN_CX, HORIZON
            x1, y1 = SUN_CX + (i * RAY_SPREAD), H + 20
            steps = max(abs(x1 - x0), abs(y1 - y0))
            for s in range(steps + 1):
                x = x0 + ((x1 - x0) * s) // steps
                y = y0 + ((y1 - y0) * s) // steps
                if 0 <= x < W and HORIZON < y < H:
                    img[y][x] = 4
    return img


def encode_run(color, n):
    return COLORCH[color - 1] + LENCH[n - 1]


def encode_skip(n):
    out = ""
    while n > 0:
        step = min(n, MAXRUN)
        out += SKIP + LENCH[step - 1]
        n -= step
    return out


def encode_base(img):
    """The full raster, top to bottom; rows advance automatically.
    Black spans become skips -- GRAPHIC's clear already painted them --
    which is what keeps the floor nearly free. A skip draws nothing and
    a black run draws black: the base wants the former, a delta erasing
    a moved grid line needs the latter."""
    out = ""
    for y in range(H):
        x = 0
        while x < W:
            c = img[y][x]
            n = 1
            while x + n < W and img[y][x + n] == c and n < MAXRUN:
                n += 1
            out += encode_skip(n) if c == 1 else encode_run(c, n)
            x += n
    return out


def encode_delta(prev, cur):
    """Row records for every row that differs, runs spanning the changed
    extent. An unchanged span of eight or more becomes a skip; anything
    shorter is simply repainted, which merges into its neighbours' runs.
    Black runs cost nothing either way -- the decoder draws nothing for
    colour 1 and just advances."""
    out = ""
    for y in range(H):
        diffs = [x for x in range(W) if prev[y][x] != cur[y][x]]
        if not diffs:
            continue
        lo, hi = diffs[0], diffs[-1]
        out += ROWREC + LENCH[y // 36] + LENCH[y % 36]
        if lo:
            out += encode_skip(lo)
        x = lo
        while x <= hi:
            n = 0
            while x + n <= hi and prev[y][x + n] == cur[y][x + n]:
                n += 1
            if n >= 8:
                out += encode_skip(n)
                x += n
                continue
            c = cur[y][x]
            n = 1
            while x + n <= hi and cur[y][x + n] == c and n < MAXRUN:
                n += 1
            out += encode_run(c, n)
            x += n
    return out


def chop(blob):
    return [blob[i:i + PAYLOAD] for i in range(0, len(blob), PAYLOAD)]


def simulate(raster, blob):
    """Apply one encoded stream to a raster exactly the way the BASIC
    decoder does, skips-draw-nothing and all."""
    x = y = p = 0
    while p < len(blob):
        c = blob[p]
        if c == ROWREC:
            y = (LENCH.index(blob[p + 1]) * 36) + LENCH.index(blob[p + 2])
            x = 0
            p += 3
            continue
        n = LENCH.index(blob[p + 1]) + 1
        ci = COLORCH.find(c) + 1
        if ci >= 1:
            for i in range(n):
                raster[y][x + i] = ci
        x += n
        if x >= W:
            x = 0
            y += 1
        p += 2
    return raster


def verify(frames, base_blob, delta_blobs):
    """The base must reproduce frame 0 exactly, and each delta must
    carry the raster exactly to the next frame. A skip leaves the cell
    the encoder promised was already right, so equality is total and
    any difference at all is an encoder bug."""
    raster = [[1] * W for _ in range(H)]
    raster = simulate(raster, base_blob)
    assert raster == frames[0], "base stream does not reproduce frame 0"
    for i, blob in enumerate(delta_blobs):
        want = frames[(i + 1) % PHASES]
        raster = simulate(raster, blob)
        assert raster == want, "delta %d does not reproduce its frame" % i


def emit_block(base_lines, delta_ranges, all_lines):
    out = []
    out.append("REM  ---- PICTURE-BEGIN (generated by vaporwave.py; do not")
    out.append("REM  ---- hand-edit -- rerun the script to change the picture)")
    out.append("DIM IM$(%d)" % len(all_lines))
    out.append("DIM VA#(%d)" % PHASES)
    out.append("DIM VB#(%d)" % PHASES)
    out.append("NS# = %d" % len(base_lines))
    for i, (a, b) in enumerate(delta_ranges):
        out.append("VA#(%d) = %d" % (i, a))
        out.append("VB#(%d) = %d" % (i, b))
    for i, s in enumerate(all_lines):
        out.append('IM$(%d) = "%s"' % (i, s))
    out.append("REM  ---- PICTURE-END")
    return out


def write_png(path, img, scale=5):
    w, h = W * scale, H * scale
    raw = bytearray()
    for y in range(h):
        raw.append(0)
        row = img[y // scale]
        for x in range(w):
            raw.extend(PALETTE[row[x // scale]])
    def chunk(tag, data):
        c = struct.pack(">I", len(data)) + tag + data
        return c + struct.pack(">I", zlib.crc32(tag + data) & 0xffffffff)
    with open(path, "wb") as f:
        f.write(b"\x89PNG\r\n\x1a\n")
        f.write(chunk(b"IHDR", struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0)))
        f.write(chunk(b"IDAT", zlib.compress(bytes(raw), 9)))
        f.write(chunk(b"IEND", b""))


def splice(path, block):
    with open(path) as f:
        text = f.read().splitlines()
    begin = next(i for i, l in enumerate(text) if "PICTURE-BEGIN" in l)
    end = next(i for i, l in enumerate(text) if "PICTURE-END" in l)
    text[begin:end + 1] = block
    with open(path, "w") as f:
        f.write("\n".join(text) + "\n")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--preview")
    ap.add_argument("--splice")
    args = ap.parse_args()

    frames = [compose(p) for p in range(PHASES)]
    base_blob = encode_base(frames[0])
    base_lines = chop(base_blob)
    all_lines = list(base_lines)
    delta_ranges = []
    delta_blobs = []
    for i in range(PHASES):
        blob = encode_delta(frames[i], frames[(i + 1) % PHASES])
        delta_blobs.append(blob)
        lines = chop(blob)
        delta_ranges.append((len(all_lines), len(all_lines) + len(lines) - 1))
        all_lines.extend(lines)
    verify(frames, base_blob, delta_blobs)
    dbytes = sum(len(b) for b in delta_blobs)
    print("base %d bytes in %d strings; video %d bytes in %d strings; "
          "%d strings total" %
          (sum(len(s) for s in base_lines), len(base_lines), dbytes,
           len(all_lines) - len(base_lines), len(all_lines)),
          file=sys.stderr)
    if args.preview:
        write_png(args.preview, compose(0, rays=True))
    if args.splice:
        splice(args.splice, emit_block(base_lines, delta_ranges, all_lines))


if __name__ == "__main__":
    main()