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

Produces sync_rtl/poly_arith/TB/vectors/poly_arith_input.hex
Each line is a 40-bit hex value (10 hex chars, no spaces):
  bits[39:28] = expected[11:0]
  bits[27:16] = coeff_b_in[11:0]
  bits[15:4]  = coeff_a_in[11:0]
  bit[3]      = mode (0=add, 1=sub)
  bits[2:0]   = padding (0)

expected = (coeff_a +/- coeff_b) mod Q,  Q = 3329

Edge cases covered:
  - zeros
  - max values (Q-1 = 3328)
  - mid-range
  - overflow/underflow (a + b >= Q, a - b < 0)
"""

import os
import sys

Q = 3329
Q_MINUS_1 = Q - 1  # 3328

# Output path (relative to project root, where script should be run)
OUTPUT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "vectors")
OUTPUT_FILE = os.path.join(OUTPUT_DIR, "poly_arith_input.hex")


def mod_add(a: int, b: int) -> int:
    """(a + b) mod Q."""
    return (a + b) % Q


def mod_sub(a: int, b: int) -> int:
    """(a - b) mod Q."""
    return (a - b) % Q


def pack_vector(expected: int, coeff_b: int, coeff_a: int, mode: int) -> int:
    """Pack a single test vector into a 40-bit value."""
    val = 0
    val |= (expected & 0xFFF) << 28   # bits 39:28
    val |= (coeff_b  & 0xFFF) << 16   # bits 27:16
    val |= (coeff_a  & 0xFFF) << 4    # bits 15:4
    val |= (mode     & 0x1)   << 3    # bit  3
    # bits 2:0 are padding = 0
    return val


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

    # Helper to add a vector
    def add_vector(a: int, b: int, mode: int) -> None:
        if mode == 0:
            exp = mod_add(a, b)
        else:
            exp = mod_sub(a, b)
        vectors.append(pack_vector(exp, b, a, mode))

    # ---- ADD mode vectors (mode=0) ----
    # Zero + zero
    add_vector(0, 0, 0)

    # Max + max = (3328 + 3328) % 3329 = 6656 % 3329 = 3327
    add_vector(Q_MINUS_1, Q_MINUS_1, 0)

    # Zero + max
    add_vector(0, Q_MINUS_1, 0)
    add_vector(Q_MINUS_1, 0, 0)

    # Mid-range values
    for a, b in [(1000, 2000), (1500, 1500), (1, 2), (42, 137)]:
        add_vector(a, b, 0)

    # Overflow: a + b > Q
    add_vector(2000, 2000, 0)  # 4000 % 3329 = 671
    add_vector(3000, 1000, 0)  # 4000 % 3329 = 671
    add_vector(3328, 1, 0)     # 3329 % 3329 = 0

    # Edge: Q-sized values that just fit
    add_vector(1664, 1665, 0)  # 3329 -> 0
    add_vector(1664, 1664, 0)  # 3328 -> 3328

    # ---- SUB mode vectors (mode=1) ----
    # Zero - zero
    add_vector(0, 0, 1)

    # Max - max
    add_vector(Q_MINUS_1, Q_MINUS_1, 1)

    # Zero - max: (0 - 3328) % 3329 = 1
    add_vector(0, Q_MINUS_1, 1)

    # Max - zero
    add_vector(Q_MINUS_1, 0, 1)

    # Mid-range values
    for a, b in [(2000, 1000), (1500, 1500), (2, 1), (137, 42)]:
        add_vector(a, b, 1)

    # Underflow: a < b -> negative result
    add_vector(1000, 2000, 1)  # (1000 - 2000) % 3329 = 2329
    add_vector(0, 1, 1)        # 3328
    add_vector(1, 3328, 1)     # (1 - 3328) % 3329 = 2
    add_vector(0, 2, 1)        # 3327

    # Edge cases
    add_vector(1, 3328, 1)     # (1 - 3328) % 3329 = 2

    # Run through all combinations of a few edge values
    for a in [0, 1, 1664, 3327, 3328]:
        for b in [0, 1, 1664, 3327, 3328]:
            add_vector(a, b, 0)
            add_vector(a, b, 1)

    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:
            # 40 bits = 10 hex digits
            f.write(f"{v:010X}\n")

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


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

    # Print a few samples for debugging
    print("\nSample vectors (first 5):")
    for i, v in enumerate(vectors[:5]):
        exp    = (v >> 28) & 0xFFF
        coeff_b = (v >> 16) & 0xFFF
        coeff_a = (v >> 4)  & 0xFFF
        mode   = (v >> 3)  & 0x1
        op = "ADD" if mode == 0 else "SUB"
        print(f"  [{i}] a={coeff_a:04d} b={coeff_b:04d} mode={op} expected={exp:04d}")

    return 0


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