#!/usr/bin/env python3
"""Second implementation of the audio scheme of the second marking layer, `wai-mark:band64`
(extensions/synthesis.md §7.10, §13), over the corpus `live-conformance/audio-mark/`.

Written from the extensions' text. It shares no code with wai-rs: the scheme (the cosine tables,
the frames, the integer logarithm, the stream, the shuffle, the signs, the mask, the embedding
with its three syntheses, the detector with its offsets and phases) and the pinned transforms of
sound are _wai_audio_mark.py's, written from synthesis §7.10 and robustness-evidence §4.1, the
code _wai_mark.py's, written from §7.5.3; every case is verified by synthesis_verify.py with §7.6's
checks; every robustness record by robustness_verify.py, which reruns the conformance record's
harness end to end. It reads from the extension the scheme's definition and digest (§7.10.1),
the shuffle's first entries (§7.10.3), the false-positive bound (§7.10.6), and the tables and the
claims of §7.10.7.

    python3 -I audio_mark_verify.py <live-conformance/audio-mark>

It checks that every audio file is named by its SHA-256; that the definition is §7.10.1's and
has its digest, and the cosine tables the digest the definition pins; that the layout is
§7.10.3's; that every embedding, detection and battery sample of the corpus is what this
implementation gives; that every case gets the outcome, the reason, the place and the statements
it states; that the robustness records verify, the conformance record reproduced; and that
§7.10.7's tables are the measurement's and its claims attacks the records show survived. Then it
mutates copies of the corpus and requires each to fail, and weakens its own verifier in turn and
requires the corpus to catch each. Exit status 0 only if all of that holds. Python 3.8 or later;
standard library only.
"""

import copy
import hashlib
import json
import math
import os
import re
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import _wai_audio_mark as A  # noqa: E402
import _wai_mark as TENT  # noqa: E402
import synthesis_verify as SV  # noqa: E402
import robustness_verify as RV  # noqa: E402


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


def sec(head):
    return SV.section(SV._TEXT, head)


def fold(s):
    return " ".join(s.split())


# ---------------------------------------------------------------- what the extension states

def text_definition():
    block = sec("#### 7.10.1 ").split("```")[1]
    return "".join(line.strip() for line in block.strip().splitlines()).encode("utf-8")


def text_shuffle_head():
    m = re.search(r"π begins ([0-9, ]+), and the corpus", fold(sec("#### 7.10.3 ")))
    return [int(x) for x in m.group(1).split(",")] if m else None


def text_fp():
    m = re.search(r"2\^−([0-9.]+), about", fold(sec("#### 7.10.6 ")))
    return "-" + m.group(1) if m else None


def text_tables():
    """§7.10.7: each table's rows, as {attack: {column: cell}}, by the table's first header."""
    out = {}
    s = sec("#### 7.10.7 ")
    head = None
    for line in s.splitlines():
        if line.startswith("| attack"):
            head = [c.strip().strip("`") for c in line.strip("|").split("|")]
            out[head[1]] = {}
        elif head and line.startswith("| `"):
            cells = [c.strip() for c in line.strip("|").split("|")]
            out[head[1]][cells[0].strip("`")] = dict(zip(head[1:], cells[1:]))
        elif not line.startswith("|"):
            head = None
    return out


def text_claims():
    """§7.10.7: the attacks it claims to survive."""
    m = re.search(r"\*\*claims to survive\*\* ([^.]*)\.", fold(sec("#### 7.10.7 ")))
    return re.findall(r"`([a-z0-9/.+-]+)`", m.group(1)) if m else None


def fp_log2(attempts):
    n, term, v = 248, 1.0, 0.0
    for i in range(30):
        v += term
        term = term * (n - i) / (i + 1)
    return "%.2f" % (math.log2(attempts) + 58 + math.log2(v) - n)


# ---------------------------------------------------------------- the corpus's audio

_MARKED = {}


def marked(rate, channels, buf, payload):
    k = (sha(buf), rate, channels, payload)
    if k not in _MARKED:
        _MARKED[k] = A.embed_bytes(rate, channels, buf, payload)
    return _MARKED[k]


_READ = {}


def detect(rate, channels, buf):
    k = (sha(buf), rate, channels)
    if k not in _READ:
        _READ[k] = A.detect(rate, channels, buf)
    return _READ[k]


def detected_json(d):
    if d is None:
        return None
    return {"payload": d[0].hex(), "offset": d[1], "phase": d[2], "errors": d[3]}


def derive(src, stored):
    """A source of §13: stored audio, marked under a payload, then a pinned transform, then one
    sample's lowest bit flipped. (rate, channels, buffer)."""
    buf = stored[src["audio"]]
    rate, ch = src["rate"], src["channels"]
    if src["marked"] is not None:
        buf = marked(rate, ch, buf, bytes.fromhex(src["marked"]))[0]
    if src["derive"] is not None:
        d = src["derive"]
        name = d["transform"].split(":", 1)[1]
        rate, chans = A.attack(name, d["params"], rate, A.unpack(buf, ch), d["seed"])
        ch = len(chans)
        buf = A.pack(chans)
    if src["flip"] is not None:
        b = bytearray(buf)
        b[2 * src["flip"]] ^= 1
        buf = bytes(b)
    return rate, ch, buf


def bits(b):
    return "".join("1" if x else "0" for x in b)


def check(doc, battery, scheme_bytes, stored, files, classes, only=None):
    fails = []
    want = lambda part: only is None or part in only  # noqa: E731
    for h, b in stored.items():
        if sha(b).hex() != h:
            fails.append("audio %s is not named by its SHA-256" % h)
    sc = doc["scheme"]
    if scheme_bytes != text_definition() or sha(scheme_bytes) != SV.SCHEME_DIGEST_A or sc["digest"] != SV.SCHEME_DIGEST_A.hex():
        fails.append("the definition is not §7.10.1's, or its digest is not the one §7.10.1 states")
    try:
        definition = json.loads(scheme_bytes)
    except ValueError:
        definition = {}
    tables = sha(A.tables_bytes()).hex()
    if definition.get("cosine") != [A.ONE, tables] or sc["tables"] != tables:
        fails.append("the cosine tables do not have the digest the definition pins")
    if definition.get("step") != A.STEP or definition.get("gate") != A.GATE or definition.get("bands") != list(A.BANDS):
        fails.append("the definition's parameters are not this implementation's")
    head = text_shuffle_head()
    if head is None or A.PERM[:len(head)] != head or sc["shuffle_first"] != A.PERM[:16] or sc["shuffle_last"] != A.PERM[-16:]:
        fails.append("the shuffle is not §7.10.3's")
    if sc["signs_first"] != A.SIGN[:16] or sc["mask"] != bits(A.MASK):
        fails.append("the signs or the mask are not §7.10.3's")
    if TENT.decode(A.MASK) is not None:
        fails.append("the mask decodes to a payload")
    fp = fp_log2(A.PERIOD * A.PHASES)
    if sc["false_positive_log2"] != fp or text_fp() != fp:
        fails.append("the false-positive bound is not %s (§7.10.6)" % fp)
    for e, l in doc["lg"]:
        if A.lg(int(e)) != l:
            fails.append("lg(%s) is not %d" % (e, l))
    for a, s in doc["shift"]:
        if A.shift(a) != s:
            fails.append("S(%d) is not %d" % (a, s))

    for e in (doc["embeds"] if want("embeds") else []):
        try:
            buf = stored[e["audio"]]
            out, gains = marked(e["rate"], e["channels"], buf, bytes.fromhex(e["payload"]))
            g = b"".join((x & 0xFFFF).to_bytes(2, "big") for x in gains).hex()
            if sha(out).hex() != e["marked"] or g != e["gains"] or len(gains) != e["cells"]:
                fails.append("embed %s: the marked samples or the gains differ" % e["name"])
        except Exception as x:  # noqa: BLE001
            fails.append("embed %s: %r" % (e["name"], x))
    for d in (doc["detects"] if want("detects") else []):
        try:
            rate, ch, buf = derive(d["source"], stored)
            got = detected_json(detect(rate, ch, buf))
            if (sha(buf).hex(), rate, ch, len(buf) // (2 * ch), got) != (d["samples"], d["rate"], d["channels"], d["length"], d["detected"]):
                fails.append("detect %s: %s" % (d["name"], got))
        except Exception as x:  # noqa: BLE001
            fails.append("detect %s: %r" % (d["name"], x))
    voice = next(e for e in doc["embeds"] if e["name"] == "voice")
    vbuf = stored[voice["audio"]]
    for x in (doc["battery_sample"] if want("battery") else []):
        try:
            buf = marked(voice["rate"], voice["channels"], vbuf, bytes.fromhex(voice["payload"]))[0] if x["which"] == "marked" else vbuf
            name, params = x["transform"].split(":", 1)[1], x["params"]
            rate, chans = A.attack(name, params, voice["rate"], A.unpack(buf, voice["channels"]), x["seed"])
            y = A.pack(chans)
            got = detected_json(detect(rate, len(chans), y))
            if (sha(y).hex(), rate, len(chans), got) != (x["samples"], x["rate"], x["channels"], x["detected"]):
                fails.append("battery sample %s of the %s voice: %s" % (x["attack"], x["which"], got))
        except Exception as e:  # noqa: BLE001
            fails.append("battery sample %s: %r" % (x.get("attack"), e))

    for v in (doc["cases"] if want("cases") else []):
        case = dict(v, room=None)
        try:
            case["decoded"] = [{"object": d["object"], "samples": derive(d["source"], stored)[2].hex()} for d in v["decoded"]]
            case["unmarked"] = [{"object": d["object"], "samples": derive(d["source"], stored)[2].hex()} for d in v["unmarked"]]
            case["priors"] = []
            got = SV.outcome(case, classes)
        except Exception as e:  # noqa: BLE001
            got = {"outcome": "unreadable: %r" % (e,)}
        if got != SV.stated(v):
            fails.append("case %s: %s" % (v["name"], json.dumps(got)[:300]))

    if want("measurement"):
        fails += check_measurement(battery, files)
    if want("records"):
        fails += check_records(files, stored)
    return fails


def check_measurement(battery, files):
    """§7.10.7: its tables are the measurement's, and its claims attacks every record shows
    survived."""
    fails = []
    t = text_tables()
    rows = {}
    for r in battery["table"]:
        rows.setdefault(r["attack"], {})[r["set"]] = r
    measured = {}
    for k in sorted(files):
        if k.startswith("robustness/measured/measurement-"):
            m = json.loads(files[k])
            for r in m["table"]:
                measured.setdefault(r["attack"], {})[r["set"]] = r
    pinned = t.get("voice", {})
    if not pinned:
        fails.append("§7.10.7 has no table of the pinned attacks")
    for attack, cells in pinned.items():
        for col, cell in cells.items():
            src = rows if col in ("voice", "music") else measured
            r = src.get(attack, {}).get(col)
            want = "%d / %d" % (r["survived"], r["items"]) if r else None
            if r and r["survived"]:
                want += " (%d)" % r["errors_max"]
            if want is None or cell != want:
                fails.append("§7.10.7 %s, %s: %s, the measurement %s" % (attack, col, cell, want))
    codecs = t.get("speech", {})
    for attack, cells in codecs.items():
        for col, cell in cells.items():
            r = measured.get(attack, {}).get(col)
            want = "%d / %d" % (r["survived"], r["items"]) if r else None
            if r and r["survived"]:
                want += " (%d)" % r["errors_max"]
            if want is None or cell != want:
                fails.append("§7.10.7 %s, %s: %s, the measurement %s" % (attack, col, cell, want))
    claims = text_claims()
    if claims is None:
        fails.append("§7.10.7 states no claim")
    else:
        everything = dict(rows)
        for a, c in measured.items():
            everything.setdefault(a, {}).update(c)
        for a in claims:
            sets = everything.get(a, {})
            if not sets or any(r["survived"] != r["items"] or r["false_positives"] for r in sets.values()):
                fails.append("§7.10.7 claims %s, which a set did not survive" % a)
    if battery["false_positive_log2"] != fp_log2(A.PERIOD * A.PHASES):
        fails.append("the measurement's bound is not §7.10.6's")
    return fails


def check_records(files, stored):
    """robustness-evidence §10: the scheme's records, each case verified by robustness_verify."""
    fails = []
    samples = {h: b for h, b in stored.items()}
    get = lambda k: files[k].decode("utf-8") if k in files else None  # noqa: E731

    def run(label, records, want):
        corpora = {i: c for i, (_, c) in enumerate(records) if c is not None}
        case = {"records": [r for r, _ in records], "corpora": corpora, "samples": samples, "implements": [A.SCHEME],
                "require_reproduced": False}
        got = RV.verify_case(case)
        if got[0] != "accept":
            return ["%s: refused %s at %d" % (label, got[1], got[2])]
        st = [s["status"] for s in got[1]]
        return [] if st == want else ["%s: %s, not %s" % (label, st, want)]
    cr, cd = get("robustness/records/band64-conformance.json"), get("robustness/corpus/band64-conformance.json")
    if cr is None or cd is None:
        return ["the conformance record or its corpus is missing"]
    fails += run("the conformance record", [(cr, cd)], ["reproduced"])
    gen = [(get("robustness/records/band64-%s.json" % s), get("robustness/corpus/band64-%s.json" % s)) for s in ("voice", "music")]
    claim = get("robustness/records/band64-claim.json")
    if any(r is None or c is None for r, c in gen) or claim is None:
        fails.append("a generated set's record, its corpus or the claim is missing")
    else:
        fails += run("the generated sets' records and claim", gen + [(claim, None)], ["stated"] * 3)
    meas = []
    for k in sorted(files):
        f = k[len("robustness/measured/"):]
        if not k.startswith("robustness/measured/") or not f.endswith(".json") or f.startswith(("measurement-", "corpus-")) or f == "claim.json" or f.startswith("negatives"):
            continue
        name = f[:-len(".json")]
        c = get("robustness/corpus/band64-%s.json" % name) or get("robustness/measured/corpus-%s.json" % name)
        if c is None:
            fails.append("%s: no corpus document" % k)
        meas.append((get(k), c))
    if meas:
        mc = get("robustness/measured/claim.json")
        fails += run("the measured records and their claim", meas + [(mc, None)], ["stated"] * (len(meas) + 1))
    return fails


# ---------------------------------------------------------------- mutations and weakenings

def mutations(doc, battery, scheme_bytes, stored, files):
    out = []

    def mutate(label, f):
        d, b, s, st, fs = copy.deepcopy(doc), copy.deepcopy(battery), scheme_bytes, dict(stored), dict(files)
        r = f(d, b, st, fs)
        if isinstance(r, bytes):
            s = r
        out.append((label, d, b, s, st, fs))
    mutate("a byte of the definition", lambda d, b, st, fs: scheme_bytes.replace(b'"step":848', b'"step":849'))
    mutate("the music's marked digest", lambda d, b, st, fs: d["embeds"][1].update(marked=SV.flip_hex_digit(d["embeds"][1]["marked"], 5)))
    mutate("a detection that reads the unmarked voice as marked",
           lambda d, b, st, fs: d["detects"][1].update(detected={"payload": d["embeds"][0]["payload"], "offset": 0, "phase": 0, "errors": 0}))
    mutate("a trimmed detection's offset", lambda d, b, st, fs: [x.update(detected=dict(x["detected"], offset=x["detected"]["offset"] + 1))
                                                                for x in d["detects"] if x["detected"] and x["detected"]["offset"]][:1])
    mutate("a battery sample's result digest", lambda d, b, st, fs: d["battery_sample"][3].update(samples=SV.flip_hex_digit(d["battery_sample"][3]["samples"], 1)))
    mutate("a case's statement", lambda d, b, st, fs: d["cases"][0]["statements"]["markings"][0].update(status="stated"))
    mutate("a tally of the measurement", lambda d, b, st, fs: b["table"][0].update(survived=b["table"][0]["survived"] - 1))
    mutate("a claim of an attack a set lost", lambda d, b, st, fs: [r.update(survived=0) for r in b["table"] if r["attack"] == "mulaw"][:1])
    return out


def weakened():
    """The verifier with one of §7.6's checks or §7.3's inheritance dropped, and the scheme with
    one of its own steps dropped: each with the parts of the corpus it is run on."""
    return [
        ("the marking not recomputed", SV, "recomputes", lambda *a: True, None),
        ("the detection not run again", SV, "reproduces", lambda *a: True, None),
        ("no agreement by inheritance", SV, "inherits", lambda c, p, by_link: False, None),
        ("one phase only", A, "PHASES", 1, {"detects"}),
        ("no gate", A, "GATE", -(1 << 40), {"embeds"}),
    ]


def main(corpus_dir):
    fails = []
    with open(os.path.join(HERE, "..", "registry", "capabilities.json"), encoding="utf-8") as f:
        classes = {c["capability"]: c["replicate"] for c in json.load(f)["capabilities"]}
    files = RV.load(corpus_dir)
    doc = json.loads(files["vectors.json"])
    battery = json.loads(files["battery-band64.json"])
    scheme_bytes = files["scheme-band64.json"]
    stored = {k[len("audio/"):-len(".pcm")]: v for k, v in files.items() if k.startswith("audio/")}

    base = check(doc, battery, scheme_bytes, stored, files, classes)
    for b in base:
        print("FAIL " + b)
    print(("PASS" if not base else "FAIL") + " the scheme, %d embeddings, %d detections, %d battery samples, %d cases and the records"
          % (len(doc["embeds"]), len(doc["detects"]), len(doc["battery_sample"]), len(doc["cases"])))
    fails += base
    seen = {v.get("reason") for v in doc["cases"]}
    missing = [r for r in SV.MARK_ONLY if r not in seen]
    print(("PASS" if not missing else "FAIL") + " every code only a scheme's corpus gives (§13) refuses a case: " + ", ".join(SV.MARK_ONLY))
    fails += ["codes no case gives: %s" % missing] if missing else []
    for label, d, b, s, st, fs in mutations(doc, battery, scheme_bytes, stored, files):
        caught = check(d, b, s, st, fs, classes)
        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 label, mod, name, value, only in weakened():
        saved = getattr(mod, name)
        setattr(mod, name, value)
        # The memos of embeddings, detections and harness runs are of the scheme's functions:
        # only a weakening of the scheme itself makes them stale.
        if mod is A:
            _READ.clear()
            _MARKED.clear()
            SV._EMBEDDED.clear()
            SV._DETECTED.clear()
            RV.HARNESS_MEMO.clear()
        try:
            caught = check(doc, battery, scheme_bytes, stored, files, classes, only)
        finally:
            setattr(mod, name, saved)
            if mod is A:
                _READ.clear()
                _MARKED.clear()
                SV._EMBEDDED.clear()
                SV._DETECTED.clear()
                RV.HARNESS_MEMO.clear()
        print(("PASS" if caught else "FAIL") + " the corpus catches a verifier with %s%s" % (label, " (%s)" % caught[0][:100] if caught else ""))
        if not caught:
            fails.append("a weakened verifier passes: %s" % label)
    for f in fails:
        print("FAIL " + f)
    print(("PASS" if not fails else "FAIL") + " audio-mark corpus")
    return 1 if fails else 0


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