"""Simulate shifted low-resolution frames from one high-resolution scene and interleave them.

Run: python multiframe_sr.py photo.jpg (any photo; needs NumPy and Pillow).
"""
import numpy as np


def observe(scene, s, dy, dx, pixel="point"):
    """One low-res frame: the camera moved by (dy, dx) fine pixels, then the image shrunk s times."""
    # leave room for the largest shift, so that every frame has the same size
    h, w = (scene.shape[0] - s) // s * s, (scene.shape[1] - s) // s * s
    shifted = scene[dy:dy + h, dx:dx + w]
    if pixel == "point":                       # each low-res pixel records one scene point
        return shifted[::s, ::s]
    blocks = shifted.reshape(h // s, s, w // s, s, -1)
    return blocks.mean(axis=(1, 3))            # each low-res pixel averages an s x s block


def interleave(frames, s, shape, pixel="point"):
    """Put every frame's pixels back where they were measured, on the scene's own grid."""
    out = np.zeros(shape, dtype=np.float64)
    centre = 0 if pixel == "point" else s // 2  # a point sits at its block's corner, an average at its centre
    for (dy, dx), frame in frames.items():
        rows, cols = frame.shape[:2]
        out[dy + centre:dy + centre + rows * s:s, dx + centre:dx + centre + cols * s:s] = frame
    return out


if __name__ == "__main__":
    # Checks: point samples come back exactly; averaging pixels give the centred s x s moving average.
    import sys
    from PIL import Image
    from numpy.lib.stride_tricks import sliding_window_view
    scene = np.asarray(Image.open(sys.argv[1]).convert("RGB"), dtype=np.float64)
    s = 16
    c = s // 2
    h, w = (scene.shape[0] - s) // s * s, (scene.shape[1] - s) // s * s
    pixel = "point"
    frames = {(dy, dx): observe(scene, s, dy, dx, pixel) for dy in range(s) for dx in range(s)}
    point = interleave(frames, s, scene.shape, pixel)
    print("point samples, max difference from the scene:", np.abs(point[:h, :w] - scene[:h, :w]).max())
    pixel = "average"
    frames = {(dy, dx): observe(scene, s, dy, dx, pixel) for dy in range(s) for dx in range(s)}
    area = interleave(frames, s, scene.shape, pixel)
    box = sliding_window_view(scene, (s, s), axis=(0, 1)).mean(axis=(-2, -1))
    print("averaging pixels, max difference from the moving average:", np.abs(area[c:c + h, c:c + w] - box[:h, :w]).max())