#!/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)

# The interpreter reads source through an 80-byte line buffer and refuses
# any line that fills it (AKBASIC_MAX_LINE_LENGTH, sink_stdio.c), so a
# stored line is at most 78 characters plus its newline. 'IM$(NN) = "' and
# the closing quote spend 12 of those; 64 keeps the emitted lines under
# the ceiling with margin to spare while the index stays two digits.
PAYLOAD = 64
MAXLINE = 78

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):
    """Split a stream into strings of at most PAYLOAD characters, cutting
    only between records. The decoder reads a record's tail characters
    with MID on the string it is walking, so a run (two characters) or a
    row record (three) that straddled two IM$ entries would decode as
    garbage; DRAWSTREAM only carries the cursor, never a partial record."""
    out, cur = [], ""
    p = 0
    while p < len(blob):
        n = 3 if blob[p] == ROWREC else 2
        if len(cur) + n > PAYLOAD:
            out.append(cur)
            cur = ""
        cur += blob[p:p + n]
        p += n
    if cur:
        out.append(cur)
    return out


def simulate(raster, blob, x=0, y=0):
    """Apply one encoded stream to a raster exactly the way the BASIC
    decoder does, skips-draw-nothing and all. The cursor comes in and
    goes back out because DRAWSTREAM carries it from one IM$ entry to
    the next -- decoding the chopped strings one at a time with the
    cursor threaded through is exactly what the demo will execute."""
    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, x, y


def simulate_lines(raster, lines):
    """One stream as its chopped strings, cursor carried across the
    boundaries the way DRAWSTREAM carries X# and Y#."""
    x = y = 0
    for line in lines:
        raster, x, y = simulate(raster, line, x, y)
    return raster


def verify(frames, base_lines, delta_line_groups):
    """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. This decodes the CHOPPED
    strings, not the blobs, so a chop that split a record would fail
    here instead of corrupting the screen."""
    raster = [[1] * W for _ in range(H)]
    raster = simulate_lines(raster, base_lines)
    assert raster == frames[0], "base stream does not reproduce frame 0"
    for i, lines in enumerate(delta_line_groups):
        want = frames[(i + 1) % PHASES]
        raster = simulate_lines(raster, lines)
        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")
    for line in out:
        assert len(line) <= MAXLINE, "emitted line over %d chars: %r" % (
            MAXLINE, line)
    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_line_groups = []
    for i in range(PHASES):
        blob = encode_delta(frames[i], frames[(i + 1) % PHASES])
        lines = chop(blob)
        delta_line_groups.append(lines)
        delta_ranges.append((len(all_lines), len(all_lines) + len(lines) - 1))
        all_lines.extend(lines)
    verify(frames, base_lines, delta_line_groups)
    dbytes = sum(len(l) for g in delta_line_groups for l in g)
    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()