#!/usr/bin/env python3
"""Extract the VC-3 constant tables (Annex D quantization weights, Annex E VLC tables) from the
text of the public SMPTE ST 2019-1:2016 standard.

Usage:
    curl -LO https://pub.smpte.org/doc/st2019-1/20160518-pub/st2019-1-20160518-pub.zip
    unzip st2019-1-20160518-pub.zip
    pdftotext -layout st2019-1-2016.pdf vc3.txt
    python3 extract_tables.py vc3.txt > ../src/spec_tables.rs

Only the standard's text is read. Every table's size is checked (63 AC weights per matrix, 4*64+1
amplitude codewords including EOB, 62 run codewords) and every VLC table must be a complete
prefix code.
"""
import re
import sys

text = open(sys.argv[1], encoding="utf-8").read()
lines = text.split("\n")


def section(start_pat, end_pat):
    s = next(i for i, l in enumerate(lines) if re.search(start_pat, l))
    e = next(i for i in range(s + 1, len(lines)) if re.search(end_pat, lines[i]))
    return lines[s + 1:e]


def weights(sec):
    """The 63 AC weights of an 8x8 table printed in raster order with '-' for the DC entry."""
    nums = []
    for l in sec:
        if "Page" in l or "SMPTE" in l or "Compression" in l:
            continue
        nums += [int(x) for x in re.findall(r"\b\d+\b", l)]
    return nums


def check_prefix(codes):
    """Kraft sum of a prefix code (1.0 when complete)."""
    for a in codes:
        for b in codes:
            if a is not b and b.startswith(a):
                raise SystemExit("not a prefix code: %s %s" % (a, b))
    return sum(2.0 ** -len(c) for c in codes)


out = ["// Generated by tools/extract_tables.py from SMPTE ST 2019-1:2016 (VC-3), Annexes D and E.",
       "// Do not edit by hand.", ""]

# ---- Annex D: weights in raster order (v*8+u), DC entry omitted ----
dtables = [("D_1", r"Table D\.1 – Compression ID 1235"), ("D_2", r"Table D\.2 – Compression ID 1237"),
           ("D_3", r"^Table D\.3 – Compression"), ("D_4", r"Table D\.4 – Compression"),
           ("D_5", r"Table D\.5 – Compression"), ("D_6", r"Table D\.6 – Compression"),
           ("D_7", r"^Table D\.7 – Compression"), ("D_8", r"Table D\.8 – Compression"),
           ("D_9", r"Table D\.9 – Compression"), ("D_10", r"Table D\.10 – Compression"),
           ("D_11", r"^Table D\.11 – Compression")]
for i, (name, pat) in enumerate(dtables):
    end = dtables[i + 1][1] if i + 1 < len(dtables) else r"^Annex E VLC Tables"
    sec = section(pat, end)
    if name == "D_11":
        sec = [l for l in sec if "Compression ID" not in l]
        nums = weights(sec)
        assert len(nums) == 63, (name, len(nums))
        out.append("pub const W_%s: [[u8; 63]; 2] = [%s, %s];" % (name, nums, nums))
        continue
    li = next(j for j, l in enumerate(sec) if "Luma Quantization" in l)
    ci = next(j for j, l in enumerate(sec) if "Chroma Quantization" in l)
    luma = weights(sec[li + 1:ci])
    chroma = weights(sec[ci + 1:])
    assert len(luma) == 63 and len(chroma) == 63, (name, len(luma), len(chroma))
    out.append("pub const W_%s: [[u8; 63]; 2] = [%s, %s];" % (name, luma, chroma))
out.append("")

# ---- Annex E ----
etables = ["E.1", "E.4", "E.7", "E.10", "E.13", "E.16"]
for n, t in enumerate(etables):
    a = int(t[2:])
    amp = section(r"^\s*Table E\.%d – Amplitude" % a, r"Table E\.%d – Run" % (a + 1))
    run = section(r"Table E\.%d – Run" % (a + 1), r"Table E\.%d – DC" % (a + 2))
    nxt = r"Table E\.%d – Amplitude" % (a + 3) if n + 1 < len(etables) else r"^Annex F System"
    dc = section(r"Table E\.%d – DC" % (a + 2), nxt)
    ent = []
    for l in amp:
        for m in re.finditer(r"\b([01]+)\s+(\d+)\s+(EOB1?|\d+)\s+([01])\s+([01])\b", l):
            code, ln, am, fr, fi = m.groups()
            assert len(code) == int(ln), m.group(0)
            ent.append((code, 0 if am.startswith("EOB") else int(am), int(fr), int(fi), am.startswith("EOB")))
    k = check_prefix([e[0] for e in ent])
    assert abs(k - 1.0) < 1e-12, (t, k, len(ent))
    assert sum(e[4] for e in ent) == 1 and len(ent) == 4 * 64 + 1, (t, len(ent))
    rent = []
    for l in run:
        for m in re.finditer(r"\b([01]+)\s+(\d+)\s+(\d+)\b", l):
            code, ln, r = m.groups()
            if len(code) != int(ln):
                continue
            rent.append((code, int(r)))
    assert len(rent) == 62 and sorted(r for _, r in rent) == list(range(1, 63)), (t, len(rent))
    assert abs(check_prefix([c for c, _ in rent]) - 1.0) < 1e-12
    dent = []
    for l in dc:
        m = re.match(r"^\s+([01]+)\s+(\d+)\s*$", l)
        if m:
            dent.append((m.group(1), int(m.group(2))))
    assert [e for _, e in dent] == list(range(len(dent))), (t, dent)
    assert abs(check_prefix([c for c, _ in dent]) - 1.0) < 1e-12, (t, dent)
    name = t.replace(".", "_")
    out.append("/// Table %s: (codeword, length, amplitude (0 = EOB), Frun, Findex)." % t)
    out.append("pub const AMP_%s: [(u16, u8, u8, u8, u8); %d] = [" % (name, len(ent)))
    for c, a_, fr, fi, _ in ent:
        out.append("    (0b%s, %d, %d, %d, %d)," % (c, len(c), a_, fr, fi))
    out.append("];")
    out.append("/// Table E.%d: (codeword, length, run)." % (a + 1))
    out.append("pub const RUN_%s: [(u16, u8, u8); 62] = [" % name)
    for c, r in rent:
        out.append("    (0b%s, %d, %d)," % (c, len(c), r))
    out.append("];")
    out.append("/// Table E.%d: (codeword, length) indexed by the DC size.", )
    out[-1] = "/// Table E.%d: (codeword, length) indexed by the DC size (eta)." % (a + 2)
    out.append("pub const DC_%s: [(u16, u8); %d] = [%s];" % (name, len(dent), ", ".join("(0b%s, %d)" % (c, len(c)) for c, _ in dent)))
    out.append("")

print("\n".join(out))