#!/usr/bin/env python3
"""Independent verifier of the faces-from-keypoints corpus (extensions/avatars.md §12).

Written from the extension's text. It shares no code with wai-rs, which generates the corpus.
It reads the three files of §7.2 to §7.4 with readers of its own; runs the driver of §7.5 and the
generator's arithmetic of §7.6 in Python's integers, operator by operator as §7.6 states them
(no operator set, no graph); verifies the origin claims, the player's log and the label policy
with the synthesis corpus's second implementation (`synthesis_verify.py`); verifies every
signature with WAI's strict Ed25519 rule (`claim-conformance/_ed25519.py`) and every link hash
with BLAKE3 written from its specification (`claim-conformance/_blake3.py`); reads §10's codes
from the extension; and implements the records of §8 and the steps of §9.2.

For every vector it requires what the corpus states: each file's reading, each frame's
digests (and, where the corpus gives them, its samples), each case's outcome, reason, place and
statements. Then it mutates copies of the corpus, each with one defect, and requires every copy
to be refused somewhere.

    python3 -I avatar_face_verify.py <live-conformance/avatar-face>

Exit status 0 only if every vector verifies and every mutated copy fails.
"""
import copy
import hashlib
import json
import os
import struct
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
WAI = os.path.normpath(os.path.join(HERE, ".."))
sys.path.insert(0, HERE)
sys.path.insert(0, os.path.join(WAI, "claim-conformance"))
import synthesis_verify as SV  # noqa: E402
import _ed25519 as ED  # noqa: E402
from _blake3 import blake3  # noqa: E402

EXTENSION = os.path.join(WAI, "extensions", "avatars.md")
REGISTRY = os.path.join(WAI, "registry", "capabilities.json")
CAPABILITY = "wai.avatar.face"
SIGN_DOMAIN = b"wai:avatar\x01"            # §8.1
LINK_DOMAIN = b"wai:avatar-link\x01"       # §8.1
OBJECT_DOMAIN = b"moq-jwp:object\x01"      # jwp-receipts §2.1
HEADER = 84                                # §7.4

Refuse, Bad, need, hexb, integer, members = SV.Refuse, SV.Bad, SV.need, SV.hexb, SV.integer, SV.members


def sha(b):
    return hashlib.sha256(b).digest()


class Form(Exception):
    """A file that does not read (§7.2 to §7.4), with its code."""

    def __init__(self, code):
        super().__init__(code)
        self.code = code


# ---------------------------------------------------------------- §7.2 to §7.4 the files

def u16(b, at):
    return b[at] | b[at + 1] << 8


def i16(b, at):
    v = u16(b, at)
    return v - 65536 if v >= 32768 else v


def side_ok(v):
    return 4 <= v <= 512 and v % 4 == 0


def have(b, at, n):
    if len(b) < at + n:
        raise Form("truncated")


def shapes(k, c1, c2, c3):
    """§7.2: (cin, cout, stride) of the ten convolutions, in the file's order."""
    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 read_generator(b):
    """§7.2, in its order."""
    have(b, 0, 4)
    if b[:4] != b"WFG1":
        raise Form("bad_magic")
    have(b, 4, 1)
    if b[4] != 1:
        raise Form("version")
    have(b, 5, 17)
    k, width, height, c1, c2, c3, gauss = b[5], u16(b, 6), u16(b, 8), u16(b, 10), u16(b, 12), u16(b, 14), u16(b, 16)
    bg = struct.unpack_from("<I", b, 18)[0]
    if not (1 <= k <= 32 and side_ok(width) and side_ok(height) and all(1 <= c <= 64 for c in (c1, c2, c3)) and gauss >= 1 and bg >= 1):
        raise Form("field_range")
    at, convs = 22, []
    for cin, cout, stride in shapes(k, c1, c2, c3):
        nw = cout * cin * 9
        have(b, at, 2 * nw + 4 * cout)
        w = list(struct.unpack_from("<%dh" % nw, b, at))
        at += 2 * nw
        bias = list(struct.unpack_from("<%di" % cout, b, at))
        at += 4 * cout
        convs.append({"cin": cin, "cout": cout, "stride": stride, "w": w, "b": bias})
    if at != len(b):
        raise Form("trailing")
    return {"k": k, "width": width, "height": height, "c1": c1, "c2": c2, "c3": c3, "gauss": gauss, "bg": bg, "convs": convs}


def read_identity(b):
    """§7.3, in §7.2's order."""
    have(b, 0, 4)
    if b[:4] != b"WFI1":
        raise Form("bad_magic")
    have(b, 4, 1)
    if b[4] != 1:
        raise Form("version")
    have(b, 5, 5)
    k, width, height = b[5], u16(b, 6), u16(b, 8)
    if not (1 <= k <= 32 and side_ok(width) and side_ok(height)):
        raise Form("field_range")
    n = 4 * k + width * height * 3
    have(b, 10, n)
    kp = [(i16(b, 10 + 4 * j), i16(b, 12 + 4 * j)) for j in range(k)]
    if len(b) != 10 + n:
        raise Form("trailing")
    return {"k": k, "width": width, "height": height, "keypoints": kp, "rgb": bytes(b[10 + 4 * k:])}


def read_segment(b):
    """§7.4, in its order."""
    have(b, 0, 4)
    if b[:4] != b"WFK1":
        raise Form("bad_magic")
    have(b, 4, 1)
    if b[4] != 1:
        raise Form("version")
    have(b, 5, HEADER - 5)
    k, step, order = b[5], b[6], b[7]
    first = struct.unpack_from("<I", b, 72)[0]
    n, num, den, reserved = u16(b, 76), u16(b, 78), u16(b, 80), u16(b, 82)
    if not (1 <= k <= 32 and 1 <= step <= 16 and order <= 15 and 1 <= n <= 1024 and num >= 1 and den >= 1 and reserved == 0):
        raise Form("field_range")
    have(b, HEADER, 4 * k)
    cur = [i16(b, HEADER + 2 * j) for j in range(2 * k)]

    def check(c):
        if not -32768 <= c * step <= 32767:
            raise Form("field_range")
    for c in cur:
        check(c)
    frames = [[(cur[2 * j] * step, cur[2 * j + 1] * step) for j in range(k)]]
    data, pos = b[HEADER + 4 * k:], [0]

    def bit():
        if pos[0] // 8 >= len(data):
            raise Form("truncated")
        v = (data[pos[0] // 8] >> (7 - pos[0] % 8)) & 1
        pos[0] += 1
        return v

    def bits(m):
        v = 0
        for _ in range(m):
            v = v << 1 | bit()
        return v

    def code():
        z = 0
        while bit() == 0:
            z += 1
            if z > 31:
                raise Form("field_range")
        m = (1 << z) | bits(z)
        return (m - 1) << order | bits(order)

    for _ in range(1, n):
        for i in range(2 * k):
            u = code()
            cur[i] += u >> 1 if u % 2 == 0 else -((u + 1) >> 1)
            check(cur[i])
        frames.append([(cur[2 * j] * step, cur[2 * j + 1] * step) for j in range(k)])
    if HEADER + 4 * k + (pos[0] + 7) // 8 != len(b):
        raise Form("trailing")
    while pos[0] % 8:
        if bit():
            raise Form("noncanonical")
    return {"k": k, "step": step, "order": order, "generator": bytes(b[8:40]), "identity": bytes(b[40:72]), "first": first,
            "rate": [num, den], "keypoints": frames}


# ---------------------------------------------------------------- §7.5 the driver

def exp_neg_q16(x):
    """integer-operators §2."""
    u = -x
    n = u // 45426
    r = x + 45426 * n
    t = r + 88670
    q = (t * t) >> 16
    y = ((23495 * q) >> 16) + 22544
    return 0 if n >= 63 else y >> n


def rdiv(n, d):
    """sgn(n)·⌊(2|n| + d) / (2d)⌋, d ≥ 1."""
    a = (2 * abs(n) + d) // (2 * d)
    return -a if n < 0 else a


def warp(ident, drv, gauss, bg):
    """§7.5: the warped picture (RGB8 bytes) and the maps ([K][h][w] integers, flat)."""
    w, h, k, ref, src = ident["width"], ident["height"], ident["k"], ident["rgb"], ident["keypoints"]
    out = bytearray(w * h * 3)
    maps = [0] * (k * w * h)
    for y in range(h):
        py = 16 * y + 8
        for x in range(w):
            px = 16 * x + 8
            den, nx, ny = bg, 0, 0
            for j in range(k):
                dx, dy = px - drv[j][0], py - drv[j][1]
                wj = exp_neg_q16(-(((dx * dx + dy * dy) * gauss + 128) >> 8))
                den += wj
                nx += wj * (src[j][0] - drv[j][0])
                ny += wj * (src[j][1] - drv[j][1])
                maps[(j * h + y) * w + x] = wj >> 8
            ux, uy = 16 * x + rdiv(nx, den), 16 * y + rdiv(ny, den)
            ix, iy = ux // 16, uy // 16
            fx, fy = ux - 16 * ix, uy - 16 * iy
            x0, x1 = min(max(ix, 0), w - 1), min(max(ix + 1, 0), w - 1)
            y0, y1 = min(max(iy, 0), h - 1), min(max(iy + 1, 0), h - 1)
            for c in range(3):
                v = ((16 - fx) * (16 - fy) * ref[(y0 * w + x0) * 3 + c] + fx * (16 - fy) * ref[(y0 * w + x1) * 3 + c]
                     + (16 - fx) * fy * ref[(y1 * w + x0) * 3 + c] + fx * fy * ref[(y1 * w + x1) * 3 + c] + 128) >> 8
                out[(y * w + x) * 3 + c] = v
    return bytes(out), maps


# ---------------------------------------------------------------- §7.6 the generator's arithmetic

def sigmoid_entry(v):
    """§7.6: T[v + 32768]."""
    if v == -32768:
        return 0
    if v < 0:
        return 32768 - sigmoid_entry(-v)
    d = 65536 + min(exp_neg_q16(-16 * v), 65536)
    return ((1 << 31) + d // 2) // d


TABLE = {}


def table(v):
    if v not in TABLE:
        TABLE[v] = sigmoid_entry(v)
    return TABLE[v]


def conv(planes, h, w, layer):
    """A 3 × 3 convolution, zero padding 1, of planes (each a flat list of h·w integers), each
    output R_half_up((sum + bias), 12)."""
    cin, cout, stride, wt, bias = layer["cin"], layer["cout"], layer["stride"], layer["w"], layer["b"]
    wp = w + 2
    n = h * wp
    padded = []
    for p in planes:
        q = [0] * ((h + 2) * wp + 2)
        for y in range(h):
            q[(y + 1) * wp + 1:(y + 1) * wp + 1 + w] = p[y * w:(y + 1) * w]
        padded.append(q)
    ho, wo = (h - 1) // stride + 1, (w - 1) // stride + 1
    out = []
    if stride == 1:
        # Each tap reads the padded plane shifted by its offset; the sums run over the
        # padded width, whose two extra columns are dropped at the end.
        for o in range(cout):
            acc = [bias[o]] * n
            for ci in range(cin):
                src = padded[ci]
                base = (o * cin + ci) * 9
                for t in range(9):
                    v = wt[base + t]
                    if v:
                        off = (t // 3) * wp + t % 3
                        acc = [a + v * s for a, s in zip(acc, src[off:off + n])]
            out.append([(acc[y * wp + x] + 2048) >> 12 for y in range(ho) for x in range(wo)])
        return out, ho, wo
    # Stride 2: each tap's samples gathered once per input channel.
    taps = [[[src[(stride * y + t // 3) * wp + stride * x + t % 3] for y in range(ho) for x in range(wo)] for t in range(9)] for src in padded]
    for o in range(cout):
        acc = [bias[o]] * (ho * wo)
        for ci in range(cin):
            base = (o * cin + ci) * 9
            for t in range(9):
                v = wt[base + t]
                if v:
                    acc = [a + v * s for a, s in zip(acc, taps[ci][t])]
        out.append([(a + 2048) >> 12 for a in acc])
    return out, ho, wo


def relu(planes):
    return [[min(max(v, 0), 32767) for v in p] for p in planes]


def shuffle(planes, h, w):
    """pixel_shuffle(2): channel 4c + 2i + j to row offset i and column offset j of channel c."""
    out = []
    for c in range(len(planes) // 4):
        p = [0] * (4 * h * w)
        for i in range(2):
            for j in range(2):
                src = planes[4 * c + 2 * i + j]
                for y in range(h):
                    p[(2 * y + i) * 2 * w + j:(2 * y + i) * 2 * w + 2 * w:2] = src[y * w:(y + 1) * w]
        out.append(p)
    return out


def generate_frame(gen, ident, drv):
    """§7.7: one frame, RGB8 bytes, and the warped picture and maps it was made from."""
    w, h, k = gen["width"], gen["height"], gen["k"]
    warped, maps = warp(ident, drv, gen["gauss"], gen["bg"])
    n = w * h
    t3 = [[(warped[i * 3 + c] - 128) * 256 for i in range(n)] for c in range(3)]
    t4 = [[(ident["rgb"][i * 3 + c] - 128) * 256 for i in range(n)] for c in range(3)]
    t5 = [[v * 128 for v in maps[j * n:(j + 1) * n]] for j in range(k)]
    cv = gen["convs"]
    e1, h1, w1 = conv(t3 + t4 + t5, h, w, cv[0])
    e1 = relu(e1)
    e2, h2, w2 = conv(e1, h1, w1, cv[1])
    e2 = relu(e2)
    e3, h3, w3 = conv(e2, h2, w2, cv[2])
    e3 = relu(e3)
    m1, _, _ = conv(e3, h3, w3, cv[3])
    m1 = relu(m1)
    m2, _, _ = conv(m1, h3, w3, cv[4])
    r = relu([[a + b for a, b in zip(p, q)] for p, q in zip(m2, e3)])
    u2, _, _ = conv(r, h3, w3, cv[5])
    u2 = relu(shuffle(u2, h3, w3))
    d2, _, _ = conv(u2 + e2, h2, w2, cv[6])
    d2 = relu(d2)
    u1, _, _ = conv(d2, h2, w2, cv[7])
    u1 = relu(shuffle(u1, h2, w2))
    d1, _, _ = conv(u1 + e1, h1, w1, cv[8])
    d1 = relu(d1)
    o, _, _ = conv(d1, h1, w1, cv[9])
    out = bytearray(3 * n)
    for i in range(n):
        m = table(min(max(o[3][i], -32767), 32767))
        for c in range(3):
            g = min(max(o[c][i], -32768), 32767)
            a = (t3[c][i] * m + 16384) >> 15
            b = (g * (32768 - m) + 16384) >> 15
            out[i * 3 + c] = min(max(((a + b + 128) >> 8) + 128, 0), 255)
    return bytes(out), warped, maps


FRAMES = {}


def frame_of(gen_bytes, ident_bytes, drv):
    """A frame, kept by the generator's and the identity's digests and the keypoints."""
    key = (sha(gen_bytes), sha(ident_bytes), tuple(drv))
    if key not in FRAMES:
        FRAMES[key] = generate_frame(read_generator(gen_bytes), read_identity(ident_bytes), drv)
    return FRAMES[key]


def maps_bytes(maps):
    return b"".join(struct.pack("<H", v) for v in maps)


# ---------------------------------------------------------------- §8 records

def read(text, kind, names, fields):
    try:
        v = SV.strict_json(text)
        members(v, names)
        need(v["kind"] == kind)
        signer, sig = hexb(v["signer_pubkey"], 32), hexb(v["sig"], 64)
        SV.date_time(v["issued_at"])
        r = fields(v)
    except (Bad, KeyError, TypeError, AttributeError, ValueError):
        raise Refuse("record_malformed")
    bad = ED.verify(signer, SIGN_DOMAIN + SV.jcs({k: x for k, x in v.items() if k != "sig"}).encode("utf-8"), sig)
    if bad:
        raise Refuse(bad)
    r.update(link=blake3(LINK_DOMAIN + SV.jcs(v).encode("utf-8")), signer=signer)
    return r


IDENTITY = ["kind", "identity", "issued_at", "signer_pubkey", "sig"]
TRACK = ["kind", "object", "identity", "seq", "prev", "frames", "issued_at", "signer_pubkey", "sig"]


def identity_fields(v):
    return {"identity": hexb(v["identity"], 32)}


def track_fields(v):
    f = v["frames"]
    if f is not None:
        need(isinstance(f, list) and 1 <= len(f) <= 1024)
        f = [hexb(x, 32) for x in f]
    return {"object": hexb(v["object"], 32), "identity": hexb(v["identity"], 32), "seq": integer(v["seq"]),
            "prev": None if v["prev"] is None else hexb(v["prev"], 32), "frames": f}


def envelope(b):
    """A WAI1 envelope's manifest and payload, or None."""
    if len(b) < 8 or b[:4] != b"WAI1":
        return None
    n = int.from_bytes(b[4:8], "little")
    try:
        m = json.loads(b[8:8 + n])
    except ValueError:
        return None
    at = 8 + n
    if len(b) < at + 4:
        return None
    ln = int.from_bytes(b[at:at + 4], "little")
    if at + 4 + ln != len(b):
        return None
    return m, b[at + 4:]


# ---------------------------------------------------------------- §9 verification

def verify(case, classes, files):
    """§9.2: the statements of §9.3, or Refuse with its reason and place."""
    def each(key, kind, names, fields, place):
        out = []
        for i, t in enumerate(case[key]):
            try:
                out.append(read(t, kind, names, fields))
            except Refuse as e:
                raise Refuse(e.reason, [place, i])
        return out

    def held(lst):
        out = []
        for h in lst:
            out.append(files[h["file"]] if "file" in h else bytes.fromhex(h["hex"]))
        return out

    # Step 1.
    ids = each("identities", "wai-face-identity", IDENTITY, identity_fields, "identity")
    tracks = each("tracks", "wai-face-track", TRACK, track_fields, "track")
    objects = {blake3(OBJECT_DOMAIN + bytes.fromhex(o)): bytes.fromhex(o) for o in case["objects"]}
    generators = {sha(g): g for g in held(case["generators"])}
    priors = {sha(p): p for p in held(case["priors"])}
    st = {"identities": [], "tracks": []}

    # Step 2.
    readable = {}
    for i, r in enumerate(ids):
        state = "not-held"
        if r["identity"] in priors:
            try:
                read_identity(priors[r["identity"]])
            except Form:
                raise Refuse("identity_invalid", ["identity", i])
            readable[r["identity"]] = priors[r["identity"]]
            state = "read"
        st["identities"].append({"link": r["link"].hex(), "identity": r["identity"].hex(), "prior": state})
    by_link = {r["link"]: r for r in ids}

    # Step 3.
    decoded, last, generator_of = [], {}, {}
    for i, t in enumerate(tracks):
        def no(reason):
            raise Refuse(reason, ["track", i])
        idr = by_link.get(t["identity"])
        if idr is None:
            no("identity_unknown")
        if t["signer"] != idr["signer"]:
            no("driver_key_mismatch")
        prev = last.get(t["identity"])
        if prev is None:
            chained = t["seq"] == 0 and t["prev"] is None
        else:
            chained = t["seq"] == prev[0] + 1 and t["prev"] == prev[1]
        if not chained:
            no("track_chain_broken")
        last[t["identity"]] = (t["seq"], t["link"])
        env = objects.get(t["object"])
        if env is None:
            no("object_unknown")
        parts = envelope(env)
        if parts is None or not isinstance(parts[0], dict) or parts[0].get("model_requirement", {}).get("capability") != CAPABILITY:
            no("payload_invalid")
        manifest, payload = parts
        try:
            seg = read_segment(payload)
        except Form:
            no("payload_invalid")
        pin = ((manifest.get("conditioning") or {}).get("prior") or {}).get("sha256")
        if not isinstance(pin, str) or pin != seg["generator"].hex():
            no("generator_mismatch")
        if seg["identity"] != idr["identity"]:
            no("identity_mismatch")
        generator_of[t["object"]] = seg["generator"]
        gen = generators.get(seg["generator"])
        if gen is not None:
            try:
                g = read_generator(gen)
            except Form:
                no("generator_invalid")
        made = None
        if gen is not None and idr["identity"] in readable:
            ident = read_identity(readable[idr["identity"]])
            if not (g["k"] == ident["k"] == seg["k"] and g["width"] == ident["width"] and g["height"] == ident["height"]):
                no("shape_mismatch")
            frames = [frame_of(gen, readable[idr["identity"]], kp)[0] for kp in seg["keypoints"]]
            if t["frames"] is not None and t["frames"] != [sha(f) for f in frames]:
                no("frame_mismatch")
            made = b"".join(frames)
        if case["policy"]["generation"] == "require-proof" and made is None:
            no("generation_not_proved")
        st["tracks"].append({"link": t["link"].hex(), "object": t["object"].hex(), "frames": len(seg["keypoints"]),
                             "status": "proved" if made is not None else "not-recomputed", "output": None if made is None else sha(made).hex()})
        if made is not None:
            decoded.append({"object": t["object"].hex(), "samples": made.hex()})

    # Step 4: synthesis.
    pol = case["policy"]
    syn = {"claims": case["claims"], "markings": [], "detections": [], "reviews": case["reviews"], "label_policy": case["label_policy"],
           "player": case["player"], "objects": case["objects"], "decoded": decoded, "priors": [], "unmarked": [], "room": None,
           "policy": {"dual_layer": pol["dual_layer"], "assistive": pol["assistive"], "second_layer": pol["second_layer"]}}
    try:
        syn_st = SV.verify(syn, classes)
    except Refuse as e:
        raise Refuse(e.reason, ["synthesis", e.at])
    st["claims"] = [{"link": c["link"], "class": c["class"]} for c in syn_st["claims"]]
    st["labels"] = syn_st["labels"]
    st["reviews_presented"] = syn_st["reviews_presented"]
    st["visual_icon"] = syn_st["visual_icon"]

    # Step 5: each segment's claim.
    claims = [SV.read(t, "wai-origin", SV.ORIGIN, SV.origin_fields) for t in case["claims"]]
    for i, t in enumerate(tracks):
        mine = [c for c in claims if c["subject"] == ("object", t["object"])]
        if not mine:
            raise Refuse("claim_missing", ["track", i])
        c = mine[0]
        if not (c["class"] == "synthesized" and c["step"] == "generate" and generator_of[t["object"]] in [m["digest"] for m in c["models"]]
                and c["record"] == t["link"]):
            raise Refuse("claim_mismatch", ["track", i])
    return st


# ---------------------------------------------------------------- the extension's tables

with open(EXTENSION, encoding="utf-8") as _f:
    _TEXT = _f.read()
_TABLES = [b for b in SV.section(_TEXT, "## 10. Reasons").split("\n\n") if any(l.startswith("| `") for l in b.splitlines())]
REASONS = [l.split("`")[1] for l in _TABLES[0].splitlines() if l.startswith("| `")]
FORM_CODES = [l.split("`")[1] for l in _TABLES[1].splitlines() if l.startswith("| `")]


def form_outcome(v):
    b = bytes.fromhex(v["hex"])
    reader = {"generator": read_generator, "identity": read_identity, "segment": read_segment}[v["form"]]
    try:
        r = reader(b)
    except Form as e:
        return {"outcome": "refuse", "reason": e.code}
    if v["form"] == "generator":
        r = {k: r[k] for k in ("k", "width", "height", "c1", "c2", "c3", "gauss", "bg")}
    elif v["form"] == "identity":
        r = {"k": r["k"], "width": r["width"], "height": r["height"], "keypoints": [list(p) for p in r["keypoints"]], "rgb_sha256": sha(r["rgb"]).hex()}
    else:
        r = dict(r, generator=r["generator"].hex(), identity=r["identity"].hex(), keypoints=[[list(p) for p in f] for f in r["keypoints"]])
    return {"outcome": "read", "read": r}


def table_outcome():
    """§7.6's table: its SHA-256 (each entry a 16-bit little-endian signed integer, T[0] first)
    and its entries at the values the corpus names."""
    entries = [sigmoid_entry(v) for v in range(-32768, 32768)]
    return {"sha256": sha(b"".join(struct.pack("<h", e) for e in entries)).hex()}


def frame_outcome(v, files):
    kp = [tuple(p) for p in v["keypoints"]]
    f, warped, maps = frame_of(files[v["generator"]], files[v["identity"]], kp)
    got = {"warped_sha256": sha(warped).hex(), "maps_sha256": sha(maps_bytes(maps)).hex(), "frame_sha256": sha(f).hex()}
    if "frame" in v:
        got.update(warped=warped.hex(), frame=f.hex())
    return got


def case_outcome(v, classes, files):
    try:
        return {"outcome": "accept", "statements": verify(v, classes, files)}
    except Refuse as e:
        return {"outcome": "refuse", "reason": e.reason, "at": e.at}


def stated(v, keys):
    return {k: v[k] for k in keys if k in v}


def check(doc, classes, files):
    """Failures of a corpus: an input other than the one pinned, or a vector that does not get
    what it states."""
    fails = []
    for e in doc["inputs"] + doc["written"]:
        if e["file"] not in files or sha(files[e["file"]]).hex() != e["sha256"]:
            fails.append("the file %s is not the one vectors.json pins" % e["file"])
    t = doc["mask_table"]
    if table_outcome()["sha256"] != t["sha256"] or any(sigmoid_entry(v) != e for v, e in t["at"]):
        fails.append("the mask's table")
    for v in doc["forms"]:
        try:
            got = form_outcome(v)
        except Exception as e:  # noqa: BLE001 — a malformed copy fails, whatever it breaks
            got = {"outcome": "unreadable: %r" % (e,)}
        if got != stated(v, ("outcome", "reason", "read")):
            fails.append("form %s: %s" % (v["name"], json.dumps(got)[:200]))
    for v in doc["frames"]:
        try:
            got = frame_outcome(v, files)
        except Exception as e:  # noqa: BLE001
            got = {"outcome": "unreadable: %r" % (e,)}
        if got != stated(v, ("warped_sha256", "maps_sha256", "frame_sha256", "warped", "frame")):
            fails.append("frame %s" % v["name"])
    for v in doc["cases"]:
        try:
            got = case_outcome(v, classes, files)
        except Exception as e:  # noqa: BLE001
            got = {"outcome": "unreadable: %r" % (e,)}
        if got != stated(v, ("outcome", "reason", "at", "statements")):
            fails.append("%s: %s" % (v["name"], json.dumps(got)[:300]))
    return fails


def flip_hex_digit(s, at):
    return s[:at] + ("0" if s[at] != "0" else "1") + s[at + 1:]


def mutations(doc, files):
    """Copies of the corpus, each with one defect a conforming verifier must catch."""
    cn = {v["name"]: k for k, v in enumerate(doc["cases"])}
    fn = {v["name"]: k for k, v in enumerate(doc["frames"])}
    on = {v["name"]: k for k, v in enumerate(doc["forms"])}

    def mutate(label, f, fs=None):
        d = copy.deepcopy(doc)
        f(d)
        return label, d, fs or files

    def relisted(label, name, at):
        """A byte of a file changed and its pin updated to the changed bytes, so that only what
        the file computes can tell."""
        fs = dict(files)
        b = bytearray(fs[name])
        b[at] ^= 0x01
        fs[name] = bytes(b)
        d = copy.deepcopy(doc)
        for e in d["inputs"] + d["written"]:
            if e["file"] == name:
                e["sha256"] = sha(fs[name]).hex()
        return label, d, fs

    def first_refusal(d):
        return next(k for k, v in enumerate(d["cases"]) if v["outcome"] == "refuse")

    first0 = fn["face1_0"]
    return [
        relisted("a weight of the seeded generator, its pin updated", "seeded.wfg", 200),
        relisted("a bias of the seeded generator, its pin updated", "seeded.wfg", len(files["seeded.wfg"]) - 3),
        relisted("a weight of the trained generator's last convolution, its pin updated", "generator.wfg", len(files["generator.wfg"]) - 40),
        relisted("a sample of the seeded identity's picture, its pin updated", "seeded.wfi", 60),
        relisted("a keypoint of the seeded identity, its pin updated", "seeded.wfi", 10),
        mutate("a byte of the first face's prior", lambda d: None, dict(files, **{"identities/face1.wfi": files["identities/face1.wfi"][:-1] + bytes([files["identities/face1.wfi"][-1] ^ 1])})),
        mutate("a keypoint of the first face's first frame moved by a sixteenth", lambda d: d["frames"][first0]["keypoints"][6].__setitem__(0, d["frames"][first0]["keypoints"][6][0] + 1)),
        mutate("a digit of the seeded frame's samples", lambda d: d["frames"][fn["seeded_moved"]].update(frame=flip_hex_digit(d["frames"][fn["seeded_moved"]]["frame"], 100))),
        mutate("the mask's table entry at 0 stated as 16382", lambda d: [a.__setitem__(1, 16382) for a in d["mask_table"]["at"] if a[0] == 0]),
        mutate("a segment's padding read as data", lambda d: d["forms"][on["segment_padding_not_zero"]].update(outcome="read")),
        mutate("the first segment stated not recomputed",
               lambda d: d["cases"][cn["one_face_two_segments"]]["statements"]["tracks"][0].update(status="not-recomputed")),
        mutate("the first segment's output stated as another",
               lambda d: d["cases"][cn["one_face_two_segments"]]["statements"]["tracks"][0].update(output=flip_hex_digit(d["cases"][cn["one_face_two_segments"]]["statements"]["tracks"][0]["output"], 5))),
        mutate("the generator withheld from the frames-unstated case", lambda d: d["cases"][cn["frames_unstated"]].update(generators=[])),
        mutate("the second participant's track signed over the first one's face",
               lambda d: d["cases"][cn["two_faces"]]["tracks"].__setitem__(1, d["cases"][cn["driver_key_mismatch"]]["tracks"][0])),
        mutate("the two segments' track records swapped",
               lambda d: d["cases"][cn["one_face_two_segments"]].update(tracks=d["cases"][cn["one_face_two_segments"]]["tracks"][::-1])),
        mutate("the label policy withheld from an unlabelled presentation", lambda d: d["cases"][cn["label_not_shown"]].update(label_policy=None)),
        mutate("the first refusal stated at another place", lambda d: d["cases"][first_refusal(d)].update(at=["track", 7])),
    ]


def read_files(corpus_dir):
    files = {}
    for sub in ("", "identities", "tracks"):
        d = os.path.join(corpus_dir, sub)
        for n in sorted(os.listdir(d)):
            p = os.path.join(d, n)
            if os.path.isdir(p) or (sub == "" and n == "README.md"):
                continue
            with open(p, "rb") as f:
                files[n if sub == "" else sub + "/" + n] = f.read()
    return files


def main(corpus_dir):
    fails = []
    st = ED._self_test()
    print(("PASS" if not st else "FAIL") + " Ed25519 strict rule self-test " + ",".join(st))
    if st:
        fails.append("ed25519 self-test")
    with open(REGISTRY, encoding="utf-8") as f:
        classes = {c["capability"]: c["replicate"] for c in json.load(f)["capabilities"]}
    files = read_files(corpus_dir)
    doc = json.loads(files["vectors.json"])
    listed = {e["file"] for e in doc["inputs"] + doc["written"]} | {"vectors.json"}
    if set(files) != listed:
        fails.append("the corpus holds other files than vectors.json names: %s" % sorted(set(files) ^ listed))
    print("%s the corpus holds exactly the files vectors.json names" % ("PASS" if set(files) == listed else "FAIL"))
    if doc["capability"] != CAPABILITY or classes.get(CAPABILITY) != "generative-exact":
        fails.append("the capability is not registered generative-exact")
    # The mask's table.
    t = doc["mask_table"]
    ok = table_outcome()["sha256"] == t["sha256"] and all(sigmoid_entry(v) == e for v, e in t["at"])
    print("%s the mask's table of §7.6 (%s)" % ("PASS" if ok else "FAIL", t["sha256"][:16]))
    if not ok:
        fails.append("the mask's table differs from the corpus's")
    # The forms.
    codes = set()
    for v in doc["forms"]:
        got = form_outcome(v)
        ok = got == stated(v, ("outcome", "reason", "read"))
        codes.add(got.get("reason"))
        print(("PASS" if ok else "FAIL") + " form %s: %s" % (v["name"], got.get("reason", "read")))
        if not ok:
            fails.append("form %s: %s" % (v["name"], json.dumps(got)[:300]))
    missing = [c for c in FORM_CODES if c not in codes]
    print("%s every code of §10's second table refuses a file" % ("PASS" if not missing else "FAIL"))
    if missing:
        fails.append("no file refused with %s" % ", ".join(missing))
    # The frames.
    for v in doc["frames"]:
        got = frame_outcome(v, files)
        ok = got == stated(v, ("warped_sha256", "maps_sha256", "frame_sha256", "warped", "frame"))
        print(("PASS" if ok else "FAIL") + " frame %s: %s" % (v["name"], got["frame_sha256"][:16]))
        if not ok:
            fails.append("frame %s" % v["name"])
    # The cases.
    seen = set()
    for v in doc["cases"]:
        try:
            got = case_outcome(v, classes, files)
        except Exception as e:  # noqa: BLE001
            got = {"outcome": "unreadable: %r" % (e,)}
        ok = got == stated(v, ("outcome", "reason", "at", "statements"))
        if v.get("reason"):
            seen.add(v["reason"])
        print(("PASS" if ok else "FAIL") + " case %s: %s" % (v["name"], got["outcome"] + (" %s at %s" % (got["reason"], got["at"]) if "reason" in got else "")))
        if not ok:
            fails.append("%s: %s, the corpus states %s" % (v["name"], json.dumps(got)[:400], json.dumps(stated(v, ("outcome", "reason", "at")))[:200]))
    missing = [r for r in REASONS if r not in seen]
    print("%s every reason of §10's first table refuses a case" % ("PASS" if not missing else "FAIL"))
    if missing:
        fails.append("no case refused with %s" % ", ".join(missing))
    for label, d, fs in mutations(doc, files):
        caught = check(d, classes, fs)
        print(("PASS" if caught else "FAIL") + " mutated copy fails: %s%s" % (label, " (%s)" % caught[0][:100] if caught else ""))
        if not caught:
            fails.append("a mutated copy verifies: %s" % label)
    for f in fails:
        print("FAIL " + f)
    print(("PASS" if not fails else "FAIL") + " avatar-face corpus: %d forms, %d frames, %d cases" % (len(doc["forms"]), len(doc["frames"]), len(doc["cases"])))
    return 1 if fails else 0


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("usage: avatar_face_verify.py <live-conformance/avatar-face>", file=sys.stderr)
        sys.exit(2)
    sys.exit(main(sys.argv[1]))