#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-3.0-or-later
"""Losslessly package NR2 Gamma-prior tables for SpectralNR.

The input is the raw table format used by the public WDSP calculus file:
two consecutive 241 x 241 planes of little-endian IEEE-754 doubles (GG,
then GGS). The output is a generated C++ include containing a compressed,
base64-encoded byte-shuffled two-dimensional delta representation.

This script does not generate the mathematical values. Build and run
generate_nr2_gamma_tables.cpp for a first-principles approximation and an
exhaustive, fail-closed comparison with a reference calculus file.
"""

from __future__ import annotations

import argparse
import base64
import hashlib
import math
from pathlib import Path
import struct
import zlib


GRID_SIZE = 241
PLANE_SIZE = GRID_SIZE * GRID_SIZE
VALUE_COUNT = 2 * PLANE_SIZE
RAW_SIZE = VALUE_COUNT * 8
UINT64_MASK = (1 << 64) - 1


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input", required=True, type=Path)
    parser.add_argument("--output", required=True, type=Path)
    return parser.parse_args()


def validate(raw: bytes) -> tuple[float, ...]:
    if len(raw) != RAW_SIZE:
        raise ValueError(
            f"expected {RAW_SIZE} bytes ({VALUE_COUNT} doubles), got {len(raw)}"
        )
    values = struct.unpack(f"<{VALUE_COUNT}d", raw)
    for index, value in enumerate(values):
        if not math.isfinite(value) or value < 0.0:
            raise ValueError(f"invalid table value at index {index}: {value}")
    return values


def encode(raw: bytes) -> bytes:
    bits = struct.unpack(f"<{VALUE_COUNT}Q", raw)
    residuals: list[int] = []
    for index, value in enumerate(bits):
        plane_index = index % PLANE_SIZE
        row, column = divmod(plane_index, GRID_SIZE)
        left = bits[index - 1] if column > 0 else 0
        up = bits[index - GRID_SIZE] if row > 0 else 0
        upper_left = (
            bits[index - GRID_SIZE - 1] if row > 0 and column > 0 else 0
        )
        residuals.append((value - left - up + upper_left) & UINT64_MASK)

    little_endian = [value.to_bytes(8, "little") for value in residuals]
    shuffled = b"".join(
        bytes(value[byte_index] for value in little_endian)
        for byte_index in range(8)
    )
    # qUncompress expects a four-byte big-endian uncompressed-size prefix.
    qcompressed = struct.pack(">I", len(shuffled)) + zlib.compress(shuffled, 9)
    return base64.b64encode(qcompressed)


def write_include(output: Path, encoded: bytes, raw_sha256: str) -> None:
    lines = [
        "/* SPDX-License-Identifier: GPL-2.0-or-later",
        " * Exact GG/GGS values from TAPR/OpenHPSDR-wdsp Source/calculus.",
        " * Generated by tools/embed_nr2_gamma_tables.py; do not edit.",
        " */",
        f'constexpr char kNr2GammaTablesRawSha256[] = "{raw_sha256}";',
        f"constexpr int kNr2GammaTablesEncodedSize = {len(encoded)};",
        "constexpr const char* const kNr2GammaTablesCompressedBase64Chunks[] = {",
    ]
    text = encoded.decode("ascii")
    lines.extend(
        f'    "{text[index:index + 100]}",'
        for index in range(0, len(text), 100)
    )
    lines.append("};")
    output.write_text("\n".join(lines) + "\n", encoding="ascii")


def main() -> int:
    args = parse_args()
    raw = args.input.read_bytes()
    values = validate(raw)
    encoded = encode(raw)
    raw_sha256 = hashlib.sha256(raw).hexdigest()
    write_include(args.output, encoded, raw_sha256)
    print(f"input values:       {len(values)}")
    print(f"input bytes:        {len(raw)}")
    print(f"encoded characters: {len(encoded)}")
    print(f"raw SHA-256:        {raw_sha256}")
    print(f"wrote:              {args.output}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
