#!python
"""brush: paint a picture with code, one stroke at a time, and film it.

    brush.py TARGET.png OUT.mp4 [--seed N] [--seconds 22] [--width 720] [--finest 2]

It looks at a target picture and paints it the way a person might: a dark ground first, then big loose
strokes, then smaller and smaller ones only where the canvas still looks wrong. Every stroke starts at a
point, takes its colour from the target there, and runs along the shapes (across the direction the
brightness changes in), so strokes follow edges instead of crossing them. The film is the canvas, saved
every few dabs. Morel, Nov 2026. The idea of layered strokes that follow the gradient is Aaron Hertzmann's
("Painterly Rendering with Curved Brush Strokes of Multiple Sizes", 1998); this is my own small version.
"""
import argparse, subprocess, sys
import numpy as np
from PIL import Image, ImageFilter

ap = argparse.ArgumentParser()
ap.add_argument("target"); ap.add_argument("out")
ap.add_argument("--seed", type=int, default=7)
ap.add_argument("--seconds", type=float, default=22)
ap.add_argument("--width", type=int, default=720)
ap.add_argument("--ground", default="141d2b")
ap.add_argument("--still", help="also save the finished painting here")
ap.add_argument("--finest", type=int, default=2, help="smallest brush radius to use; 3 or 5 leaves it looser")
a = ap.parse_args()
rng = np.random.default_rng(a.seed)
FPS = 30

img = Image.open(a.target).convert("RGB")
W = a.width // 2 * 2
H = round(img.height * W / img.width) // 2 * 2
img = img.resize((W, H), Image.LANCZOS)
ground = np.array([int(a.ground[i:i + 2], 16) for i in (0, 2, 4)], np.float32) / 255
canvas = np.empty((H, W, 3), np.float32); canvas[:] = ground
canvas += rng.normal(0, 0.012, (H, W, 1)).astype(np.float32)          # a little tooth, like primed board

ff = subprocess.Popen(["ffmpeg", "-loglevel", "error", "-y", "-f", "rawvideo", "-pix_fmt", "rgb24", "-s", f"{W}x{H}", "-r", str(FPS),
                       "-i", "-", "-c:v", "libx264", "-preset", "slow", "-crf", "17", "-pix_fmt", "yuv420p",
                       "-movflags", "+faststart", a.out], stdin=subprocess.PIPE)
frames = 0
def shoot(n=1):
    global frames
    b = (np.clip(canvas, 0, 1) * 255 + 0.5).astype(np.uint8).tobytes()
    for _ in range(n):
        ff.stdin.write(b)
    frames += n

def dab_mask(r):
    """a soft round dab with a few bristle streaks in it"""
    y, x = np.mgrid[-r:r + 1, -r:r + 1].astype(np.float32)
    d = np.sqrt(x * x + y * y) / max(r, 1)
    return np.clip(1.15 - d, 0, 1) ** 0.6 * (d <= 1)

def blur(arr, radius):
    return np.asarray(Image.fromarray((np.clip(arr, 0, 1) * 255).astype(np.uint8)).filter(ImageFilter.GaussianBlur(radius)), np.float32) / 255

target = np.asarray(img, np.float32) / 255
LAYERS = [(34, 0.10, 0.9), (18, 0.075, 0.9), (9, 0.06, 0.88), (5, 0.05, 0.85), (3, 0.045, 0.85), (2, 0.04, 0.9)]   # radius, how wrong before repainting, opacity
SHARE = [0.13, 0.19, 0.22, 0.20, 0.15, 0.11]                                # of the film each layer gets
keep = [i for i, layer in enumerate(LAYERS) if layer[0] >= a.finest]
LAYERS = [LAYERS[i] for i in keep]
SHARE = [SHARE[i] / sum(SHARE[j] for j in keep) for i in keep]
shoot(FPS // 2)
total_frames = int(a.seconds * FPS)
strokes_made = 0

for (r, tol, alpha), share in zip(LAYERS, SHARE):
    ref = blur(target, r * 0.55) if r > 2 else target
    lum = ref @ np.array([0.30, 0.59, 0.11], np.float32)
    gy, gx = np.gradient(blur(np.repeat(lum[..., None], 3, 2), max(r * 0.5, 1))[..., 0])
    mask = dab_mask(r)[..., None]
    # where is the canvas still wrong? look at the average error in each cell of a grid
    err = np.sqrt(((canvas - ref) ** 2).sum(2))
    step = max(r, 2)
    pts = []
    for y0 in range(0, H, step):
        for x0 in range(0, W, step):
            cell = err[y0:y0 + step, x0:x0 + step]
            if cell.mean() > tol:
                k = np.unravel_index(cell.argmax(), cell.shape)
                pts.append((y0 + k[0], x0 + k[1]))
    rng.shuffle(pts)
    if r >= 9:                          # big strokes go on darkest first, the way an underpainting does
        pts.sort(key=lambda p: lum[p] + rng.normal(0, 0.12))
    want = max(int(total_frames * share), 1)
    every = max(len(pts) // want, 1)
    max_len = 4 + int(48 / r ** 0.5)
    for n, (y, x) in enumerate(pts):
        colour = ref[y, x] + rng.normal(0, 0.018, 3).astype(np.float32)
        colour += rng.normal(0, 0.012)                                      # each load of paint is a touch lighter or darker
        fy, fx, ldx, ldy = float(y), float(x), 0.0, 0.0
        for s in range(max_len):
            iy, ix = int(fy), int(fx)
            if not (0 <= iy < H and 0 <= ix < W):
                break
            if s > 2 and np.abs(ref[iy, ix] - colour).sum() > np.abs(ref[iy, ix] - canvas[iy, ix]).sum():
                break                                                       # the canvas is already closer than this paint
            y0, y1, x0, x1 = max(iy - r, 0), min(iy + r + 1, H), max(ix - r, 0), min(ix + r + 1, W)
            m = mask[y0 - iy + r:y1 - iy + r, x0 - ix + r:x1 - ix + r] * (alpha * (1 - 0.35 * s / max_len))   # paint runs out along the stroke
            canvas[y0:y1, x0:x1] += (colour - canvas[y0:y1, x0:x1]) * m
            g = np.hypot(gx[iy, ix], gy[iy, ix])
            if g < 1e-4:
                dx, dy = (ldx, ldy) if s else (np.cos(t := rng.uniform(0, 6.283)), np.sin(t))
            else:
                dx, dy = -gy[iy, ix] / g, gx[iy, ix] / g                    # along the edge, not across it
                if s and dx * ldx + dy * ldy < 0:
                    dx, dy = -dx, -dy
                elif not s and rng.random() < 0.5:
                    dx, dy = -dx, -dy
                dx, dy = 0.6 * dx + 0.4 * ldx if s else dx, 0.6 * dy + 0.4 * ldy if s else dy   # a wrist doesn't turn sharply
                k = np.hypot(dx, dy) or 1.0
                dx, dy = dx / k, dy / k
            ldx, ldy = dx, dy
            fx += dx * r * 0.7; fy += dy * r * 0.7
        strokes_made += 1
        if n % every == 0:
            shoot()
    shoot(4)
    print(f"brush {r:>2}: {len(pts)} strokes", file=sys.stderr)

shoot(int(FPS * 2.5))                                                       # stand back and look
ff.stdin.close(); ff.wait()
if a.still:
    Image.fromarray((np.clip(canvas, 0, 1) * 255 + 0.5).astype(np.uint8)).save(a.still)
print(f"{strokes_made} strokes, {frames} frames, {frames / FPS:.1f}s, {W}x{H}", file=sys.stderr)