qrmorph.py

#!/usr/bin/env python3
"""
qrmorph.py

Given a string, generate a QR code of each character and create a looping video.

Dependencies:
    pip install qrcode[pil] moviepy numpy
    (also install FFmpeg: https://ffmpeg.org/)
"""

import argparse
import math
import random
from typing import List, Tuple, Dict

import numpy as np
import qrcode
from qrcode.constants import ERROR_CORRECT_M, ERROR_CORRECT_H
from moviepy import VideoClip

# ----------------------------
# QR helpers
# ----------------------------

def qr_matrix_for_char(ch: str, version: int, error_correct: str = "H", border: int = 4) -> np.ndarray:
    """Return a boolean matrix (H x W) of the QR code for a single character, with a fixed version and quiet zone border."""
    ecc = ERROR_CORRECT_H if error_correct.upper() == "H" else ERROR_CORRECT_M
    qr = qrcode.QRCode(
        version=version,
        error_correction=ecc,
        box_size=1,   # box_size is irrelevant here; we use our own drawing
        border=0      # we'll add border ourselves as matrix padding
    )
    qr.add_data(ch)
    # fit=False ensures we stick to the chosen version
    qr.make(fit=False)
    core = qr.get_matrix()  # list of lists of booleans (no quiet zone)
    core = np.array(core, dtype=bool)

    if border > 0:
        core = np.pad(core, pad_width=border, mode='constant', constant_values=False)

    return core

# ----------------------------
# Matching and movement planning
# ----------------------------

Coord = Tuple[int, int]

def manhattan(a: Coord, b: Coord) -> int:
    return abs(a[0] - b[0]) + abs(a[1] - b[1])

def greedy_match(srcs: List[Coord], dsts: List[Coord]) -> Tuple[List[Tuple[int, int]], List[int], List[int]]:
    """
    Greedy nearest-neighbor matching between src and dst coordinates.
    Returns:
        pairs: list of (src_index, dst_index)
        unmatched_src_indices
        unmatched_dst_indices
    """
    if not srcs or not dsts:
        return [], list(range(len(srcs))), list(range(len(dsts)))

    unmatched_src = set(range(len(srcs)))
    unmatched_dst = set(range(len(dsts)))
    pairs = []

    # Iterate targets in a spatially reasonable order (sort by row, then col)
    dst_order = sorted(list(range(len(dsts))), key=lambda i: (dsts[i][0], dsts[i][1]))

    for di in dst_order:
        if not unmatched_src:
            break
        best_s = None
        best_d = 10**9
        dr, dc = dsts[di]
        # Linear scan; fast enough for typical sizes
        for si in list(unmatched_src):
            d = abs(srcs[si][0] - dr) + abs(srcs[si][1] - dc)
            if d < best_d:
                best_d = d
                best_s = si
                if d == 0:
                    break
        if best_s is not None:
            pairs.append((best_s, di))
            unmatched_src.remove(best_s)
            unmatched_dst.remove(di)

    return pairs, sorted(list(unmatched_src)), sorted(list(unmatched_dst))

def plan_transition_moves(
    A: np.ndarray, B: np.ndarray, rng: random.Random
) -> List[Dict]:
    """
    Plan tile movements from QR matrix A to B.
    We treat black modules as sliding tiles; if counts differ, extras fade out or in.
    Returns a list of movement dicts:
        {
            'start': (r0, c0),
            'end':   (r1, c1),
            'fade':  'none' | 'in' | 'out',
            'order': 'H' | 'V',   # which leg goes first in the L-path
            'alpha0': float (0 or 1),
            'alpha1': float (0 or 1)
        }
    """
    # Positions of black and white in A and B
    H, W = A.shape
    A_black = [(r, c) for r in range(H) for c in range(W) if A[r, c]]
    B_black = [(r, c) for r in range(H) for c in range(W) if B[r, c]]

    A_white = [(r, c) for r in range(H) for c in range(W) if not A[r, c]]
    B_white = [(r, c) for r in range(H) for c in range(W) if not B[r, c]]

    moves = []

    # Match black->black
    pairs, A_black_unmatched, B_black_unmatched = greedy_match(A_black, B_black)
    for si, di in pairs:
        r0, c0 = A_black[si]
        r1, c1 = B_black[di]
        order = 'H' if ((r0 + c0) % 2 == 0) else 'V'
        moves.append({
            'start': (r0, c0),
            'end':   (r1, c1),
            'fade':  'none',
            'order': order,
            'alpha0': 1.0,
            'alpha1': 1.0,
        })

    # Any extra black in A must fade out to nearby white positions in B or stay in place
    if A_black_unmatched:
        # For fade-out, we want tiles to stay close to their original positions
        # Rather than finding optimal matches, find nearby white positions or just fade in place
        for si in A_black_unmatched:
            r0, c0 = A_black[si]
            
            # Find a nearby white position in B, within a reasonable distance
            best_target = None
            best_dist = float('inf')
            max_search_dist = 5  # Don't move more than 5 cells for fade-out
            
            for r1, c1 in B_white:
                dist = abs(r1 - r0) + abs(c1 - c0)
                if dist <= max_search_dist and dist < best_dist:
                    best_dist = dist
                    best_target = (r1, c1)
            
            # If no nearby white position found, just fade out in place
            if best_target is None:
                r1, c1 = r0, c0
            else:
                r1, c1 = best_target
                
            order = 'V' if ((r0 + c0) % 2 == 0) else 'H'
            moves.append({
                'start': (r0, c0),
                'end':   (r1, c1),
                'fade':  'out',
                'order': order,
                'alpha0': 1.0,
                'alpha1': 0.0,
            })

    # Any shortage of black in A must fade in from white positions in A to black positions in B
    if B_black_unmatched:
        temp_pairs, _, _ = greedy_match(A_white, [B_black[i] for i in B_black_unmatched])
        for si, di_rel in temp_pairs:
            r0, c0 = A_white[si]
            r1, c1 = B_black[B_black_unmatched[di_rel]]
            order = 'H' if ((r0 + c0) % 2 == 0) else 'V'
            moves.append({
                'start': (r0, c0),
                'end':   (r1, c1),
                'fade':  'in',
                'order': order,
                'alpha0': 0.0,
                'alpha1': 1.0,
            })

    # Slight shuffle for visual variety (but deterministic given seed)
    rng.shuffle(moves)
    return moves

# ----------------------------
# Rendering helpers
# ----------------------------

def smoothstep(x: float) -> float:
    x = 0.0 if x < 0 else 1.0 if x > 1 else x
    return x * x * (3.0 - 2.0 * x)


def parse_color(s: str) -> Tuple[int, int, int]:
    """Parse a color string into an (R,G,B) tuple.

    Accepts:
      - Hex: "#RRGGBB" or "RRGGBB"
      - CSV: "R,G,B"
      - Single gray value: "128"
    """
    if s is None:
        return (255, 255, 255)
    s = s.strip()
    # Hex
    if s.startswith('#'):
        s = s[1:]
    if len(s) == 6 and all(c in '0123456789abcdefABCDEF' for c in s):
        r = int(s[0:2], 16)
        g = int(s[2:4], 16)
        b = int(s[4:6], 16)
        return (r, g, b)

    # CSV
    if ',' in s:
        parts = [p.strip() for p in s.split(',')]
        if len(parts) == 3:
            vals = []
            for p in parts:
                v = int(p)
                if v < 0 or v > 255:
                    raise ValueError(f"Color channel out of range: {p}")
                vals.append(v)
            return tuple(vals)

    # Single gray
    try:
        v = int(s)
        if 0 <= v <= 255:
            return (v, v, v)
    except Exception:
        pass

    raise ValueError(f"Unrecognized color format: {s}")

def draw_qr_static(matrix: np.ndarray, tile_px: int, pad_px: int, frame_pad_px: int, shrink: float = 0.86, bg_color: Tuple[int,int,int] = (255,255,255), square_color: Tuple[int,int,int] = (0,0,0)) -> np.ndarray:
    """Draw a QR code with shrunk squares as a numpy uint8 RGB image.

    bg_color and square_color are (R,G,B) tuples.
    """
    H, W = matrix.shape
    img_h = H * tile_px + 2 * frame_pad_px
    img_w = W * tile_px + 2 * frame_pad_px
    frame = np.empty((img_h, img_w, 3), dtype=np.uint8)
    frame[:] = bg_color

    # Shrink tiles to show as squares with gutters
    inner = max(1, int(round(tile_px * shrink)))
    pad_each_side = max(0, (tile_px - inner) // 2)

    for r in range(H):
        row = matrix[r]
        # Get black module positions
        cols = np.where(row)[0]
        for c in cols:
            # Calculate shrunk square position
            x0 = frame_pad_px + c * tile_px + pad_each_side
            y0 = frame_pad_px + r * tile_px + pad_each_side
            x1 = x0 + inner
            y1 = y0 + inner
            
            # Clamp to frame bounds
            x0 = max(0, x0)
            y0 = max(0, y0)
            x1 = min(frame.shape[1], x1)
            y1 = min(frame.shape[0], y1)
            
            if x1 > x0 and y1 > y0:
                frame[y0:y1, x0:x1] = square_color
    return frame

def draw_transition(
    moves: List[Dict],
    p: float,                 # 0..1
    H: int, W: int,
    tile_px: int,
    shrink: float,
    frame_pad_px: int,
    bg_color: Tuple[int,int,int] = (255,255,255),
    square_color: Tuple[int,int,int] = (0,0,0)
) -> np.ndarray:
    """Draw a transition frame (p in [0,1]) by sliding tiles along L-paths.

    bg_color and square_color are (R,G,B) tuples.
    """
    # Background
    img_h = H * tile_px + 2 * frame_pad_px
    img_w = W * tile_px + 2 * frame_pad_px
    frame = np.empty((img_h, img_w, 3), dtype=np.uint8)
    frame[:] = bg_color

    # Smooth movement/alpha
    q = smoothstep(p)

    # Shrink tiles to leave gutters while moving
    inner = max(1, int(round(tile_px * shrink)))
    pad_each_side = max(0, (tile_px - inner) // 2)

    for mv in moves:
        (r0, c0) = mv['start']
        (r1, c1) = mv['end']
        fade = mv['fade']
        order = mv['order']

        dr = r1 - r0
        dc = c1 - c0
        step_r = 1 if dr >= 0 else -1
        step_c = 1 if dc >= 0 else -1
        dist_r = abs(dr)
        dist_c = abs(dc)
        total = dist_r + dist_c

        # Compute how far along the L-path we are (in "cells" traveled)
        if total == 0:
            traveled = 0.0
        else:
            traveled = q * total

        # Resolve L path
        if order == 'H':
            # Move along columns first, then rows
            if traveled <= dist_c:
                cur_c = c0 + step_c * traveled
                cur_r = float(r0)
            else:
                cur_c = float(c1)  # Ensure we end exactly at destination
                cur_r = r0 + step_r * (traveled - dist_c)
        else:
            # Move along rows first, then columns
            if traveled <= dist_r:
                cur_r = r0 + step_r * traveled
                cur_c = float(c0)
            else:
                cur_r = float(r1)  # Ensure we end exactly at destination
                cur_c = c0 + step_c * (traveled - dist_r)
        
        # Clamp coordinates to valid grid bounds to prevent "flying" tiles
        cur_r = max(0, min(H - 1, cur_r))
        cur_c = max(0, min(W - 1, cur_c))

        # Alpha (fade) progress
        if fade == 'none':
            alpha = 1.0
        elif fade == 'in':
            alpha = q
        else:  # 'out'
            alpha = 1.0 - q

        if alpha <= 0.001:
            continue

        # Pixel bbox (compute raw and then clamp to valid frame region)
        x0_raw = frame_pad_px + int(round(cur_c * tile_px)) + pad_each_side
        y0_raw = frame_pad_px + int(round(cur_r * tile_px)) + pad_each_side
        x0 = max(0, x0_raw)
        y0 = max(0, y0_raw)
        x1 = min(frame.shape[1], x0_raw + inner)
        y1 = min(frame.shape[0], y0_raw + inner)

        # If the clamped box is empty, skip drawing
        if x1 <= x0 or y1 <= y0:
            continue

        # Alpha blend between bg_color and square_color per channel
        if isinstance(bg_color, tuple) and isinstance(square_color, tuple):
            blended = [int(round(bg_color[ch] * (1.0 - alpha) + square_color[ch] * alpha)) for ch in range(3)]
            frame[y0:y1, x0:x1] = blended
        else:
            # Fallback grayscale behavior
            val = int(round(255 * (1.0 - alpha)))
            frame[y0:y1, x0:x1] = val

    return frame

# ----------------------------
# Main
# ----------------------------

def build_timeline(mats: List[np.ndarray], seed: int = 1234) -> List[Dict]:
    """Precompute transition moves for each consecutive pair (including last->first)."""
    rng = random.Random(seed)
    transitions = []
    n = len(mats)
    for i in range(n):
        A = mats[i]
        B = mats[(i + 1) % n]
        moves = plan_transition_moves(A, B, rng)
        transitions.append(moves)
    return transitions

def main():
    parser = argparse.ArgumentParser(description="Generate a looping QR video from a string.")
    parser.add_argument("-i", "--input", required=True, help="Input string (one QR per character)")
    parser.add_argument("-o", "--output", default="qr_puzzle_loop.mp4", help="Output video filename (e.g., .mp4, .gif)")
    parser.add_argument("--format", choices=["mp4", "gif"], default=None, help="Output format (auto-detected from filename if not specified)")
    parser.add_argument("-v", "--version", type=int, default=4, help="Fixed QR version (1..40). All codes use this.")
    parser.add_argument("--ecc", choices=["M", "H"], default="H", help="Error correction level (M or H).")
    parser.add_argument("--border", type=int, default=4, help="Quiet zone (modules) to pad around the QR (default 4).")
    parser.add_argument("--tile", type=int, default=20, help="Pixels per module (tile size).")
    parser.add_argument("--shrink", type=float, default=0.86, help="Tile shrink during motion (0.7..1.0).")
    parser.add_argument("--framepad", type=int, default=0, help="Extra frame padding pixels around the QR.")
    parser.add_argument("--fps", type=int, default=30, help="Frames per second.")
    parser.add_argument("--bg-color", type=str, default="#ffffff", help="Background color (hex #RRGGBB, R,G,B or gray 0-255).")
    parser.add_argument("--square-color", type=str, default="#000000", help="Square color (hex #RRGGBB, R,G,B or gray 0-255).")
    parser.add_argument("-th", "--hold", type=float, default=0.6, help="Hold duration per QR (seconds).")
    parser.add_argument("-tt", "--transition", type=float, default=1.2, help="Transition duration between QRs (seconds).")
    parser.add_argument("--seed", type=int, default=42, help="Random seed for path variety.")
    args = parser.parse_args()

    s = args.input
    if len(s) == 0:
        raise SystemExit("Input string is empty.")

    # Determine output format
    output_format = args.format
    if output_format is None:
        # Auto-detect from file extension
        if args.output.lower().endswith('.gif'):
            output_format = 'gif'
        else:
            output_format = 'mp4'

    # Generate matrices (fixed version -> equal sizes)
    mats = [qr_matrix_for_char(ch, version=args.version, error_correct=args.ecc, border=args.border) for ch in s]
    H, W = mats[0].shape
    for i, M in enumerate(mats):
        if M.shape != (H, W):
            raise RuntimeError(f"Matrix size mismatch at char index {i}. Ensure fixed version is sufficient.")

    # Precompute transitions for all consecutive pairs (including last->first)
    transitions = build_timeline(mats, seed=args.seed)

    # Build timeline segments: [Hold A0] -> Transition A0->A1 -> Hold A1 -> ... -> Transition last->first
    segs = []
    t = 0.0
    n = len(mats)
    for i in range(n):
        # Hold segment
        segs.append({
            'type': 'hold',
            'index': i,
            'start': t,
            'end': t + args.hold
        })
        t += args.hold

        # Transition to next
        segs.append({
            'type': 'transition',
            'from': i,
            'to': (i + 1) % n,
            'moves': transitions[i],
            'start': t,
            'end': t + args.transition
        })
        t += args.transition

    duration = t  # end of the last transition is the end; that ends on the first QR (loopable)

    tile_px = args.tile
    frame_pad_px = args.framepad

    # Parse colors
    try:
        bg_color = parse_color(args.bg_color)
    except Exception as e:
        raise SystemExit(f"Invalid --bg-color: {e}")
    try:
        square_color = parse_color(args.square_color)
    except Exception as e:
        raise SystemExit(f"Invalid --square-color: {e}")

    # Prepare static frames for holds to save recomputing
    static_frames = {}
    for i, M in enumerate(mats):
        static_frames[i] = draw_qr_static(M, tile_px, pad_px=0, frame_pad_px=frame_pad_px, shrink=args.shrink, bg_color=bg_color, square_color=square_color)

    img_h = H * tile_px + 2 * frame_pad_px
    img_w = W * tile_px + 2 * frame_pad_px

    def make_frame(t_sec: float) -> np.ndarray:
        # Keep t inside bounds [0, duration)
        if t_sec >= duration:
            t_here = duration - 1e-6
        elif t_sec < 0:
            t_here = 0.0
        else:
            t_here = t_sec

        # Find segment
        # Small linear search is fine; segments count = 2*N
        for seg in segs:
            if seg['start'] <= t_here < seg['end']:
                if seg['type'] == 'hold':
                    idx = seg['index']
                    return static_frames[idx].copy()
                else:
                    # Transition
                    start, end = seg['start'], seg['end']
                    p = (t_here - start) / (end - start)
                    moves = seg['moves']
                    frame = draw_transition(
                        moves=moves,
                        p=p,
                        H=H, W=W,
                        tile_px=tile_px,
                        shrink=args.shrink,
                        frame_pad_px=frame_pad_px,
                        bg_color=bg_color,
                        square_color=square_color
                    )
                    return frame

        # Fallback (should not happen): return first frame
        return static_frames[0].copy()

    clip = VideoClip(make_frame, duration=duration)
    # moviepy newer versions expose `with_fps` instead of `set_fps`.
    # Also write_videofile accepts an fps parameter, but set fps on the clip
    # so downstream operations that inspect clip.fps get the correct value.
    clip = clip.with_fps(args.fps)

    # Write to file based on format
    if output_format == 'gif':
        print(f"Generating GIF: {args.output}")
        # For GIF output with basic parameters
        clip.write_gif(
            args.output,
            fps=args.fps
        )
    else:
        print(f"Generating MP4: {args.output}")
        # For widely compatible mp4: codec="libx264", audio=False, preset="medium"
        clip.write_videofile(
            args.output,
            fps=args.fps,
            codec="libx264",
            audio=False,
            preset="medium",
            threads=4
        )

if __name__ == "__main__":
    main()