#!/usr/bin/env python3
"""Extract HEVC syntax / decoding tables from the text of the public ITU-T H.265 specification.

Usage:
    curl -L -o h265.pdf 'https://www.itu.int/rec/dologin_pub.asp?lang=e&id=T-REC-H.265-202108-S!!PDF-E&type=items'
    pdftotext -layout h265.pdf h265.txt
    python3 extract_tables.py h265.txt > ../src/spec_tables.rs

Only the specification text is read. Extracted: CABAC initValue tables (Tables 9-5..9-35) arranged per
initType following Table 9-4, default scaling lists (Table 7-6), intraPredAngle / invAngle (Tables 8-5,
8-6), QpC (Table 8-10), beta'/tC' (Table 8-12), luma / chroma interpolation filter taps
(equations 8-236..8-238 and 8-249..8-261) and the 32x32 DCT matrix (equations 8-319, 8-321).
"""
import re
import sys
from pathlib import Path


def _spec_path(arg):
    """Resolve a CLI path argument and require it to be an existing regular file."""
    path = Path(arg).resolve()
    if not path.is_file():
        sys.exit(f"not a readable file: {arg}")
    return path


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


def body_after(title_regex, stop_regex=r"^\s*(Table \d+-\d+ - |\d+\s+Rec\. ITU|Rec\. ITU)", skip_first=True):
    """Lines following the (last) occurrence of a table title, until the next title."""
    idx = [i for i, l in enumerate(lines) if re.search(title_regex, l)]
    assert idx, title_regex
    start = idx[-1] + 1  # last occurrence = the table itself (first is the table of contents)
    out = []
    for l in lines[start:]:
        if re.search(r"^\s*Table \d+-\d+ - ", l):
            break
        out.append(l)
    return out


def init_values(table):
    vals = []
    for l in body_after(r"Table %s - Values of initValue" % re.escape(table)):
        m = re.match(r"^\s*initValue\s+(.*)$", l)
        if m:
            vals += [int(v) for v in m.group(1).split()]
    return vals


# (name, table, ranges per initType 0,1,2 as lists of ctxIdx; None = element not used for that initType)
R = lambda a, b: list(range(a, b + 1))
CABAC = [
    ("SAO_MERGE", "9-5", [[0], [1], [2]]),
    ("SAO_TYPE", "9-6", [[0], [1], [2]]),
    ("SPLIT_CU", "9-7", [R(0, 2), R(3, 5), R(6, 8)]),
    ("CU_TRANSQUANT_BYPASS", "9-8", [[0], [1], [2]]),
    ("CU_SKIP", "9-9", [None, R(0, 2), R(3, 5)]),
    ("PRED_MODE", "9-10", [None, [0], [1]]),
    ("PART_MODE", "9-11", [[0], R(1, 4), R(5, 8)]),
    ("PREV_INTRA_LUMA", "9-12", [[0], [1], [2]]),
    ("INTRA_CHROMA", "9-13", [[0], [1], [2]]),
    ("RQT_ROOT_CBF", "9-14", [None, [0], [1]]),
    ("MERGE_FLAG", "9-15", [None, [0], [1]]),
    ("MERGE_IDX", "9-16", [None, [0], [1]]),
    ("INTER_PRED_IDC", "9-17", [None, R(0, 4), R(5, 9)]),
    ("REF_IDX", "9-18", [None, R(0, 1), R(2, 3)]),
    ("MVP_FLAG", "9-19", [None, [0], [1]]),
    ("SPLIT_TRANSFORM", "9-20", [R(0, 2), R(3, 5), R(6, 8)]),
    ("CBF_LUMA", "9-21", [R(0, 1), R(2, 3), R(4, 5)]),
    ("CBF_CHROMA", "9-22", [R(0, 3) + [12], R(4, 7) + [13], R(8, 11) + [14]]),
    ("ABS_MVD_GT0", "9-23", [None, [0], [2]]),
    ("ABS_MVD_GT1", "9-23", [None, [1], [3]]),
    ("CU_QP_DELTA_ABS", "9-24", [R(0, 1), R(2, 3), R(4, 5)]),
    ("TRANSFORM_SKIP", "9-25", [[0, 3], [1, 4], [2, 5]]),
    ("LAST_X_PREFIX", "9-26", [R(0, 17), R(18, 35), R(36, 53)]),
    ("LAST_Y_PREFIX", "9-27", [R(0, 17), R(18, 35), R(36, 53)]),
    ("CODED_SUB_BLOCK", "9-28", [R(0, 3), R(4, 7), R(8, 11)]),
    ("SIG_COEFF", "9-29", [R(0, 41) + R(126, 127), R(42, 83) + R(128, 129), R(84, 125) + R(130, 131)]),
    ("GT1", "9-30", [R(0, 23), R(24, 47), R(48, 71)]),
    ("GT2", "9-31", [R(0, 5), R(6, 11), R(12, 17)]),
    ("EXPLICIT_RDPCM", "9-32", [None, [0, 2], [1, 3]]),
    ("EXPLICIT_RDPCM_DIR", "9-33", [None, [0, 2], [1, 3]]),
    ("CU_CHROMA_QP_OFFSET_FLAG", "9-34", [[0], [1], [2]]),
    ("CU_CHROMA_QP_OFFSET_IDX", "9-35", [[0], [1], [2]]),
]

out = []
w = out.append
w("// Generated by crates/hevc/tools/extract_tables.py from the text of Rec. ITU-T H.265. Do not edit.")
w("")
w("// CABAC context layout: offset of each syntax element in the context array (Table 9-4).")
offset = 0
init = [[], [], []]
for name, table, ranges in CABAC:
    vals = init_values(table)
    n = max(len(r) for r in ranges if r is not None)
    for t in range(3):
        r = ranges[t]
        if r is None:
            init[t] += [154] * n  # element not present in this initType (value unused)
        else:
            # shorter ranges (part_mode in I slices) are padded with an unused value
            init[t] += [vals[i] for i in r] + [154] * (n - len(r))
    w("pub const %s: usize = %d;" % (name, offset))
    offset += n
w("pub const NUM_CTX: usize = %d;" % offset)
w("")
w("/// initValue per initType (9.3.2.2).")
w("pub static CABAC_INIT: [[u8; NUM_CTX]; 3] = [")
for t in range(3):
    w("    [%s]," % ", ".join(str(v) for v in init[t]))
w("];")
w("")

# Table 7-6: default 8x8 scaling lists (intra matrixId 0..2, inter 3..5)
intra, inter = [], []
for l in body_after(r"Table 7-6 - Specification of default values of ScalingList"):
    m = re.match(r"^\s*ScalingList\[ 1\.\.3 \]\[ (0\.\.2|3\.\.5) \]\[ i \]\s+(.*)$", l)
    if m:
        (intra if m.group(1) == "0..2" else inter).extend(int(v) for v in m.group(2).split())
assert len(intra) == 64 and len(inter) == 64
w("/// Table 7-6: default ScalingList[1..3][0..2] (intra) and [3..5] (inter), in up-right diagonal order.")
w("pub static DEFAULT_SCALING_INTRA: [u8; 64] = [%s];" % ", ".join(map(str, intra)))
w("pub static DEFAULT_SCALING_INTER: [u8; 64] = [%s];" % ", ".join(map(str, inter)))
w("")

# Table 8-5 / 8-6
angle = {}
for l in body_after(r"Table 8-5 - Specification of intraPredAngle"):
    if l.strip().startswith("predModeIntra"):
        modes = [int(v) for v in l.split()[1:]]
    elif l.strip().startswith("intraPredAngle"):
        vs = l.split()[1:]
        for mo, v in zip(modes, vs):
            if v != "-":
                angle[mo] = int(v)
assert sorted(angle) == list(range(2, 35))
w("/// Table 8-5: intraPredAngle for predModeIntra 0..34 (0 for planar / DC).")
w("pub static INTRA_PRED_ANGLE: [i8; 35] = [%s];" % ", ".join(str(angle.get(i, 0)) for i in range(35)))
inv = {}
for l in body_after(r"Table 8-6 - Specification of invAngle"):
    s = l.strip()
    if s.startswith("predModeIntra"):
        modes = [int(v) for v in s.split()[1:]]
    elif s.startswith("invAngle"):
        vs = [int(v.replace(" ", "")) for v in re.findall(r"-\d+(?: \d{3})?", s)]
        for mo, v in zip(modes, vs):
            inv[mo] = v
assert sorted(inv) == list(range(11, 26)), inv
w("/// Table 8-6: invAngle for predModeIntra 11..25 (index = mode - 11).")
w("pub static INV_ANGLE: [i16; 15] = [%s];" % ", ".join(str(inv[i]) for i in range(11, 26)))
w("")

# Table 8-10: QpC for qPi 30..43
for l in body_after(r"Table 8-10 - Specification of QpC as a function of qPi"):
    s = l.strip()
    if s.startswith("qPi"):
        qpi = s.split()
    elif s.startswith("QpC"):
        qpc = s.split()
qpi_n = [int(v) for v in qpi[3:-2]]
qpc_n = [int(v) for v in qpc[3:-4]]
assert qpi_n == list(range(30, 44)) and len(qpc_n) == 14, (qpi_n, qpc_n)
w("/// Table 8-10: QpC for qPi = 30..43 (qPi < 30: QpC = qPi; qPi > 43: QpC = qPi - 6).")
w("pub static QPC_30_43: [u8; 14] = [%s];" % ", ".join(map(str, qpc_n)))
w("")

# Table 8-12
beta, tc = [], []
for l in body_after(r"Table 8-12 - Derivation of threshold variables"):
    s = l.strip()
    if s.startswith("β"):
        beta += [int(v) for v in s.split()[1:] if v != "-"]
    elif s.startswith("t C"):
        tc += [int(v) for v in s.split()[2:]]
assert len(beta) == 52 and len(tc) == 54, (len(beta), len(tc))
w("/// Table 8-12: beta' for Q = 0..51.")
w("pub static BETA_TABLE: [u8; 52] = [%s];" % ", ".join(map(str, beta)))
w("/// Table 8-12: tC' for Q = 0..53.")
w("pub static TC_TABLE: [u8; 54] = [%s];" % ", ".join(map(str, tc)))
w("")


def taps(eq_no, var, n):
    for l in lines:
        if "(%s)" % eq_no in l:
            expr = l.split("=", 1)[1].split(">>")[0]
            coef = {}
            for sign, c, k in re.findall(r"([+-]?)\s*(?:(\d+)\s*\*\s*)?%s0,\s*(-?\d)" % var, expr):
                coef[int(k)] = (-1 if sign == "-" else 1) * (int(c) if c else 1)
            return coef
    raise SystemExit("equation %s not found" % eq_no)


luma = []
for eq in ["8-236", "8-237", "8-238"]:
    c = taps(eq, "c", 8)
    luma.append([c.get(k, 0) for k in range(-3, 5)])
w("/// Luma 8-tap interpolation filter fL[xFrac][k] (8-236..8-238), taps at offsets -3..4.")
w("pub static LUMA_FILTER: [[i8; 8]; 4] = [[0, 0, 0, 64, 0, 0, 0, 0], %s];" % ", ".join("[%s]" % ", ".join(map(str, t)) for t in luma))
chroma = []
for eq in ["8-255", "8-256", "8-257", "8-258", "8-259", "8-260", "8-261"]:
    c = taps(eq, "aX", 4)
    chroma.append([c.get(k, 0) for k in range(-1, 3)])
w("/// Chroma 4-tap interpolation filter fC[xFrac][k] (8-255..8-261), taps at offsets -1..2.")
w("pub static CHROMA_FILTER: [[i8; 4]; 8] = [[0, 64, 0, 0], %s];" % ", ".join("[%s]" % ", ".join(map(str, t)) for t in chroma))
w("")

# DCT matrix: transMatrix[m][n] (m = sample column 0..31, n = row 0..31); printed as 32 rows of 16 columns twice.
def matrix_rows(eq):
    i = [k for k, l in enumerate(lines) if "(%s)" % eq in l][-1]
    rows = []
    for l in lines[i + 1:]:
        m = re.match(r"^\s*\{([-\d\s]+)\}\s*$", l)
        if m:
            rows.append([int(v) for v in m.group(1).split()])
        elif rows and l.strip() == "}":
            break
    rows = rows[:32]
    assert len(rows) == 32 and all(len(r) == 16 for r in rows), (eq, len(rows))
    return rows


a, b = matrix_rows("8-319"), matrix_rows("8-321")
w("/// DCT matrix (8-319, 8-321): TRANS_MATRIX[n][m] = transMatrix[m][n] (row n = basis function).")
w("pub static TRANS_MATRIX: [[i8; 32]; 32] = [")
for n in range(32):
    w("    [%s]," % ", ".join(str(v) for v in a[n] + b[n]))
w("];")
print("\n".join(out))