#!/usr/bin/env python3
# shakeMeasure.py - read shake.screen's effect out of a recording: Sierra's
# (a DOSBox AVI, every 70 Hz frame) or JoeyAGI's (agiHost's NAME.frames
# list with a NAME-HASH.ppm per different frame).
#
# Each frame is compared with the last frame that was in place: the same
# picture moved by (dx, dy) is a shaken frame, anything else a new
# in-place frame. Consecutive shaken frames make a beat; beats less than
# GAP_MAX frames apart make one shake.
#
#   shakeMeasure.py RECORDING [OTHER]
# prints one line per shake (beats, frames moved / in place per beat, the
# offset); with OTHER it compares the two and exits 1 if they differ.
import os
import subprocess
import sys

from PIL import Image

WIDTH = 320
HEIGHT = 200
SEARCH_X = 16
SEARCH_Y = 8
SAMPLE = 3
GAP_MAX = 20


def framesFromAvi(path):
    proc = subprocess.Popen(["ffmpeg", "-loglevel", "error", "-i", path, "-f", "rawvideo",
                             "-pix_fmt", "rgb24", "-"], stdout=subprocess.PIPE)
    size = WIDTH * HEIGHT * 3
    while True:
        buf = proc.stdout.read(size)
        if len(buf) < size:
            break
        yield buf
    proc.wait()


def framesFromHost(path):
    base = path[:-len(".frames")]
    cache = {}
    with open(path) as f:
        for line in f:
            _, h = line.split()
            if h not in cache:
                cache[h] = Image.open("%s-%s.ppm" % (base, h)).convert("RGB").tobytes()
            yield cache[h]


def pixel(buf, x, y):
    i = (y * WIDTH + x) * 3
    return buf[i:i + 3]


def offset(frame, ref):
    # The (dx, dy) that maps ref onto frame exactly, or None.
    for dy in range(-SEARCH_Y, SEARCH_Y + 1):
        for dx in range(-SEARCH_X, SEARCH_X + 1):
            ok = True
            for y in range(max(0, dy), min(HEIGHT, HEIGHT + dy), SAMPLE):
                for x in range(max(0, dx), min(WIDTH, WIDTH + dx), SAMPLE):
                    if pixel(frame, x, y) != pixel(ref, x - dx, y - dy):
                        ok = False
                        break
                if not ok:
                    break
            if ok:
                return (dx, dy)
    return None


def measure(frames):
    # A list of (moved?, offset) per frame.
    ref = None
    known = {}
    marks = []
    for frame in frames:
        if ref is None or frame == ref:
            ref = frame
            marks.append(None)
            continue
        key = hash(frame)
        if key not in known:
            known[key] = offset(frame, ref)
        o = known[key]
        if o is None or o == (0, 0):
            ref = frame
            known.clear()
            marks.append(None)
        else:
            marks.append(o)
    return marks


def shakes(marks):
    beats = []
    i = 0
    while i < len(marks):
        if marks[i] is None:
            i += 1
            continue
        start = i
        while i < len(marks) and marks[i] is not None:
            i += 1
        beats.append((start, i - start, marks[start]))
    groups = []
    for b in beats:
        if groups and b[0] - (groups[-1][-1][0] + groups[-1][-1][1]) <= GAP_MAX:
            groups[-1].append(b)
        else:
            groups.append([b])
    lines = []
    for g in groups:
        moved = sorted(set(b[1] for b in g))
        still = sorted(set(g[k + 1][0] - g[k][0] - g[k][1] for k in range(len(g) - 1)))
        offsets = sorted(set(b[2] for b in g))
        lines.append("%d beats, moved %s frames, in place %s frames, offset %s" % (len(g), moved, still, offsets))
    return lines


def load(path):
    if path.endswith(".frames"):
        return framesFromHost(path)
    return framesFromAvi(path)


if __name__ == "__main__":
    if len(sys.argv) < 2:
        sys.exit("usage: shakeMeasure.py RECORDING [OTHER]")
    first = shakes(measure(load(sys.argv[1])))
    print("%s:" % os.path.basename(sys.argv[1]))
    for line in first:
        print("  " + line)
    if len(sys.argv) > 2:
        second = shakes(measure(load(sys.argv[2])))
        print("%s:" % os.path.basename(sys.argv[2]))
        for line in second:
            print("  " + line)
        sys.exit(0 if first == second else 1)