#!/usr/bin/env python3
"""The README still for the PACKED player -- ROADMAP K3, FINDINGS 64.

    python3 tools/media/make_packed_media.py [container.dlxp]
            [--snap tmp/snap_packed_gate] [--map tmp/packed_snaps_gate.csv]
            [--src tmp/fr_singe] [--frame N] [--out docs/img/packed-player.png]

Blu-ray source | what the emulated 68000 actually put on screen. The right-hand
panel is MAME's own snapshot, de-double-scanned and cropped to the picture -- not
a re-render, not `dlxp.render`. It is the same rule the codec's still is built on
(tools/media/make_readme_media.py) and it is the only reason the picture is worth
printing: an encoder can be checked against its own inverse, and a screen cannot.

THE FRAME IS CHOSEN, NOT PICKED. --frame defaults to the one whose PSNR against
the source is CLOSEST TO THE MEAN over the whole gated window, so the still is
representative rather than flattering. The chosen frame and its distance from the
mean are printed, so a reader can see it was not the best one.
"""
import argparse, csv, os, shutil, struct, subprocess, sys, wave
sys.path.insert(0, "tools/encoder")
import numpy as np
from PIL import Image, ImageDraw
from dlxp import DLXP

ap = argparse.ArgumentParser()
ap.add_argument("container", nargs="?", default="tmp/packed_singe.dlxp")
ap.add_argument("--snap", default="tmp/snap_packed_gate")
ap.add_argument("--map", default="tmp/packed_snaps_gate.csv")
ap.add_argument("--src", default="tmp/fr_singe")
ap.add_argument("--frame", type=int, default=None)
ap.add_argument("--out", default="docs/img/packed-player.png")
ap.add_argument("--webm", default="docs/img/packed-player.webm")
ap.add_argument("--wav", default="tmp/packed_aud.wav",
                help="MAME's own -wavwrite capture from run 5 of "
                     "packed_run.sh; --no-audio drops it")
ap.add_argument("--no-audio", action="store_true")
ap.add_argument("--no-webm", action="store_true")
ap.add_argument("--fps", type=float, default=None,
                help="clip rate; defaults to the container's own")
a = ap.parse_args()

d = DLXP(a.container)
SNAP_W, SNAP_H = 256, 512
with open(a.map) as fh:
    shot = {int(r["frame"]): r["snapshot"] for r in csv.DictReader(fh)}
if not shot:
    sys.exit(f"{a.map} is empty -- run tools/bench/packed_run.sh first")


def screen(fr):
    """The 256x192 picture out of one MAME native snapshot."""
    p = f"{a.snap}/x68000/{shot[fr]}.png"
    s = np.asarray(Image.open(p).convert("RGB"))
    if s.shape[:2] != (SNAP_H, SNAP_W):
        sys.exit(f"{p}: expected {SNAP_W}x{SNAP_H}, got {s.shape[1]}x{s.shape[0]}")
    g = s[0::2]                                   # undo the double scan
    y = (g.shape[0] - d.H) // 2                   # the picture is centred
    return g[y:y + d.H]


def source(fr):
    p = f"{a.src}/f{fr+1:04d}.png"
    if not os.path.exists(p):
        sys.exit(f"missing {p} -- re-extract the frames the container was built "
                 f"from, or point --src at them")
    return np.asarray(Image.open(p).convert("RGB"))


def psnr(x, y):
    e = ((x.astype(float) - y.astype(float)) ** 2).mean()
    return float("inf") if e == 0 else 10 * np.log10(255.0 ** 2 / e)


frames = sorted(shot)
scores = {f: psnr(source(f), screen(f)) for f in frames}
mean = float(np.mean(list(scores.values())))
if a.frame is None:
    pick = min(scores, key=lambda f: abs(scores[f] - mean))
else:
    pick = a.frame
    if pick not in scores:
        sys.exit(f"frame {pick} was not sampled by that run")

# THE PANEL IS GATED, not just drawn. A still of the player is a claim that the
# player drew it, and the snapshot has to still be pixel-exact against the
# container for that claim to hold -- verify_packed.py checks all of them and
# this checks the one being printed, so the picture cannot outlive the result.
ref = d.render(pick)
if not np.array_equal(screen(pick), ref):
    sys.exit(f"frame {pick} is NOT pixel-exact against {a.container}. The still "
             f"is not being written: it would be a picture of a failure with a "
             f"caption saying otherwise.")

Z, BAR = 2, 22


def captioned(img, text):
    up = np.repeat(np.repeat(img, Z, 0), Z, 1)
    out = Image.new("RGB", (up.shape[1], up.shape[0] + BAR), (16, 16, 18))
    out.paste(Image.fromarray(up), (0, BAR))
    ImageDraw.Draw(out).text((6, 6), text, fill=(190, 190, 196))
    return out


left = captioned(source(pick), "Blu-ray source, cropped 256x192")
right = captioned(screen(pick),
                  "emulated 68000, MAME's own snapshot, no decoder")
out = Image.new("RGB", (left.width + right.width + 8, left.height), (16, 16, 18))
out.paste(left, (0, 0))
out.paste(right, (left.width + 8, 0))
os.makedirs(os.path.dirname(a.out), exist_ok=True)
out.save(a.out)
print(f"{a.out}: frame {pick} of {d.nframes}, {scores[pick]:.2f} dB against the "
      f"24-bit source")
print(f"  chosen as the frame CLOSEST TO THE MEAN ({mean:.2f} dB over "
      f"{len(frames)} gated frames), {abs(scores[pick]-mean):.3f} dB from it -- "
      f"best in the window is {max(scores.values()):.2f}, worst "
      f"{min(scores.values()):.2f}")
print(f"  and it is pixel-exact against {a.container}, checked before writing")


# ---------------------------------------------------------------------------
# THE CLIP.  ROADMAP K3's result is a SEQUENCE -- 120 records off a real volume,
# every one of them a literal -- and a still cannot show the one property that
# distinguishes this branch from the codec's: there is no recursion here, so
# frame 119 says nothing about frame 60 and every frame has to be its own claim.
# This writes all of them, and gates all of them before writing any (64.1).
#
# WHAT IT IS, EXACTLY, because two runs of packed_run.sh are in it:
#   picture   run 1, the GATE run -- paced at half the container's rate so the
#             snapshot lands inside the write window, cycle stealing, no sound.
#             Every panel is MAME's own snapshot, 2x nearest, no filtering.
#   sound     run 5, the AUDIO run -- the same container at 12 fps with the
#             MSM6258 on channel 3, captured by MAME's -wavwrite off the
#             speaker. It is the chip's stream, not the encoder's.
# They are two runs because they have to be: the gate run is at 6 fps and audio
# cut at 12 fps played at 6 is not this scene. The clip is therefore a
# COMPOSITE, and saying so is the point -- what it is NOT is a real-time capture
# of the shipping configuration, which at the container's own burst rate would
# show a blank layer for 99.5% of every slot (FINDINGS 64.2).
if not a.no_webm:
    if not shutil.which("ffmpeg"):
        sys.exit("ffmpeg not found -- needed for the webm (--no-webm skips it)")
    fps = a.fps or d.fps

    # EVERY frame gated, not the printed one. A packed frame is a literal: the
    # codec's last-frame test audits 120 through its own recursion and nothing
    # here does, so a clip of 120 frames is 120 separate claims.
    bad = [f for f in frames if not np.array_equal(screen(f), d.render(f))]
    if bad:
        sys.exit(f"{len(bad)} of {len(frames)} frames are NOT pixel-exact "
                 f"against {a.container} (first {bad[0]}). The clip is not "
                 f"being written: it would be a recording of a failure.")

    aud = None
    if not a.no_audio and d.has_audio and os.path.exists(a.wav):
        # THE ALIGNMENT IS FOUND, NOT ASSUMED, and then CHECKED. The capture
        # opens with the machine booting, so the stream starts at the first
        # non-zero sample -- and "first non-zero" is exactly the kind of thing
        # that is off by one byte forever with every counter in the player
        # agreeing (FINDINGS 71.5). So lump 0 is decoded with the FOUR AXES OUT
        # OF THE CONTAINER'S OWN HEADER and required to be sample-exact from
        # there. Only lump 0: past it the seams need the walk in
        # tools/bench/verify_packed_audio.py, which is what gates the whole
        # stream in check.sh. This gates the cut.
        sys.path.insert(0, "tools/encoder")
        import adpcm
        SCALE = 8                        # okim6258's signal<<4 at gain 0.50
        w = wave.open(a.wav)
        nfr, ch, rate, sw = (w.getnframes(), w.getnchannels(),
                             w.getframerate(), w.getsampwidth())
        if rate != d.aud_hz or sw != 2:
            sys.exit(f"{a.wav}: {rate} Hz / {sw*8}-bit -- the capture has to be "
                     f"the chip's own {d.aud_hz} Hz or a resampler is in the "
                     f"measurement")
        raw = w.readframes(nfr)
        left = struct.unpack("<%dh" % (nfr * ch), raw)[0::ch]
        rec = [round(v / SCALE) for v in left]
        start = next((i for i, v in enumerate(rec) if v), None)
        if start is None:
            sys.exit(f"{a.wav} is silent -- run 5 of packed_run.sh writes it")
        dec = d.decoder()
        l0 = d.lump(0)
        want = adpcm.decode(adpcm.unpack(l0, order=dec["order"]),
                            variant=dec["variant"], init=dec["init"],
                            bits=dec["bits"])
        got = rec[start:start + len(want)]
        if got != list(want):
            n = sum(x != y for x, y in zip(got, want))
            sys.exit(f"the cut at capture sample {start:,} does not decode as "
                     f"lump 0: {n:,} of {len(want):,} samples differ. The clip "
                     f"is not being written -- the sound would be the right "
                     f"scene from the wrong byte.")
        n_out = int(round(len(frames) / fps * rate))
        aud = "tmp/_packed_media_audio.wav"
        ow = wave.open(aud, "wb")
        ow.setnchannels(ch); ow.setsampwidth(sw); ow.setframerate(rate)
        ow.writeframes(raw[start * ch * sw:(start + n_out) * ch * sw])
        ow.close()
        print(f"  audio: {a.wav}, cut at sample {start:,} "
              f"({start/rate:.2f} s of boot dropped), {n_out/rate:.2f} s -- "
              f"lump 0 sample-exact against the container's own axes "
              f"({dec['variant']}/{dec['order']}, {dec['bits']}-bit, "
              f"init {dec['init']})")
    elif not a.no_audio:
        print(f"  no audio: {a.wav} is missing or the container is silent")

    tmpd = "tmp/_packed_media_frames"
    shutil.rmtree(tmpd, ignore_errors=True)
    os.makedirs(tmpd)
    for n, f in enumerate(frames):
        l = captioned(source(f), "Blu-ray source, cropped 256x192")
        r = captioned(screen(f), "emulated 68000, MAME's own snapshot, "
                                 "no decoder in the machine")
        im = Image.new("RGB", (l.width + r.width + 8, l.height), (16, 16, 18))
        im.paste(l, (0, 0)); im.paste(r, (l.width + 8, 0))
        im.save(f"{tmpd}/{n:04d}.png")

    # VP9 near-lossless: this is 256x192 palettised pixel art scaled by an
    # integer, and a codec that smooths a colour boundary would be editorialising
    # about the one thing the picture is evidence of.
    cmd = ["ffmpeg", "-y", "-loglevel", "error", "-framerate", f"{fps:g}",
           "-i", f"{tmpd}/%04d.png"]
    if aud:
        cmd += ["-i", aud, "-c:a", "libopus", "-b:a", "96k", "-shortest"]
    cmd += ["-c:v", "libvpx-vp9", "-crf", "12", "-b:v", "0",
            "-pix_fmt", "yuv444p", "-row-mt", "1", a.webm]
    subprocess.run(cmd, check=True)
    shutil.rmtree(tmpd)
    if aud:
        os.remove(aud)
    print(f"{a.webm}: {len(frames)} frames @ {fps:g} fps, "
          f"{os.path.getsize(a.webm)/1024:.0f} KB, every frame pixel-exact "
          f"against {a.container}"
          + (" -- with the chip's own audio" if aud else " -- silent"))