#!/usr/bin/env python3
"""gen_vectors.py - Generate test vectors for comp_decomp_sync XSIM testbench.

Produces sync_rtl/comp_decomp/TB/vectors/comp_decomp_input.hex
Each line is a 32-bit hex value (8 hex chars, no spaces):
  bits[31:26] = padding (0)
  bits[25:14] = expected[11:0]
  bits[13:2]  = coeff_in[11:0]
  bit[1]      = d[0]   (d[4:0] split across bits [6:2] and bit[1])
  bit[0]      = mode (0=compress, 1=decompress)

Revised layout (cleaner, 32 bits):
  bits[31:20] = expected[11:0]
  bits[19:8]  = coeff_in[11:0]
  bits[7:3]   = d[4:0]
  bit[2]      = mode (0=compress, 1=decompress)
  bits[1:0]   = padding (0)

FIPS 203 formulas:
  Compress_q(x, d) = round((2^d / Q) * x) mod 2^d
                   = ((x * 2^d + Q//2) // Q) & ((1 << d) - 1)
  Decompress_q(y, d) = round((Q / 2^d) * y)
                     = (y * Q + (1 << (d-1))) // (1 << d)

Tests cover d in {4, 5, 10, 11} (ML-KEM standard values) plus edge d=1.
"""

import os
import sys

Q = 3329

OUTPUT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "vectors")
OUTPUT_FILE = os.path.join(OUTPUT_DIR, "comp_decomp_input.hex")


def compress_q(x: int, d: int) -> int:
    """Compress_q(x, d) per FIPS 203."""
    if d == 0:
        return 0
    two_d = 1 << d
    # round((2^d / Q) * x) = floor((x * 2^d + Q/2) / Q)
    val = (x * two_d + Q // 2) // Q
    return val & (two_d - 1)


def decompress_q(y: int, d: int) -> int:
    """Decompress_q(y, d) per FIPS 203."""
    if d == 0:
        return 0
    two_d = 1 << d
    # round((Q / 2^d) * y) = floor((y * Q + 2^(d-1)) / 2^d)
    val = (y * Q + (two_d >> 1)) // two_d
    return val % Q


def pack_vector(expected: int, coeff_in: int, d: int, mode: int) -> int:
    """Pack a single test vector into a 32-bit value."""
    val = 0
    val |= (expected & 0xFFF) << 20   # bits 31:20
    val |= (coeff_in & 0xFFF) << 8    # bits 19:8
    val |= (d       & 0x1F)  << 3    # bits 7:3
    val |= (mode    & 0x1)   << 2    # bit  2
    # bits 1:0 are padding = 0
    return val


def generate() -> list[int]:
    """Generate all test vectors. Returns list of packed 32-bit values."""
    vectors: list[int] = []

    def add_vector(x: int, d: int, mode: int) -> None:
        if mode == 0:
            exp = compress_q(x, d)
        else:
            exp = decompress_q(x, d)
        vectors.append(pack_vector(exp, x, d, mode))

    # Standard ML-KEM d values
    d_values = [4, 5, 10, 11]

    # Test coefficient values: edges and mid-range
    coeffs = [
        0,            # zero
        1,            # minimal
        3328,         # max (Q-1)
        1000,         # mid-range
        2000,         # mid-range
        1664,         # Q/2
        42,           # small
    ]

    # ---- COMPRESS (mode=0) ----
    for d in d_values:
        for c in coeffs:
            add_vector(c, d, 0)

        # Compress edge: max input
        add_vector(3328, d, 0)

        # Compress: some systematic sweep
        for c in [0, 500, 1000, 1500, 2000, 2500, 3000, 3328]:
            add_vector(c, d, 0)

    # ---- DECOMPRESS (mode=1) ----
    # For decompress, input is in [0, 2^d-1]
    for d in d_values:
        two_d_mask = (1 << d) - 1

        # Zero
        add_vector(0, d, 1)

        # Max in range
        add_vector(two_d_mask, d, 1)

        # Mid-range
        mid = two_d_mask // 2
        if mid > 0:
            add_vector(mid, d, 1)
            add_vector(mid - 1, d, 1) if mid > 1 else None
            add_vector(mid + 1, d, 1) if mid + 1 <= two_d_mask else None

        # Systematic sweep through valid range
        step = max(1, two_d_mask // 8)
        for y in range(0, two_d_mask + 1, step):
            add_vector(y, d, 1)

    # ---- Edge case: d=1 (minimum non-zero) ----
    add_vector(0, 1, 0)
    add_vector(3328, 1, 0)
    add_vector(0, 1, 1)
    add_vector(1, 1, 1)

    # ---- d=12 (max for 12-bit operands, though not in ML-KEM) ----
    add_vector(0, 12, 0)
    add_vector(3328, 12, 0)

    return vectors


def write_vectors(vectors: list[int]) -> None:
    """Write vectors to hex file."""
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    with open(OUTPUT_FILE, "w") as f:
        for v in vectors:
            # 32 bits = 8 hex digits
            f.write(f"{v:08X}\n")

    print(f"Generated {len(vectors)} test vectors -> {OUTPUT_FILE}")


def main() -> int:
    vectors = generate()
    write_vectors(vectors)

    # Print statistics and samples
    print(f"\nTotal vectors: {len(vectors)}")
    compress_count = sum(1 for v in vectors if ((v >> 2) & 1) == 0)
    decompress_count = sum(1 for v in vectors if ((v >> 2) & 1) == 1)
    print(f"  Compress:   {compress_count}")
    print(f"  Decompress: {decompress_count}")

    print("\nSample vectors (first 5):")
    for i, v in enumerate(vectors[:5]):
        exp     = (v >> 20) & 0xFFF
        coeff   = (v >> 8)  & 0xFFF
        d       = (v >> 3)  & 0x1F
        mode    = (v >> 2)  & 0x1
        op = "COMPRESS" if mode == 0 else "DECOMPRESS"
        if mode == 0:
            print(f"  [{i}] {op} x={coeff:04d} d={d:02d} expected={exp:04d}")
        else:
            print(f"  [{i}] {op} y={coeff:04d} d={d:02d} expected={exp:04d}")

    return 0


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