#!/usr/bin/env python3
"""Train the integer face generator of `extensions/avatars.md` (§7, the lowering `face-int/1`),
measure it, and write the inputs of the corpus `live-conformance/avatar-face/`.

    python tools/wai_avatar_face_train.py --out <dir> [--steps 30000] [--qat-steps 8000]

Needs torch and numpy. Everything it trains on is generated here: `render` draws a face from a
seeded parameter set (head shape, skin, hair, eyes, brows, nose, mouth, background) at a pose
and an expression (head translation, scale, roll, yaw and pitch, eyes open or shut, gaze, brows,
mouth open and its width), at four times the output size, and box-filters it down. Each picture
comes with the sixteen keypoints the parameters put on it, exactly. No picture, recording or
footage from outside this repository is read. The generated set carries this repository's
licence.

It writes, under <dir>:

- `generator.wfg`: the generator (§7.2), trained in floating point, then fine-tuned with the
  integer arithmetic emulated in the loop (weights on the 2^-12 grid, activations on the
  2^-15 grid, the sigmoid table), and written by rounding each float32 weight times 2^12 and
  each bias times 2^27 half to even: exact from the float's bits;
- `identities/<name>.wfi`: identity priors (§7.3) of faces held out of training, and
  `tracks/<name>.json`: keypoint tracks (§7.4) of held-out motion, which the corpus reads;
- `report.json`: the seeds, the steps, the losses, and the measurement of §7.9: for the
  held-out set, PSNR and SSIM of the integer generator against the floating-point network and
  against the rendered truth, the keypoint error of a separately trained float keypoint
  detector over each, and the bitrate of the keypoint tracks under §7.4's coding.

Every seed is fixed (`--seed`), but float training is not claimed to be bit-reproducible: its
sums depend on the device and the thread count of the machine that runs it. What is pinned is
the generator's bytes, by SHA-256. The integer generation over those bytes is exact (§7.6),
whatever machine trained them. `int_forward` below is the arithmetic of §7.5 to §7.7 in numpy
int64: the integers this tool reports are those, not the float network's.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
import struct
import time
from pathlib import Path

import numpy as np

# ---------------------------------------------------------------------------------------
# The arithmetic of §7.5 (the driver's warp and keypoint maps) in numpy int64.
# ---------------------------------------------------------------------------------------

K = 16
FW = 12          # weights are integers at 2^12
FA = 15          # activations are integers at 2^15
FB = FW + FA     # biases are integers at 2^27
ACT_HI = 32767


def exp_neg_q16(x: np.ndarray) -> np.ndarray:
    """integer-operators §2: exp_neg_q16 of non-positive Q16 values, elementwise."""
    x = x.astype(np.int64)
    u = -x
    n = u // 45426
    g = 45426 * n
    r = x + g
    t = r + 88670
    t2 = t * t
    q = t2 >> 16
    c = 23495 * q
    y = (c >> 16) + 22544
    out = np.where(n >= 63, 0, y >> np.minimum(n, 62))
    return out.astype(np.int64)


def rdiv(n: np.ndarray, d: np.ndarray) -> np.ndarray:
    """sgn(n)·⌊(2|n| + d) / (2d)⌋ for d ≥ 1: half away from zero (integer-operators E12)."""
    a = np.abs(n)
    return np.sign(n) * ((2 * a + d) // (2 * d))


def warp_and_maps(ref: np.ndarray, src: np.ndarray, drv: np.ndarray, gauss: int, bg: int):
    """§7.5: the warped reference [h, w, 3] u8 and the keypoint maps [K, h, w] in [0, 256].

    `ref` is the identity's picture [h, w, 3] u8, `src` its keypoints and `drv` the driving
    keypoints, each [K, 2] integers in sixteenths of a pixel (x, y)."""
    h, w, _ = ref.shape
    ys, xs = np.mgrid[0:h, 0:w]
    px = (16 * xs + 8).astype(np.int64)[None]
    py = (16 * ys + 8).astype(np.int64)[None]
    dx = px - drv[:, 0].astype(np.int64)[:, None, None]
    dy = py - drv[:, 1].astype(np.int64)[:, None, None]
    d2 = dx * dx + dy * dy
    e = (d2 * gauss + 128) >> 8
    wk = exp_neg_q16(-e)                                # [K, h, w], each in [0, 65553]
    den = bg + wk.sum(axis=0)
    mx = (src[:, 0] - drv[:, 0]).astype(np.int64)[:, None, None]
    my = (src[:, 1] - drv[:, 1]).astype(np.int64)[:, None, None]
    ddx = rdiv((wk * mx).sum(axis=0), den)
    ddy = rdiv((wk * my).sum(axis=0), den)
    ux = 16 * xs + ddx
    uy = 16 * ys + ddy
    ix = ux // 16
    iy = uy // 16
    fx = ux - 16 * ix
    fy = uy - 16 * iy
    x0 = np.clip(ix, 0, w - 1)
    x1 = np.clip(ix + 1, 0, w - 1)
    y0 = np.clip(iy, 0, h - 1)
    y1 = np.clip(iy + 1, 0, h - 1)
    r = ref.astype(np.int64)
    p00, p01, p10, p11 = r[y0, x0], r[y0, x1], r[y1, x0], r[y1, x1]
    fx3, fy3 = fx[..., None], fy[..., None]
    v = ((16 - fx3) * (16 - fy3) * p00 + fx3 * (16 - fy3) * p01 + (16 - fx3) * fy3 * p10 + fx3 * fy3 * p11 + 128) >> 8
    return v.astype(np.uint8), (wk >> 8).astype(np.int64)


def sigmoid_table() -> np.ndarray:
    """§7.6: the `lut16` table of the mask, T[v + 32768] for v in −32768..32767."""
    v = np.arange(0, 32769, dtype=np.int64)
    d = 65536 + np.minimum(exp_neg_q16(-16 * v), 65536)
    s = ((1 << 31) + d // 2) // d                       # S(v) for v ≥ 0
    t = np.zeros(65536, dtype=np.int64)
    t[32768:] = s[:32768]                               # v = 0 .. 32767
    t[1:32768] = 32768 - s[32767:0:-1]                  # v = −32767 .. −1
    t[0] = 0                                            # v = −32768, never read
    return t


# ---------------------------------------------------------------------------------------
# The generator file (§7.2) and its integer forward pass (§7.6, §7.7) in numpy int64.
# ---------------------------------------------------------------------------------------


def layer_shapes(c1: int, c2: int, c3: int, k: int = K):
    """(cin, cout, stride) of the ten convolutions, in the file's order (§7.2)."""
    return [
        (6 + k, c1, 1), (c1, c2, 2), (c2, c3, 2), (c3, c3, 1), (c3, c3, 1),
        (c3, 4 * c2, 1), (2 * c2, c2, 1), (c2, 4 * c1, 1), (2 * c1, c1, 1), (c1, 4, 1),
    ]


def write_wfg(hdr: dict, layers) -> bytes:
    o = bytearray(b"WFG1")
    o += struct.pack("<BBHHHHHHI", 1, hdr["k"], hdr["width"], hdr["height"], hdr["c1"], hdr["c2"], hdr["c3"], hdr["gauss"], hdr["bg"])
    for w, b in layers:
        o += np.asarray(w, dtype="<i2").tobytes()
        o += np.asarray(b, dtype="<i4").tobytes()
    return bytes(o)


def read_wfg(b: bytes):
    assert b[:4] == b"WFG1" and b[4] == 1
    _, k, width, height, c1, c2, c3, gauss, bg = struct.unpack_from("<BBHHHHHHI", b, 4)
    hdr = dict(k=k, width=width, height=height, c1=c1, c2=c2, c3=c3, gauss=gauss, bg=bg)
    at, layers = 22, []
    for cin, cout, s in layer_shapes(c1, c2, c3, k):
        nw = cout * cin * 9
        w = np.frombuffer(b, dtype="<i2", count=nw, offset=at).astype(np.int64).reshape(cout, cin, 3, 3)
        at += 2 * nw
        bb = np.frombuffer(b, dtype="<i4", count=cout, offset=at).astype(np.int64)
        at += 4 * cout
        layers.append((w, bb, s))
    assert at == len(b)
    return hdr, layers


def half_up(p: np.ndarray, s: int) -> np.ndarray:
    return p if s == 0 else (p + (1 << (s - 1))) >> s


def conv_int(x: np.ndarray, w: np.ndarray, b: np.ndarray, stride: int) -> np.ndarray:
    """conv 3×3, zero padding 1, then rq: pre bias, shift 12, half up (§7.6)."""
    c, h, wd = x.shape
    xp = np.zeros((c, h + 2, wd + 2), np.int64)
    xp[:, 1:-1, 1:-1] = x
    ho, wo = (h - 1) // stride + 1, (wd - 1) // stride + 1
    acc = np.zeros((w.shape[0], ho, wo), np.int64)
    for ky in range(3):
        for kx in range(3):
            patch = xp[:, ky:ky + stride * (ho - 1) + 1:stride, kx:kx + stride * (wo - 1) + 1:stride]
            acc += np.einsum("oc,chw->ohw", w[:, :, ky, kx], patch)
    return half_up(acc + b[:, None, None], FW)


def shuffle2(x: np.ndarray) -> np.ndarray:
    c4, h, w = x.shape
    c = c4 // 4
    return x.reshape(c, 2, 2, h, w).transpose(0, 3, 1, 4, 2).reshape(c, 2 * h, 2 * w)


TABLE = sigmoid_table()


def int_forward(layers, warped: np.ndarray, ref: np.ndarray, heat: np.ndarray) -> np.ndarray:
    """§7.6 and §7.7 over one frame: the generated picture [h, w, 3] u8."""
    relu = lambda v: np.clip(v, 0, ACT_HI)  # noqa: E731
    t3 = (warped.transpose(2, 0, 1).astype(np.int64) - 128) * 256
    t4 = (ref.transpose(2, 0, 1).astype(np.int64) - 128) * 256
    t5 = heat.astype(np.int64) * 128
    x = np.concatenate([t3, t4, t5])
    (w1, b1, s1), (w2, b2, s2), (w3, b3, s3), (w4, b4, _), (w5, b5, _), (w6, b6, _), (w7, b7, _), (w8, b8, _), (w9, b9, _), (w10, b10, _) = layers
    e1 = relu(conv_int(x, w1, b1, s1))
    e2 = relu(conv_int(e1, w2, b2, s2))
    e3 = relu(conv_int(e2, w3, b3, s3))
    m1 = relu(conv_int(e3, w4, b4, 1))
    r = relu(conv_int(m1, w5, b5, 1) + e3)
    u2 = relu(shuffle2(conv_int(r, w6, b6, 1)))
    d2 = relu(conv_int(np.concatenate([u2, e2]), w7, b7, 1))
    u1 = relu(shuffle2(conv_int(d2, w8, b8, 1)))
    d1 = relu(conv_int(np.concatenate([u1, e1]), w9, b9, 1))
    o = conv_int(d1, w10, b10, 1)
    g = np.clip(o[:3], -32768, 32767)
    z = np.clip(o[3:4], -32767, 32767)
    m = TABLE[z + 32768]
    a = half_up(t3 * m, 15)
    om = 32768 - m
    bb = half_up(g * om, 15)
    s = a + bb
    out = np.clip(half_up(s, 8) + 128, 0, 255)
    return out.transpose(1, 2, 0).astype(np.uint8)


# ---------------------------------------------------------------------------------------
# The keypoint track coding (§7.4): what the bitrate measurement counts.
# ---------------------------------------------------------------------------------------


def zigzag(v: int) -> int:
    return 2 * v if v >= 0 else -2 * v - 1


def eg_bits(u: int, k: int) -> int:
    """The length of Exp-Golomb order k of u."""
    m = (u >> k) + 1
    return 2 * m.bit_length() - 1 + k


def encode_track(frames: np.ndarray, step: int, gen: bytes, ident: bytes, first: int, rate=(25, 1)) -> bytes:
    """A `WFK1` payload of the frames [n, K, 2] (sixteenths of a pixel), the coordinates
    rounded to multiples of `step` sixteenths, with the order of the deltas' code chosen to
    make the payload shortest."""
    q = np.floor(frames / step + 0.5).astype(np.int64)
    n = q.shape[0]
    deltas = [zigzag(int(v)) for v in (q[1:] - q[:-1]).reshape(-1)]
    best = min(range(16), key=lambda k: sum(eg_bits(u, k) for u in deltas))
    out = bytearray(b"WFK1")
    out += struct.pack("<BBBB", 1, q.shape[1], step, best)
    out += gen + ident
    out += struct.pack("<IHHH", first, n, rate[0], rate[1])
    out += q[0].reshape(-1).astype("<i2").tobytes()
    bits = []
    for u in deltas:
        m = (u >> best) + 1
        nb = m.bit_length()
        bits += [0] * (nb - 1)
        bits += [int(c) for c in format(m, "b")]
        if best:
            bits += [int(c) for c in format(u & ((1 << best) - 1), "0%db" % best)]
    while len(bits) % 8:
        bits.append(0)
    for i in range(0, len(bits), 8):
        out.append(int("".join(map(str, bits[i:i + 8])), 2))
    return bytes(out)


# ---------------------------------------------------------------------------------------
# The generated faces.
# ---------------------------------------------------------------------------------------

KEYPOINTS = [
    "left eye", "right eye", "left upper lid", "right upper lid", "left pupil", "right pupil",
    "nose tip", "mouth left", "mouth right", "upper lip", "lower lip", "chin",
    "left brow", "right brow", "left cheek", "right cheek",
]


def identity_params(rng: np.random.Generator) -> dict:
    skin = np.array([rng.uniform(0.35, 0.95), 0, 0])
    skin[1] = skin[0] * rng.uniform(0.62, 0.82)
    skin[2] = skin[0] * rng.uniform(0.45, 0.70)
    hair = rng.uniform(0.02, 0.6) * np.array([1.0, rng.uniform(0.6, 0.95), rng.uniform(0.3, 0.8)])
    if rng.random() < 0.12:
        hair = rng.uniform(0.6, 0.85) * np.ones(3)
    return dict(
        bg0=rng.uniform(0.1, 0.9, 3), bg1=rng.uniform(0.1, 0.9, 3), bgdir=rng.uniform(0, 2 * math.pi),
        bgfreq=rng.uniform(2, 9), bgamp=rng.uniform(0, 0.08),
        skin=skin, hair=hair,
        a=rng.uniform(0.25, 0.31), b=rng.uniform(0.32, 0.38),
        hairline=rng.uniform(-0.62, -0.38), hairvol=rng.uniform(0.0, 0.10), sidehair=rng.uniform(0, 1),
        eye_x=rng.uniform(0.34, 0.44), eye_y=rng.uniform(0.0, 0.14), eye_w=rng.uniform(0.40, 0.52),
        eye_h=rng.uniform(0.020, 0.030), iris=rng.uniform(0.1, 0.6, 3) * np.array([1, 1, 1.3]),
        iris_r=rng.uniform(0.85, 1.15), brow_t=rng.uniform(0.008, 0.016), brow_y=rng.uniform(0.16, 0.24),
        nose_l=rng.uniform(0.20, 0.30), nose_w=rng.uniform(0.05, 0.08),
        mouth_y=rng.uniform(0.46, 0.58), mouth_w=rng.uniform(0.55, 0.78), lip=np.clip(skin * np.array([rng.uniform(0.75, 1.0), 0.55, 0.55]), 0, 1),
        glasses=rng.random() < 0.2, gl_col=rng.uniform(0.0, 0.5, 3),
        neck=rng.uniform(0.35, 0.55), shirt=rng.uniform(0.05, 0.95, 3),
    )


def pose_params(rng: np.random.Generator, neutral: bool = False) -> dict:
    s = 0.35 if neutral else 1.0
    return dict(
        tx=rng.uniform(-0.06, 0.06) * s, ty=rng.uniform(-0.05, 0.05) * s, scale=1 + rng.uniform(-0.08, 0.08) * s,
        roll=rng.uniform(-0.25, 0.25) * s, yaw=rng.uniform(-0.45, 0.45) * s, pitch=rng.uniform(-0.25, 0.25) * s,
        open=0.0 if neutral and rng.random() < 0.7 else float(np.clip(rng.uniform(-0.3, 1.0), 0, 1)),
        smile=rng.uniform(-0.2, 0.45) * s,
        blink_l=0.0 if neutral else float(np.clip(rng.uniform(-1.5, 1.0), 0, 1)),
        blink_r=None if rng.random() < 0.85 else float(np.clip(rng.uniform(-1.0, 1.0), 0, 1)),
        gx=rng.uniform(-1, 1) * s, gy=rng.uniform(-0.7, 0.7) * s, brow=rng.uniform(-1, 1) * s,
    )


def rot(roll, yaw, pitch):
    cr, sr, cy, sy, cp, sp = math.cos(roll), math.sin(roll), math.cos(yaw), math.sin(yaw), math.cos(pitch), math.sin(pitch)
    rz = np.array([[cr, -sr, 0], [sr, cr, 0], [0, 0, 1]])
    ry = np.array([[cy, 0, sy], [0, 1, 0], [-sy, 0, cy]])
    rx = np.array([[1, 0, 0], [0, cp, -sp], [0, sp, cp]])
    return rz @ ry @ rx


def render(idp: dict, pp: dict, size: int = 64, ss: int = 4):
    """The face of `idp` at the pose and expression `pp`: [size, size, 3] u8 and its
    keypoints [K, 2] in pixels (x, y) of the output."""
    n = size * ss
    yy, xx = np.mgrid[0:n, 0:n].astype(np.float64)
    u = (xx + 0.5) / n - 0.5
    v = (yy + 0.5) / n - 0.5
    img = np.zeros((n, n, 3))
    # Background: a gradient and a soft grating.
    d = math.cos(idp["bgdir"]) * u + math.sin(idp["bgdir"]) * v
    tt = np.clip(d + 0.5, 0, 1)[..., None]
    img[:] = idp["bg0"] * (1 - tt) + idp["bg1"] * tt
    img += idp["bgamp"] * np.sin(2 * math.pi * idp["bgfreq"] * (u + 0.7 * v))[..., None]

    R = rot(pp["roll"], pp["yaw"], pp["pitch"])
    sc = pp["scale"]
    cx, cy = pp["tx"], pp["ty"] - 0.02
    A, B, C = idp["a"], idp["b"], idp["a"] * 0.95

    def proj(p):
        q = R @ np.asarray(p, dtype=np.float64)
        return np.array([cx + sc * q[0], cy + sc * q[1]]), q[2]

    def surf(x, y):
        z2 = 1 - (x / A) ** 2 - (y / B) ** 2
        return C * math.sqrt(max(z2, 0.0))

    def ell(c0, ax, ay, ang=0.0):
        ca, sa = math.cos(ang), math.sin(ang)
        du, dv = u - c0[0], v - c0[1]
        ru = (ca * du + sa * dv) / max(ax, 1e-6)
        rv = (-sa * du + ca * dv) / max(ay, 1e-6)
        return ru * ru + rv * rv

    def paint(mask, col, soft=1.5 / n):
        nonlocal img
        m = np.clip(mask, 0, 1)[..., None]
        img = img * (1 - m) + np.asarray(col) * m

    def inside(q, edge):
        # q < 1 is inside; a soft edge of about `edge` in the normalised measure.
        return np.clip((1 - q) / edge + 0.5, 0, 1)

    roll = pp["roll"]
    # Neck and shirt.
    nk, _ = proj([0, B * 0.85, -C * 0.2])
    paint(inside(ell(nk + np.array([0, 0.12]), idp["neck"] * A * sc, 0.25 * sc, roll * 0.3), 0.08), idp["skin"] * 0.82)
    paint(np.clip((v - (nk[1] + 0.22 * sc)) / 0.01, 0, 1) * inside(ell([nk[0], 0.62], 0.5, 0.3), 0.1), idp["shirt"])
    # Hair volume behind the head.
    hc, _ = proj([0, -B * 0.25, -C * 0.3])
    paint(inside(ell(hc, (A + idp["hairvol"] + 0.02 * idp["sidehair"]) * sc, (B * 0.95 + idp["hairvol"]) * sc, roll), 0.06), idp["hair"])
    # The head: the silhouette of the rotated ellipsoid, approximated by its projected centre
    # and axes, and shaded by a light the head turns under.
    hcen, _ = proj([0, 0, 0])
    q = ell(hcen, A * sc, B * sc, roll)
    head = inside(q, 0.05)
    du, dv = u - hcen[0], v - hcen[1]
    ca, sa = math.cos(roll), math.sin(roll)
    lu = (ca * du + sa * dv) / (A * sc)
    lv = (-sa * du + ca * dv) / (B * sc)
    nz = np.sqrt(np.clip(1 - lu * lu - lv * lv, 0, 1))
    light = np.array([-0.45, -0.5, 0.74])
    lam = np.clip(lu * light[0] * math.cos(pp["yaw"]) + lv * light[1] + nz * light[2] + 0.25 * math.sin(pp["yaw"]), 0, 1)
    face_col = idp["skin"][None, None] * (0.55 + 0.5 * lam[..., None])
    m = head[..., None]
    img = img * (1 - m) + face_col * m
    # Hair over the top of the head, above a hairline that turns with it.
    hl = idp["hairline"] + 0.08 * pp["pitch"]
    hairmask = head * np.clip((hl - lv) / 0.06 + 0.5, 0, 1)
    side = idp["sidehair"] * np.clip((np.abs(lu) - 0.82) / 0.05, 0, 1) * np.clip((0.2 - lv) / 0.1, 0, 1)
    paint(np.maximum(hairmask, head * side), idp["hair"] * (0.8 + 0.3 * lam[..., None]))

    kp = np.zeros((K, 2))

    def visible_w(xside):
        # Foreshortening of a feature at x on the surface: |n·view| for a turned head.
        return max(0.25, math.cos(pp["yaw"] + math.asin(max(-0.95, min(0.95, xside / A)))))

    # Eyes.
    ey = -idp["eye_y"] * B
    for side, (ki, li, pi_) in ((-1, (0, 2, 4)), (1, (1, 3, 5))):
        ex = side * idp["eye_x"] * A
        p3 = [ex, ey, surf(ex, ey) - 0.01]
        c2, _ = proj(p3)
        fw = visible_w(ex)
        blink = pp["blink_l"] if (side < 0 or pp["blink_r"] is None) else pp["blink_r"]
        ew = idp["eye_w"] * A * sc * fw
        eh = idp["eye_h"] * sc * (1 - 0.92 * blink)
        socket = ell(c2, ew * 0.75, idp["eye_h"] * sc * 1.1, roll)
        paint(inside(socket, 0.3) * 0.25, idp["skin"] * 0.6)
        eyem = inside(ell(c2, ew * 0.5, eh, roll), 0.12)
        paint(eyem, [0.95, 0.94, 0.92])
        g = np.array([pp["gx"] * ew * 0.22, pp["gy"] * eh * 0.35])
        g = np.array([ca * g[0] - sa * g[1], sa * g[0] + ca * g[1]])
        ic = c2 + g
        ir = idp["iris_r"] * idp["eye_h"] * sc
        paint(inside(ell(ic, ir * fw, ir, roll), 0.15) * eyem, idp["iris"])
        paint(inside(ell(ic, ir * 0.45 * fw, ir * 0.45, roll), 0.2) * eyem, [0.03, 0.03, 0.04])
        lid = c2 + np.array([sa * eh, -ca * eh])
        paint(inside(ell(c2, ew * 0.52, eh * 1.06, roll), 0.12) * (1 - inside(ell(c2, ew * 0.48, eh * 0.9, roll), 0.12)), idp["skin"] * 0.35)
        kp[ki] = c2
        kp[li] = lid
        kp[pi_] = ic
        # Brows.
        bx = ex
        by = ey - idp["brow_y"] * B - 0.012 * pp["brow"]
        b2, _ = proj([bx, by, surf(bx, by)])
        paint(inside(ell(b2, ew * 0.62, idp["brow_t"] * sc, roll + side * 0.08 * (pp["brow"] - 0.3)), 0.25), idp["hair"] * 0.7)
        kp[12 if side < 0 else 13] = b2
    # Glasses.
    if idp["glasses"]:
        for ki in (0, 1):
            r0 = idp["eye_w"] * A * sc * 0.62
            ring = ell(kp[ki], r0, r0 * 0.8, roll)
            paint(np.clip(1 - np.abs(ring - 1) / 0.12, 0, 1), idp["gl_col"])
    # Nose: shading along its ridge and two nostrils, with the tip standing out of the face.
    ny = idp["nose_l"] * B
    tip, _ = proj([0, ny, surf(0, ny) + 0.05])
    root, _ = proj([0, ey, surf(0, ey)])
    paint(inside(ell((tip + root) / 2, idp["nose_w"] * 0.35 * sc, np.linalg.norm(tip - root) / 2 + 0.01, roll), 0.4) * 0.25, idp["skin"] * 0.7)
    for side in (-1, 1):
        nn, _ = proj([side * idp["nose_w"] * 0.45, ny + 0.008, surf(0, ny) + 0.03])
        paint(inside(ell(nn, 0.010 * sc, 0.006 * sc, roll), 0.4), idp["skin"] * 0.3)
    kp[6] = tip
    # Mouth: lips, the opening and teeth.
    my = idp["mouth_y"] * B
    mw = idp["mouth_w"] * A * (1 + 0.35 * pp["smile"])
    lc, _ = proj([-mw / 2, my - 0.012 * pp["smile"], surf(-mw / 2, my)])
    rc, _ = proj([mw / 2, my - 0.012 * pp["smile"], surf(mw / 2, my)])
    mc, _ = proj([0, my, surf(0, my) + 0.01])
    opn = pp["open"] * 0.045
    up, _ = proj([0, my - 0.008 - opn * 0.2, surf(0, my) + 0.012])
    lo, _ = proj([0, my + 0.009 + opn, surf(0, my) + 0.01])
    mid = (up + lo) / 2
    half_w = np.linalg.norm(rc - lc) / 2
    ang = math.atan2(rc[1] - lc[1], rc[0] - lc[0])
    hh = np.linalg.norm(lo - up) / 2
    paint(inside(ell(mid, half_w, hh + 0.005, ang), 0.15), idp["lip"])
    if pp["open"] > 0.05:
        inner = inside(ell(mid, half_w * 0.85, max(hh - 0.006, 0.002), ang), 0.2)
        paint(inner, [0.25, 0.06, 0.07])
        teeth = inner * np.clip(((up[1] + 0.005 * sc + 0.01 * pp["open"]) - v) / 0.004 + 0.5, 0, 1)
        paint(teeth, [0.93, 0.92, 0.88])
    kp[7], kp[8], kp[9], kp[10] = lc, rc, up, lo
    # Chin and cheeks.
    chin, _ = proj([0, B * 0.98, surf(0, B * 0.98) + 0.01])
    kp[11] = chin
    for ki, side in ((14, -1), (15, 1)):
        cpt, _ = proj([side * A * 0.97, 0.15 * B, surf(side * A * 0.97, 0.15 * B)])
        kp[ki] = cpt
    img = np.clip(img, 0, 1)
    small = img.reshape(size, ss, size, ss, 3).mean(axis=(1, 3))
    out = np.clip(np.floor(small * 255 + 0.5), 0, 255).astype(np.uint8)
    kpx = (kp + 0.5) * size  # u, v in [-0.5, 0.5) -> pixels
    return out, kpx


def kp_q4(kp: np.ndarray) -> np.ndarray:
    """Keypoints in pixels to sixteenths of a pixel, rounded half to even."""
    return np.round(kp * 16).astype(np.int64)


# ---------------------------------------------------------------------------------------
# Training (torch).
# ---------------------------------------------------------------------------------------


def make_sample(seed: int, size: int, gauss: int, bg: int):
    rng = np.random.default_rng(seed)
    idp = identity_params(rng)
    ref, rk = render(idp, pose_params(rng, neutral=True), size)
    tgt, tk = render(idp, pose_params(rng), size)
    rq, tq = kp_q4(rk), kp_q4(tk)
    warped, heat = warp_and_maps(ref, rq, tq, gauss, bg)
    return ref, warped, heat, tgt, rq, tq


def sequence(idp_seed: int, n: int, size: int, fps: int = 25):
    """A held-out motion: smooth head motion, blinks, speech-like mouth, glances."""
    rng = np.random.default_rng(idp_seed)
    idp = identity_params(rng)
    t = np.arange(n) / fps

    def smooth(amp, fmin=0.15, fmax=0.6):
        f = rng.uniform(fmin, fmax, 3)
        ph = rng.uniform(0, 2 * math.pi, 3)
        return amp * sum(np.sin(2 * math.pi * f[i] * t + ph[i]) for i in range(3)) / 3

    blink = np.zeros(n)
    for b0 in rng.uniform(0, n / fps, max(1, int(n / fps / 3))):
        blink = np.maximum(blink, np.clip(1 - np.abs(t - b0) / 0.12, 0, 1))
    speech = np.clip(0.5 + 0.5 * np.sin(2 * math.pi * rng.uniform(3, 5) * t + rng.uniform(0, 6)) * (0.6 + 0.4 * np.sin(2 * math.pi * 0.7 * t)), 0, 1)
    poses = []
    tx, ty, sc, ro, yw, pi_ = smooth(0.05), smooth(0.04), 1 + smooth(0.06), smooth(0.18), smooth(0.35), smooth(0.18)
    sm, gx, gy, br = smooth(0.25) + 0.1, smooth(0.8, 0.2, 1.0), smooth(0.5, 0.2, 1.0), smooth(0.7, 0.2, 0.8)
    for i in range(n):
        poses.append(dict(tx=tx[i], ty=ty[i], scale=sc[i], roll=ro[i], yaw=yw[i], pitch=pi_[i], open=float(speech[i]) * 0.9,
                          smile=sm[i], blink_l=float(blink[i]), blink_r=None, gx=gx[i], gy=gy[i], brow=br[i]))
    rng2 = np.random.default_rng(idp_seed + 1)
    ref, rk = render(idp, pose_params(rng2, neutral=True), size)
    frames, kps = zip(*(render(idp, p, size) for p in poses))
    return ref, rk, np.stack(frames), np.stack(kps)


try:
    import torch as _torch
    from torch.utils.data import IterableDataset as _Iterable
except ImportError:  # the measurement-free parts of this tool run without torch
    _torch, _Iterable = None, object


POSES = 9  # renders per training identity: a near-neutral one and eight others


def _render_identity(args):
    seed, size = args
    rng = np.random.default_rng(seed)
    idp = identity_params(rng)
    imgs, kps = [], []
    for j in range(POSES):
        img, kp = render(idp, pose_params(rng, neutral=(j == 0)), size)
        imgs.append(img)
        kps.append(kp)
    return np.stack(imgs), np.stack(kps)


def render_pool(path: Path, n_id: int, size: int, workers: int, base: int):
    """Render `n_id` training identities, POSES each, into `imgs.npy` and `kps.npy` (memory
    mapped, so the data loader's workers share them)."""
    import multiprocessing as mp

    imgs = np.lib.format.open_memmap(path / "imgs.npy", mode="w+", dtype=np.uint8, shape=(n_id, POSES, size, size, 3))
    kps = np.lib.format.open_memmap(path / "kps.npy", mode="w+", dtype=np.float64, shape=(n_id, POSES, K, 2))
    t0 = time.time()
    with mp.get_context("spawn").Pool(workers) as pool:
        for i, (im, kp) in enumerate(pool.imap(_render_identity, [(base + i, size) for i in range(n_id)], chunksize=16)):
            imgs[i] = im
            kps[i] = kp
            if (i + 1) % 2000 == 0:
                print(f"rendered {i + 1}/{n_id} identities ({time.time() - t0:.0f} s)", flush=True)
    imgs.flush()
    kps.flush()
    del imgs, kps


class FaceStream(_Iterable):
    """Training samples, made in the data loader's workers from the rendered pool: a
    reference face, its warp to another pose of the same face, the keypoint maps, and the
    face rendered at that pose."""

    def __init__(self, pool: Path, base, gauss, bg):
        self.pool, self.base, self.gauss, self.bg = str(pool), base, gauss, bg

    def __iter__(self):
        info = _torch.utils.data.get_worker_info()
        wid = info.id if info else 0
        imgs = np.load(self.pool + "/imgs.npy", mmap_mode="r")
        kps = np.load(self.pool + "/kps.npy", mmap_mode="r")
        rng = np.random.default_rng(self.base + 1_000_003 * wid)
        f = lambda a: _torch.from_numpy(np.ascontiguousarray(a.transpose(2, 0, 1)).astype(np.float32))  # noqa: E731
        while True:
            i = int(rng.integers(0, imgs.shape[0]))
            r = 0 if rng.random() < 0.6 else int(rng.integers(1, POSES))
            t = int(rng.integers(0, POSES - 1))
            t = t + 1 if t >= r else t
            ref, tgt = np.asarray(imgs[i, r]), np.asarray(imgs[i, t])
            warped, heat = warp_and_maps(ref, kp_q4(np.asarray(kps[i, r])), kp_q4(np.asarray(kps[i, t])), self.gauss, self.bg)
            yield f(warped), f(ref), _torch.from_numpy(heat.astype(np.float32)), f(tgt)


def train(args, out: Path, report: dict):
    import torch
    import torch.nn as nn
    import torch.nn.functional as F

    dev = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
    torch.manual_seed(args.seed)
    size, gauss, bg = args.size, args.gauss, args.bg
    c1, c2, c3 = args.c1, args.c2, args.c3
    shapes = layer_shapes(c1, c2, c3)

    class Q(torch.autograd.Function):
        """Round half up onto a grid, gradient straight through."""

        @staticmethod
        def forward(ctx, x, scale):
            return torch.floor(x * scale + 0.5) / scale

        @staticmethod
        def backward(ctx, g):
            return g, None

    class QE(torch.autograd.Function):
        """Round half to even onto a grid (the export's rule), gradient straight through."""

        @staticmethod
        def forward(ctx, x, scale):
            return torch.round(x * scale) / scale

        @staticmethod
        def backward(ctx, g):
            return g, None

    class Net(nn.Module):
        def __init__(self):
            super().__init__()
            self.convs = nn.ModuleList([nn.Conv2d(ci, co, 3, stride=s, padding=1) for ci, co, s in shapes])
            self.qat = False

        def conv(self, i, x):
            c = self.convs[i]
            w, b = c.weight, c.bias
            if self.qat:
                w = QE.apply(torch.clamp(w, -32768 / 4096, 32767 / 4096), 4096.0)
                b = QE.apply(b, float(1 << FB))
            y = F.conv2d(x, w, b, stride=c.stride, padding=1)
            if self.qat:
                y = Q.apply(y, 32768.0)
            return y

        def forward(self, warped, ref, heat):
            act = lambda v: torch.clamp(v, 0, ACT_HI / 32768)  # noqa: E731
            t3 = (warped - 128) / 128
            t4 = (ref - 128) / 128
            t5 = heat / 256
            x = torch.cat([t3, t4, t5], 1)
            e1 = act(self.conv(0, x))
            e2 = act(self.conv(1, e1))
            e3 = act(self.conv(2, e2))
            m1 = act(self.conv(3, e3))
            r = act(self.conv(4, m1) + e3)
            u2 = act(F.pixel_shuffle(self.conv(5, r), 2))
            d2 = act(self.conv(6, torch.cat([u2, e2], 1)))
            u1 = act(F.pixel_shuffle(self.conv(7, d2), 2))
            d1 = act(self.conv(8, torch.cat([u1, e1], 1)))
            o = self.conv(9, d1)
            g = torch.clamp(o[:, :3], -1, 32767 / 32768)
            z = o[:, 3:4]
            m = torch.sigmoid(8 * z)
            if self.qat:
                m = Q.apply(m, 32768.0)
            s = m * t3 + (1 - m) * g
            return s * 128 + 128, g * 128 + 128, 8 * z

    pool = out / "pool"
    pool.mkdir(exist_ok=True)
    if not (pool / "imgs.npy").exists() or np.load(pool / "imgs.npy", mmap_mode="r").shape[0] != args.identities:
        render_pool(pool, args.identities, size, args.workers, args.seed * 104729)
    if args.measure_only:
        # The float network and the generator of an earlier run of this tool, from --out.
        fnet = Net().to(dev)
        fnet.load_state_dict(torch.load(out / "float.pt", map_location=dev))
        fnet.eval()
        blob = (out / "generator.wfg").read_bytes()
        report["generator"] = {"sha256": hashlib.sha256(blob).hexdigest(), "bytes": len(blob)}
        prev = json.loads((out / "report.json").read_text()) if (out / "report.json").exists() else {}
        for k in ("losses", "train_seconds"):
            if k in prev:
                report[k] = prev[k]
        return fnet, blob
    net = Net().to(dev)
    with torch.no_grad():
        # Start with an even mask, and smaller weights.
        net.convs[9].bias[3] = 0.0
        for c in net.convs:
            c.weight.mul_(0.7)

    from torch.utils.data import DataLoader

    def batches(base):
        dl = DataLoader(FaceStream(pool, base, gauss, bg), batch_size=args.batch, num_workers=args.workers, persistent_workers=False)
        return iter(dl)

    def loss_fn(y, t):
        l1 = (y - t).abs().mean()
        gy = lambda a: a[..., 1:, :] - a[..., :-1, :]  # noqa: E731
        gx = lambda a: a[..., :, 1:] - a[..., :, :-1]  # noqa: E731
        return l1 + 1.0 * ((gy(y) - gy(t)).abs().mean() + (gx(y) - gx(t)).abs().mean())

    losses = []
    phases = [("float", args.steps, args.lr, False), ("qat", args.qat_steps, args.lr * 0.2, True)]
    it = batches(args.seed * 7919)
    t0 = time.time()
    for name, steps, lr, qat in phases:
        net.qat = qat
        opt = torch.optim.Adam(net.parameters(), lr=lr)
        sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, max(steps, 1))
        run = 0.0
        for step in range(steps):
            w_, r_, h_, t_ = (a.to(dev) for a in next(it))
            y, g_, z_ = net(w_, r_, h_)
            # The blend's own loss, the generated picture's (so the generator learns a whole
            # face whatever the mask chooses), and a penalty on mask logits past ±4 (so the
            # mask never saturates where its gradient vanishes).
            loss = loss_fn(y, t_) + 0.25 * (g_ - t_).abs().mean() + 0.1 * torch.relu(z_.abs() - 4).mean()
            opt.zero_grad()
            loss.backward()
            opt.step()
            sched.step()
            run += loss.item()
            if (step + 1) % 500 == 0:
                losses.append([name, step + 1, run / 500])
                print(f"{name} {step + 1}/{steps} loss {run / 500:.3f} ({time.time() - t0:.0f} s)", flush=True)
                run = 0.0
        if name == "float":
            float_state = {k: v.detach().clone() for k, v in net.state_dict().items()}
            torch.save(float_state, out / "float.pt")
    report["losses"] = losses
    report["train_seconds"] = round(time.time() - t0)

    # Export: each float32 weight times 2^12 and each bias times 2^27, rounded half to even.
    layers = []
    for c in net.convs:
        w = c.weight.detach().cpu().numpy().astype(np.float32).astype(np.float64)
        b = c.bias.detach().cpu().numpy().astype(np.float32).astype(np.float64)
        wi = np.clip(np.round(w * 4096), -32768, 32767).astype(np.int64)
        bi = np.clip(np.round(b * (1 << FB)), -(1 << 31), (1 << 31) - 1).astype(np.int64)
        layers.append((wi, bi))
    hdr = dict(k=K, width=size, height=size, c1=c1, c2=c2, c3=c3, gauss=gauss, bg=bg)
    blob = write_wfg(hdr, layers)
    (out / "generator.wfg").write_bytes(blob)
    report["generator"] = {"sha256": hashlib.sha256(blob).hexdigest(), "bytes": len(blob)}

    # The float network, for the measurement: the float phase's weights, run in float32.
    fnet = Net().to(dev)
    fnet.load_state_dict(float_state)
    fnet.qat = False
    fnet.eval()
    return fnet, blob


def train_detector(args, dev, pool: Path):
    """A float keypoint detector, trained on the rendered training faces and their keypoints,
    for the measurement's keypoint error: it reads a picture and gives its sixteen keypoints."""
    import torch
    import torch.nn as nn

    torch.manual_seed(args.seed + 1)
    size = args.size

    class Det(nn.Module):
        """Keypoint maps at half the picture's size, each read by a spatial softmax as the
        expected position: sub-pixel by construction."""

        def __init__(self):
            super().__init__()

            def block(a, b, s):
                return [nn.Conv2d(a, b, 3, stride=s, padding=1), nn.BatchNorm2d(b), nn.ReLU()]
            self.down = nn.Sequential(*block(3, 32, 1), *block(32, 64, 2), *block(64, 64, 1), *block(64, 128, 2), *block(128, 128, 1), *block(128, 128, 1))
            self.up = nn.Sequential(nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False), *block(128, 64, 1), nn.Conv2d(64, K, 3, padding=1))
            half = size // 2
            g = (torch.arange(half, dtype=torch.float32) + 0.5) * 2
            self.register_buffer("gx", g.view(1, 1, 1, half))
            self.register_buffer("gy", g.view(1, 1, half, 1))

        def forward(self, x):
            m = self.up(self.down(x / 255 - 0.5))
            n, k, h, w = m.shape
            p = torch.softmax(m.reshape(n, k, -1), -1).reshape(n, k, h, w)
            return torch.stack([(p * self.gx).sum((2, 3)), (p * self.gy).sum((2, 3))], -1)

    det = Det().to(dev)
    opt = torch.optim.Adam(det.parameters(), lr=1e-3)
    det.train()
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, args.det_steps)
    rng = np.random.default_rng(args.seed + 99)
    all_imgs = np.load(pool / "imgs.npy", mmap_mode="r")
    all_kps = np.load(pool / "kps.npy", mmap_mode="r")
    n = min(args.det_pool, all_imgs.shape[0])
    imgs = torch.from_numpy(np.ascontiguousarray(np.asarray(all_imgs[:n]).reshape(-1, size, size, 3).transpose(0, 3, 1, 2)).astype(np.float32)).to(dev)
    kps = torch.from_numpy(np.asarray(all_kps[:n]).reshape(-1, K, 2).astype(np.float32)).to(dev)
    for step in range(args.det_steps):
        idx = torch.from_numpy(rng.integers(0, imgs.shape[0], args.batch)).to(dev)
        loss = (det(imgs[idx]) - kps[idx]).abs().mean()
        opt.zero_grad()
        loss.backward()
        opt.step()
        sched.step()
        if (step + 1) % 1000 == 0:
            print(f"detector {step + 1}/{args.det_steps} loss {loss.item():.3f} px", flush=True)
    det.eval()
    return det


def psnr(a: np.ndarray, b: np.ndarray) -> float:
    mse = np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2)
    return float("inf") if mse == 0 else 10 * math.log10(255 * 255 / mse)


def ssim(a: np.ndarray, b: np.ndarray) -> float:
    """SSIM (11×11 Gaussian window, σ 1.5, K1 0.01, K2 0.03) on each channel, averaged."""
    g = np.exp(-((np.arange(11) - 5) ** 2) / (2 * 1.5 ** 2))
    g /= g.sum()

    def blur(x):
        x = np.apply_along_axis(lambda r: np.convolve(r, g, mode="valid"), 0, x)
        return np.apply_along_axis(lambda r: np.convolve(r, g, mode="valid"), 1, x)

    c1, c2 = (0.01 * 255) ** 2, (0.03 * 255) ** 2
    vals = []
    for ch in range(3):
        x, y = a[..., ch].astype(np.float64), b[..., ch].astype(np.float64)
        mx, my = blur(x), blur(y)
        sxx = blur(x * x) - mx * mx
        syy = blur(y * y) - my * my
        sxy = blur(x * y) - mx * my
        s = ((2 * mx * my + c1) * (2 * sxy + c2)) / ((mx * mx + my * my + c1) * (sxx + syy + c2))
        vals.append(s.mean())
    return float(np.mean(vals))


def measure(args, out: Path, report: dict, fnet, blob: bytes):
    import torch

    dev = next(fnet.parameters()).device
    det = train_detector(args, dev, out / "pool")
    hdr, layers = read_wfg(blob)
    ilayers = [(w, b, s) for (w, b, s) in layers]
    size, gauss, bg = args.size, args.gauss, args.bg
    rows = {"int_float": [], "int_truth": [], "float_truth": [], "warp_truth": []}
    ss = {"int_float": [], "int_truth": [], "float_truth": [], "warp_truth": []}
    kerr = {"truth": [], "int": [], "float": []}
    mouth = {"truth": [], "int": [], "float": []}
    bitrate = {"truth_step4": [], "truth_step1": [], "detected_step4": [], "detected_step1": []}
    gen_sha = hashlib.sha256(blob).digest()
    nseq, nfr = args.measure_identities, args.measure_frames
    t_int = 0.0
    for s in range(nseq):
        ref, rk, frames, kps = sequence(50_000_000 + 7 * s, nfr, size)
        rq = kp_q4(rk)
        ident = hashlib.sha256(b"identity" + struct.pack("<I", s)).digest()
        dq_all = kp_q4(kps)
        for name, src in (("truth", dq_all),):
            for step in (4, 1):
                tot = 0
                for f0 in range(0, nfr, 25):
                    tot += len(encode_track(src[f0:f0 + 25].astype(np.float64), step, gen_sha, ident, f0))
                bitrate[f"{name}_step{step}"].append(tot * 8 / (nfr / 25) / 1000)
        with torch.no_grad():
            dk = det(torch.from_numpy(np.ascontiguousarray(frames.transpose(0, 3, 1, 2)).astype(np.float32)).to(dev)).cpu().numpy()
        for step in (4, 1):
            tot = 0
            for f0 in range(0, nfr, 25):
                tot += len(encode_track(kp_q4(dk[f0:f0 + 25]).astype(np.float64), step, gen_sha, ident, f0))
            bitrate[f"detected_step{step}"].append(tot * 8 / (nfr / 25) / 1000)
        ints, floats = [], []
        for i in range(nfr):
            dq = dq_all[i]
            warped, heat = warp_and_maps(ref, rq, dq, gauss, bg)
            t1 = time.time()
            yi = int_forward(ilayers, warped, ref, heat)
            t_int += time.time() - t1
            with torch.no_grad():
                f = lambda a: torch.from_numpy(a.transpose(2, 0, 1)[None].astype(np.float32)).to(dev)  # noqa: E731
                yf = fnet(f(warped), f(ref), torch.from_numpy(heat[None].astype(np.float32)).to(dev))[0]
            yf = np.clip(np.floor(yf[0].cpu().numpy().transpose(1, 2, 0) + 0.5), 0, 255).astype(np.uint8)
            truth = frames[i]
            for key, (a, b) in {"int_float": (yi, yf), "int_truth": (yi, truth), "float_truth": (yf, truth), "warp_truth": (warped, truth)}.items():
                rows[key].append(psnr(a, b))
                ss[key].append(ssim(a, b))
            ints.append(yi)
            floats.append(yf)
        with torch.no_grad():
            for key, imgs in (("int", np.stack(ints)), ("float", np.stack(floats)), ("truth", frames)):
                d = det(torch.from_numpy(np.ascontiguousarray(imgs.transpose(0, 3, 1, 2)).astype(np.float32)).to(dev)).cpu().numpy()
                err = np.sqrt(((d - kps) ** 2).sum(-1))
                kerr[key].append(float(np.sqrt((err ** 2).mean())))
                mouth[key].append(float(np.sqrt((err[:, 7:11] ** 2).mean())))
        print(f"measured sequence {s + 1}/{nseq}", flush=True)
    finite = lambda v: [x for x in v if math.isfinite(x)]  # noqa: E731
    report["measurement"] = {
        "identities": nseq, "frames_each": nfr, "fps": 25, "size": [size, size],
        "psnr_db": {k: {"mean": round(float(np.mean(finite(v))), 2), "min": round(float(np.min(v)), 2), "identical_frames": int(sum(1 for x in v if not math.isfinite(x)))} for k, v in rows.items()},
        "ssim": {k: {"mean": round(float(np.mean(v)), 4), "min": round(float(np.min(v)), 4)} for k, v in ss.items()},
        "keypoint_rms_px": {k: round(float(np.mean(v)), 3) for k, v in kerr.items()},
        "mouth_keypoint_rms_px": {k: round(float(np.mean(v)), 3) for k, v in mouth.items()},
        "track_kbit_s": {k: round(float(np.mean(v)), 3) for k, v in bitrate.items()},
        "numpy_int_seconds_per_frame": round(t_int / (nseq * nfr), 4),
    }


def corpus_inputs(args, out: Path, report: dict, blob: bytes):
    """Identity priors and keypoint tracks of held-out faces, for the corpus (§9)."""
    gen = hashlib.sha256(blob).hexdigest()
    (out / "identities").mkdir(exist_ok=True)
    (out / "tracks").mkdir(exist_ok=True)
    files = {}
    for name, seed, nfr in (("face1", 70_000_001, args.corpus_frames), ("face2", 70_000_017, args.corpus_frames)):
        ref, rk, frames, kps = sequence(seed, nfr, args.size)
        rq = kp_q4(rk)
        wfi = bytearray(b"WFI1")
        wfi += struct.pack("<BBHH", 1, K, args.size, args.size)
        wfi += rq.astype("<i2").tobytes()
        wfi += ref.tobytes()
        (out / "identities" / f"{name}.wfi").write_bytes(bytes(wfi))
        track = {"identity": name, "keypoints": kp_q4(kps).tolist()}
        (out / "tracks" / f"{name}.json").write_text(json.dumps(track, separators=(",", ":")) + "\n")
        files[name] = {"wfi_sha256": hashlib.sha256(bytes(wfi)).hexdigest(), "frames": nfr}
    report["corpus_inputs"] = {"generator": gen, "faces": files}


def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--out", type=Path, required=True)
    ap.add_argument("--seed", type=int, default=20261022)
    ap.add_argument("--size", type=int, default=64)
    ap.add_argument("--c1", type=int, default=24)
    ap.add_argument("--c2", type=int, default=48)
    ap.add_argument("--c3", type=int, default=64)
    ap.add_argument("--gauss", type=int, default=1024, help="§7.5's g: 32768/σ², σ in pixels")
    ap.add_argument("--bg", type=int, default=6554, help="§7.5's background weight, Q16")
    ap.add_argument("--identities", type=int, default=6000)
    ap.add_argument("--steps", type=int, default=60000)
    ap.add_argument("--qat-steps", type=int, default=12000)
    ap.add_argument("--lr", type=float, default=1e-3)
    ap.add_argument("--batch", type=int, default=32)
    ap.add_argument("--workers", type=int, default=12)
    ap.add_argument("--det-steps", type=int, default=12000)
    ap.add_argument("--det-pool", type=int, default=6000, help="training identities the detector reads, all their poses")
    ap.add_argument("--measure-identities", type=int, default=12)
    ap.add_argument("--measure-frames", type=int, default=100)
    ap.add_argument("--corpus-frames", type=int, default=3)
    ap.add_argument("--preview", action="store_true", help="only write preview pictures of the renderer and the warp")
    ap.add_argument("--measure-only", action="store_true", help="measure the generator and float network an earlier run wrote to --out")
    args = ap.parse_args()
    args.out.mkdir(parents=True, exist_ok=True)
    if args.preview:
        preview(args)
        return
    report = {"seed": args.seed, "args": {k: v for k, v in vars(args).items() if k != "out"}, "keypoints": KEYPOINTS}
    fnet, blob = train(args, args.out, report)
    corpus_inputs(args, args.out, report, blob)
    measure(args, args.out, report, fnet, blob)
    (args.out / "report.json").write_text(json.dumps(report, indent=1, sort_keys=True) + "\n")
    print(json.dumps(report["measurement"], indent=1))


def preview(args):
    """A contact sheet of generated faces and warps, as a binary PPM (no image library)."""
    tiles = []
    for i in range(8):
        ref, warped, heat, tgt, rq, tq = make_sample(1000 + i, args.size, args.gauss, args.bg)
        hm = np.clip(heat.max(0) * 255 // 256, 0, 255).astype(np.uint8)
        tiles.append(np.concatenate([ref, warped, tgt, np.stack([hm] * 3, -1)], 1))
    sheet = np.concatenate(tiles, 0)
    big = np.repeat(np.repeat(sheet, 3, 0), 3, 1)
    h, w, _ = big.shape
    (args.out / "preview.ppm").write_bytes(b"P6 %d %d 255\n" % (w, h) + big.tobytes())
    _, _, frames, kps = sequence(50_000_000, 16, args.size)
    seq = np.concatenate(list(frames), 1)
    seq = np.repeat(np.repeat(seq, 3, 0), 3, 1)
    h, w, _ = seq.shape
    (args.out / "sequence.ppm").write_bytes(b"P6 %d %d 255\n" % (w, h) + seq.tobytes())
    print("wrote", args.out / "preview.ppm", args.out / "sequence.ppm")


if __name__ == "__main__":
    main()