#!/usr/bin/env python3
"""Independent verifier of the loudness corpus (extensions/packaging.md §12).

Written from the extension's text. It shares no code with wai-rs, which generates the corpus, and
never runs it. It takes Ed25519 verification as SPEC §7 *Signatures* states it from the routine
every Python verifier of WAI's corpora shares (`../claim-conformance/_ed25519.py`). Everything
§12 defines is written here: a programme's samples from its formula (§12.8), the integer
K-weighting, blocks and gates (§12.2), the integer logarithm (§12.3), the gain (§12.4), and the
statement's form, signature and checks (§12.5, §12.6).

For every case it requires the outcome the corpus states, field for field; that every reason of
§12.7's table appears; then it mutates copies of the corpus, one defect each, and requires every
copy to fail; and it drops each of its own checks in turn and requires the corpus to catch the
weakened verifier.

    python3 loudness_verify.py <pack-conformance/loudness>

Exit status 0 only if every case verifies, every mutated copy fails and every weakened verifier is
caught. Python 3.8 or later and the standard library.
"""

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

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "claim-conformance"))
import _ed25519 as ED  # noqa: E402

EXTENSION = os.path.join(HERE, "..", "extensions", "packaging.md")
DOMAIN = b"wai:pack-loudness\x01"
RATE, BLOCK, STEP = 48000, 19200, 4800
STAGE1 = (1648327767, -2890186776, 1286764483, -1815331593, 786495243)
STAGE2 = (1073741824, -2147483648, 1073741824, -2136797184, 1063081984)
LOG10_2 = 1292913986
WEIGHTS = {1: (100,), 2: (100, 100), 6: (100, 100, 100, 0, 141, 141)}
MEMBERS = {"blocks", "ceiling_milli", "component", "gain_milli", "gated", "integrated_milli", "kind", "limited", "method",
           "peak", "peak_milli", "receipt_hash", "sig", "signer_pubkey", "source", "target_milli", "version"}
SINT = re.compile(r"0|-?[1-9][0-9]*")
HEX = re.compile(r"[0-9a-f]*")
COMPONENT = re.compile(r"[A-Za-z0-9._-]{1,255}")

# A weakened run drops one of these checks (see `WEAKENINGS`); None checks everything.
DROPPED = None


def drops(name):
    return DROPPED == name


def reasons_table():
    """§12.7's codes, in order."""
    text = open(EXTENSION, encoding="utf-8").read()
    at = text.index("\n**12.7**")
    ends = [e for e in (text.find("\n**12.", at + 1), text.find("\n## ", at + 1)) if e >= 0]
    rows = text[at:min(ends) if ends else len(text)].splitlines()
    return [r.split("`")[1] for r in rows if r.startswith("| `")]


def filter_of_text():
    """§12.2's two sections' coefficients, as the text states them."""
    text = open(EXTENSION, encoding="utf-8").read()

    def read(after):
        at = text.index(after) + len(after)
        body = text[at:text.index(")", at)].replace("\u2212", "-")
        return tuple(int("".join(v.split())) for v in body.split(","))

    return read("the first section (b0, b1, b2, a1, a2) = ("), read("the second (")


def tone_table():
    """§12.8's table of the tone's period."""
    text = open(EXTENSION, encoding="utf-8").read()
    m = re.search(r"sine_1k = \[([^\]]*)\]", text)
    return [int(v) for v in m.group(1).replace("\n", " ").split(",")]


class Refused(Exception):
    def __init__(self, reason):
        super().__init__(reason)
        self.reason = reason


# ------------------------------------------------------------------ §12.8: a programme's samples

def lcg(x):
    return (x * 1103515245 + 12345) % (1 << 31)


def source(s, frames, table):
    k = s["kind"]
    if k == "zero":
        return [0] * frames
    if k == "tone":
        return [table[i % 48] * s["num"] // s["den"] for i in range(frames)]
    out, x = [], s["seed"]
    for i in range(frames):
        x = lcg(x)
        v = (x >> 8) % (2 * s["amp"] + 1) - s["amp"]
        out.append(32767 if k == "spike" and i == s["at"] else v)
    return out


def samples(p, table):
    """The programme's channels, each a list."""
    chans = [[] for _ in range(p["channels"])]
    for seg in p["segments"]:
        for c, s in zip(chans, seg["sources"]):
            c.extend(source(s, seg["frames"], table))
    return chans


# ------------------------------------------------------------------ §12.2, §12.3

def log2_q32(x):
    n = x.bit_length() - 1
    m = x << (62 - n) if n <= 62 else x >> (n - 62)
    r = n << 32
    for k in range(31, 0 if drops("log_bit") else -1, -1):
        m = (m * m) >> 62
        if m >= 1 << 63:
            m >>= 1
            r += 1 << k
    return r


def milli_db(p, d):
    return (10000 * (log2_q32(p) - log2_q32(d)) * LOG10_2) // (1 << 64)


def loudness(p, blocks):
    return -691 + milli_db(p, 100 * BLOCK * (1 << 46) * blocks)


def stage(xs, c):
    b0, b1, b2, a1, a2 = c
    x1 = x2 = y1 = y2 = 0
    out = []
    for x in xs:
        y = (b0 * x + b1 * x1 + b2 * x2 - a1 * y1 - a2 * y2 + (1 << 29)) >> 30
        x2, x1, y2, y1 = x1, x, y1, y
        out.append(y)
    return out


_MEASURED = {}


def measure(chans, rate):
    """§12.2: (frames, blocks, gated, integrated or None, peak, peak figure or None)."""
    key = (DROPPED, rate, tuple(tuple(c) for c in chans))
    if key in _MEASURED:
        return _MEASURED[key]
    n = len(chans)
    if rate != RATE or n not in WEIGHTS:
        raise Refused("loudness-unsupported")
    w = WEIGHTS[n]
    if drops("weights"):
        w = tuple(0 if g == 0 else 100 for g in w)
    if drops("lfe"):
        w = tuple(100 if g == 0 else g for g in w)
    frames = len(chans[0])
    peak = max((abs(v) for c in chans for v in c), default=0)
    peak_milli = 2 * milli_db(peak, 32768) if peak else None
    step = BLOCK if drops("overlap") else STEP
    blocks = (frames - BLOCK) // step + 1 if frames >= BLOCK else 0
    z = []
    for c in chans:
        x = [v << 8 for v in c]
        z.append(stage(x if drops("kfilter") else stage(x, STAGE1), STAGE2))
    sq = [None if g == 0 else [v * v for v in zc] for zc, g in zip(z, w)]
    prefix = [None if s is None else [0] for s in sq]
    for p, s in zip(prefix, sq):
        if s is not None:
            acc = 0
            for v in s:
                acc += v
                p.append(acc)
    power = []
    for j in range(blocks):
        a = j * step
        power.append(sum(g * (p[a + BLOCK] - p[a]) for p, g in zip(prefix, w) if g))
    level = [loudness(p, 1) if p > 0 else None for p in power]
    above = [j for j in range(blocks) if level[j] is not None and (drops("absolute_gate") or level[j] > -70000)]
    if not above:
        gated, integrated = 0, None
    else:
        mean = loudness(sum(power[j] for j in above), len(above))
        kept = above if drops("relative_gate") else [j for j in above if level[j] > mean - 10000]
        gated, integrated = len(kept), loudness(sum(power[j] for j in kept), len(kept))
    out = (frames, blocks, gated, integrated, peak, peak_milli)
    _MEASURED[key] = out
    return out


def gain(integrated, peak_milli, target, ceiling):
    """§12.4: (gain or None, limited)."""
    if integrated is None:
        return None, False
    want = target - integrated
    if peak_milli is not None and want > ceiling - peak_milli and not drops("ceiling"):
        return ceiling - peak_milli, True
    return want, False


def sint_out(v):
    return None if v is None else str(v)


# ------------------------------------------------------------------ §12.5, §12.6: statements

def jcs(v):
    return json.dumps(v, sort_keys=True, separators=(",", ":"), ensure_ascii=False)


def no_dupes(pairs):
    out = {}
    for k, v in pairs:
        if k in out:
            raise ValueError("a member twice")
        out[k] = v
    return out


def read_sint(v):
    if not isinstance(v, str) or not SINT.fullmatch(v) or abs(int(v)) >= 1 << 53:
        raise Refused("loudness-malformed")
    return int(v)


def read_sint_opt(v):
    return None if v is None else read_sint(v)


def uint(v):
    if not isinstance(v, int) or isinstance(v, bool) or not 0 <= v < 1 << 53:
        raise Refused("loudness-malformed")
    return v


def hexn(v, n):
    if not isinstance(v, str) or len(v) != 2 * n or not HEX.fullmatch(v):
        raise Refused("loudness-malformed")
    return bytes.fromhex(v)


def verify(text, receipt, src):
    """§12.6, in its order: the statement's summary, or Refused."""
    bad = Refused("loudness-malformed")
    try:
        v = json.loads(text, object_pairs_hook=no_dupes)
    except ValueError:
        raise bad
    if not isinstance(v, dict) or set(v) != MEMBERS or not isinstance(v["source"], dict) \
            or set(v["source"]) != {"channels", "frames", "rate", "sha256"}:
        raise bad
    if v["kind"] != "wai.pack.loudness" or type(v["version"]) is not int or v["version"] != 1 \
            or not isinstance(v["method"], str):
        raise bad
    if not isinstance(v["component"], str) or not COMPONENT.fullmatch(v["component"]) or v["limited"] not in ("yes", "no"):
        raise bad
    s = v["source"]
    frames, rate, channels, sha = uint(s["frames"]), uint(s["rate"]), uint(s["channels"]), hexn(s["sha256"], 32)
    blocks, gated, peak = uint(v["blocks"]), uint(v["gated"]), uint(v["peak"])
    if peak > 32768:
        raise bad
    integrated, peak_milli = read_sint_opt(v["integrated_milli"]), read_sint_opt(v["peak_milli"])
    target, ceiling, g = read_sint(v["target_milli"]), read_sint(v["ceiling_milli"]), read_sint_opt(v["gain_milli"])
    rh = None if v["receipt_hash"] is None else hexn(v["receipt_hash"], 32)
    pk, sig = hexn(v["signer_pubkey"], 32), hexn(v["sig"], 64)
    body = {k: x for k, x in v.items() if k != "sig"}
    if not drops("signature") and ED.verify(pk, DOMAIN + jcs(body).encode("utf-8"), sig) is not None:
        raise Refused("loudness-signature")
    if v["method"] != "bs1770-int/1" or rate != RATE or channels not in WEIGHTS:
        raise Refused("loudness-unsupported")
    want_blocks = (frames - BLOCK) // STEP + 1 if frames >= BLOCK else 0
    want_peak = 2 * milli_db(peak, 32768) if peak else None
    if not drops("own_figures") and (blocks != want_blocks or gated > blocks or (gated == 0) != (integrated is None)
                                     or peak_milli != want_peak
                                     or gain(integrated, peak_milli, target, ceiling) != (g, v["limited"] == "yes")):
        raise Refused("loudness-gain")
    if receipt is not None and rh != receipt and not drops("receipt"):
        raise Refused("loudness-receipt-mismatch")
    if src is not None:
        chans, srate = src
        inter = bytearray()
        for k in range(len(chans[0])):
            for c in chans:
                inter += (c[k] & 0xFFFF).to_bytes(2, "little")
        if not drops("digest") and (hashlib.sha256(bytes(inter)).digest() != sha or srate != rate or len(chans) != channels):
            raise Refused("loudness-source-mismatch")
        got = measure(chans, srate)
        if got != (frames, blocks, gated, integrated, peak, peak_milli):
            raise Refused("loudness-not-recomputed")
    return {"component": v["component"], "integrated_milli": sint_out(integrated), "gain_milli": sint_out(g),
            "limited": v["limited"]}


# ------------------------------------------------------------------ the corpus

def outcome_of(m):
    frames, blocks, gated, integrated, peak, peak_milli = m
    return {"frames": frames, "blocks": blocks, "gated": gated, "integrated_milli": sint_out(integrated), "peak": peak,
            "peak_milli": sint_out(peak_milli)}


def check(doc, quiet=False):
    fails = []
    table = doc["sine_1k"]
    if table != tone_table():
        fails.append("sine_1k is not the table of §12.8")
    if filter_of_text() != (STAGE1, STAGE2):
        fails.append("the filter is not §12.2's")
    for x, want in doc["logarithms"]:
        if str(log2_q32(int(x))) != want:
            fails.append(f"log2 of {x}: stated {want}, got {log2_q32(int(x))}")
    srcs = {}
    for p in doc["programmes"]:
        try:
            chans = samples(p, table)
            got = outcome_of(measure(chans, p["rate"]))
            srcs[p["name"]] = (chans, p["rate"])
        except (Refused, KeyError, TypeError, ValueError, IndexError, ZeroDivisionError) as e:
            got = {"error": repr(e)}
        if got != p["outcome"]:
            fails.append(p["name"])
            if not quiet:
                print(f"FAIL {p['name']}\n  stated {json.dumps(p['outcome'])}\n  got    {json.dumps(got)}")
    given = set()
    for c in doc["statements"]:
        try:
            receipt = None if c["receipt_hash"] is None else bytes.fromhex(c["receipt_hash"])
            src = None if c["source"] is None else srcs[c["source"]]
            got = {"verified": verify(c["statement"], receipt, src)}
        except Refused as r:
            got = {"refuse": r.reason}
        except (KeyError, TypeError, ValueError, IndexError) as e:
            got = {"error": repr(e)}
        if got != c["outcome"]:
            fails.append(c["name"])
            if not quiet:
                print(f"FAIL {c['name']}\n  stated {json.dumps(c['outcome'])}\n  got    {json.dumps(got)}")
        if "refuse" in c["outcome"]:
            given.add(c["outcome"]["refuse"])
    for code in reasons_table():
        if code not in given:
            fails.append("no case gives " + code)
    return fails


def mutations(doc):
    def prog(d, name):
        return next(p for p in d["programmes"] if p["name"] == name)

    def stmt(d, name):
        return next(c for c in d["statements"] if c["name"] == name)

    def m_table(d):
        d["sine_1k"][12] -= 1

    def m_frames(d):
        prog(d, "stereo_minus_20")["segments"][0]["frames"] += 4800

    def m_scale(d):
        prog(d, "surround")["segments"][0]["sources"][4]["den"] = 9

    def m_seed(d):
        prog(d, "relative_gate")["segments"][0]["sources"][0]["seed"] += 1

    def m_stated(d):
        o = prog(d, "mono_noise")["outcome"]
        o["integrated_milli"] = str(int(o["integrated_milli"]) + 1)

    def m_target(d):
        c = stmt(d, "to_target")
        v = json.loads(c["statement"])
        v["target_milli"] = "-23000"
        c["statement"] = jcs(v)

    def m_source(d):
        stmt(d, "to_target")["source"] = "mono_noise"

    def m_receipt(d):
        stmt(d, "bound_to_a_receipt")["receipt_hash"] = "00" * 32

    def m_log(d):
        d["logarithms"][2][1] = str(int(d["logarithms"][2][1]) + 1)

    def m_reason_gone(d):
        d["statements"].remove(stmt(d, "loudness_off_by_one"))
        d["statements"].remove(stmt(d, "peak_off_by_one"))

    return [m_table, m_frames, m_scale, m_seed, m_stated, m_target, m_source, m_receipt, m_log, m_reason_gone]


WEAKENINGS = ["weights", "lfe", "overlap", "kfilter", "absolute_gate", "relative_gate", "ceiling", "signature",
              "own_figures", "receipt", "digest", "log_bit"]


def main(argv):
    global DROPPED
    if len(argv) != 2:
        print(__doc__)
        return 2
    doc = json.load(open(os.path.join(argv[1], "vectors.json"), encoding="utf-8"))
    extra = sorted(set(os.listdir(argv[1])) - {"vectors.json", "README.md"})
    fails = check(doc)
    if extra:
        fails.append(f"files the corpus does not hold: {extra}")
    print(f"{len(doc['programmes'])} programmes, {len(doc['statements'])} statements: {len(fails)} failure(s)")
    caught = 0
    for m in mutations(doc):
        d = copy.deepcopy(doc)
        m(d)
        if check(d, quiet=True):
            caught += 1
        else:
            fails.append(f"mutation {m.__name__} was not caught")
    print(f"{caught} of {len(mutations(doc))} mutated copies fail")
    weak = 0
    for w in WEAKENINGS:
        DROPPED = w
        if check(doc, quiet=True):
            weak += 1
        else:
            fails.append(f"a verifier without its {w} check passes the corpus")
    DROPPED = None
    print(f"{weak} of {len(WEAKENINGS)} weakened verifiers are caught")
    for f in fails:
        print("FAIL", f)
    return 1 if fails else 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))