#!/usr/bin/env python3
"""Render a motion-direction video from an MP4's own H.264 motion vectors.

Each pixel in the output is colored by the direction of the codec motion
vector that produced it: hue encodes the angle of motion (full circle) and
brightness encodes its magnitude. Sections that move differently from the
dominant flow (e.g. diagonally vs. straight down) show up as a distinct hue.

Usage:
    python3 motion.py INPUT.mp4 [OUTPUT.mp4]

If OUTPUT is omitted, it defaults to "<input>_direction.mp4".
"""

import sys
import os
import numpy as np
import cv2
import av
import imageio.v2 as iio


def extract_motion(path):
    """Decode the video and return per-frame dense content-motion maps.

    Returns (dx, dy) float32 arrays of shape (T, H, W) where each pixel holds
    the content displacement (px/frame) implied by the codec motion vector for
    its macroblock. Content displacement is -motion/scale (the MV points to the
    source the block was copied from).
    """
    container = av.open(path)
    stream = container.streams.video[0]
    # Ask the H.264 decoder to export per-macroblock motion vectors.
    stream.codec_context.options = {"flags2": "+export_mvs"}
    W = stream.codec_context.width
    H = stream.codec_context.height

    dxmaps, dymaps = [], []
    for frame in container.decode(stream):
        dx = np.zeros((H, W), np.float32)
        dy = np.zeros((H, W), np.float32)
        for d in frame.side_data:
            if d.type.name != "MOTION_VECTORS":
                continue
            a = d.to_ndarray()
            cmx = -a["motion_x"] / a["motion_scale"]   # content moves in +x
            cmy = -a["motion_y"] / a["motion_scale"]   # content moves in +y
            for i in range(len(a)):
                x0 = int(a["dst_x"][i] - a["w"][i] / 2)
                y0 = int(a["dst_y"][i] - a["h"][i] / 2)
                x1 = x0 + int(a["w"][i])
                y1 = y0 + int(a["h"][i])
                x0 = max(0, x0); y0 = max(0, y0)
                x1 = min(W, x1); y1 = min(H, y1)
                dx[y0:y1, x0:x1] = cmx[i]
                dy[y0:y1, x0:x1] = cmy[i]
        dxmaps.append(dx)
        dymaps.append(dy)
    return np.array(dxmaps), np.array(dymaps)


def render(dx, dy, out_path, fps=60, scale=6, temporal_window=15):
    """Color each pixel by motion direction and write the output video."""
    T, H, W = dx.shape

    # Per-frame codec MVs on noisy content are jittery; smooth them in time so
    # the persistent motion structure is visible.
    if temporal_window > 1:
        ker = np.ones(temporal_window) / temporal_window
        smooth = lambda a: np.apply_along_axis(
            lambda m: np.convolve(m, ker, mode="same"), 0, a)
        dx, dy = smooth(dx), smooth(dy)

    frames = []
    for t in range(T):
        fx = cv2.GaussianBlur(dx[t], (0, 0), 2)
        fy = cv2.GaussianBlur(dy[t], (0, 0), 2)
        ang = np.arctan2(fy, fx)                       # 0=right, +pi/2=down
        hue = ((ang / np.pi * 0.5 + 0.5) % 1.0 * 179).astype(np.uint8)
        mag = np.sqrt(fx ** 2 + fy ** 2)
        val = np.clip(mag / 3.5 * 255, 40, 255).astype(np.uint8)
        hsv = np.dstack([hue, np.full((H, W), 255, np.uint8), val])
        bgr = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
        rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
        frames.append(cv2.resize(rgb, (W * scale, H * scale),
                                 interpolation=cv2.INTER_NEAREST))

    iio.mimwrite(out_path, frames, fps=fps, codec="libx264", quality=8)
    return len(frames)


def main(argv):
    if len(argv) < 2:
        print(__doc__)
        sys.exit(1)
    in_path = argv[1]
    if len(argv) >= 3:
        out_path = argv[2]
    else:
        base, _ = os.path.splitext(in_path)
        out_path = base + "_direction.mp4"

    dx, dy = extract_motion(in_path)
    n = render(dx, dy, out_path)
    print("wrote %s (%d frames)" % (out_path, n))


if __name__ == "__main__":
    main(sys.argv)