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

Generates input vectors for poly_mul_sync module in hex format.
Each line encodes {256x12-bit poly A, 256x12-bit poly B} = 6144 bits
packed as a single hex number (1536 hex chars).

The poly_mul_sync module uses basecase_mul (degree-1 Karatsuba
multiplication) with zeta from poly_mul_zeta_rom.

For each pair of degree-1 polynomials (A[2k], A[2k+1]) and (B[2k], B[2k+1]):
  c0 = (a0*b0 + a1*b1*zeta) mod Q
  c1 = (a0*b1 + a1*b0) mod Q

Usage:
    python3 gen_vectors.py
Output:
    vectors/poly_mul_input.hex
"""

import sys
import os
import random

# ============================================================
# Constants
# ============================================================
Q = 3329          # ML-KEM modulus
N = 256           # Number of coefficients per polynomial

# poly_mul_zeta_rom values (zeta = zeta_bitRev[k]^2 * 17 mod Q)
# These match sync_rtl/poly_mul/poly_mul_zeta_rom.v exactly
POLY_MUL_ZETA = [
    17,   3312, 2761, 568,  583,  2746, 2649, 680,
    1637, 1692, 723,  2606, 2288, 1041, 1100, 2229,
    1409, 1920, 2662, 667,  3281, 48,   233,  3096,
    756,  2573, 2156, 1173, 3015, 314,  3050, 279,
    1703, 1626, 1651, 1678, 2789, 540,  1789, 1540,
    1847, 1482, 952,  2377, 1461, 1868, 2687, 642,
    939,  2390, 2308, 1021, 2437, 892,  2388, 941,
    733,  2596, 2337, 992,  268,  3061, 641,  2688,
    1584, 1745, 2298, 1031, 2037, 1292, 3220, 109,
    375,  2954, 2549, 780,  2090, 1239, 1645, 1684,
    1063, 2266, 319,  3010, 2773, 556,  757,  2572,
    2099, 1230, 561,  2768, 2466, 863,  2594, 735,
    2804, 525,  1092, 2237, 403,  2926, 1026, 2303,
    1143, 2186, 2150, 1179, 2775, 554,  886,  2443,
    1722, 1607, 1212, 2117, 1874, 1455, 1029, 2300,
    2110, 1219, 2935, 394,  885,  2444, 2154, 1175,
]


def barrett_mul(a: int, b: int) -> int:
    """Barrett modular multiplication: (a * b) mod Q."""
    return (a * b) % Q


def compute_expected(a_coeffs: list, b_coeffs: list) -> list:
    """Compute expected output of poly_mul_sync.

    The poly_mul_sync module processes coefficient pairs using basecase_mul:
    For each k in 0..127:
      a0=A[2k], a1=A[2k+1], b0=B[2k], b1=B[2k+1], zeta=zeta_rom[k]
      c0 = (a0*b0 + a1*b1*zeta) mod Q
      c1 = (a0*b1 + a1*b0) mod Q
    """
    result = [0] * N
    for k in range(N // 2):
        a0 = a_coeffs[2 * k]
        a1 = a_coeffs[2 * k + 1]
        b0 = b_coeffs[2 * k]
        b1 = b_coeffs[2 * k + 1]
        zeta = POLY_MUL_ZETA[k]

        # c0 = (a0*b0 + a1*b1*zeta) mod Q
        t1 = barrett_mul(a0, b0)
        t2 = barrett_mul(a1, b1)
        t2_zeta = barrett_mul(t2, zeta)
        c0 = (t1 + t2_zeta) % Q

        # c1 = (a0*b1 + a1*b0) mod Q
        t3 = barrett_mul(a0, b1)
        t4 = barrett_mul(a1, b0)
        c1 = (t3 + t4) % Q

        result[2 * k] = c0
        result[2 * k + 1] = c1
    return result


def coeffs_pair_to_hex(a_coeffs: list, b_coeffs: list) -> str:
    """Convert 256+256 12-bit coefficients to a 1536-char hex string.

    A coeffs in the MSB half, B coeffs in the LSB half.
    A[0] at the very top, B[255] at the very bottom.
    """
    result = 0
    for c in a_coeffs:
        result = (result << 12) | (c & 0xFFF)
    for c in b_coeffs:
        result = (result << 12) | (c & 0xFFF)
    return f"{result:01536X}"


def write_vector(f, a_coeffs: list, b_coeffs: list, label: str):
    """Write a single test vector to the hex file."""
    hex_str = coeffs_pair_to_hex(a_coeffs, b_coeffs)
    f.write(f"// {label}\n")
    f.write(hex_str + "\n")


def generate_vectors():
    """Generate test vectors for poly_mul_sync."""
    os.makedirs("vectors", exist_ok=True)
    hex_path = os.path.join("vectors", "poly_mul_input.hex")

    tests = []

    # --- Test 0: All zeros ---
    zeros = [0] * N
    expected = [0] * N
    tests.append((zeros, zeros, expected, "A=0, B=0"))

    # --- Test 1: A = 1 at all positions, B = all zeros ---
    a_ones = [1] * N
    tests.append((a_ones, zeros, [0] * N, "A=1, B=0"))

    # --- Test 2: A = [1,0,1,0,...], B = [0,1,0,1,...] ---
    a_10 = [1 if i % 2 == 0 else 0 for i in range(N)]
    b_01 = [0 if i % 2 == 0 else 1 for i in range(N)]
    expected = compute_expected(a_10, b_01)
    tests.append((a_10, b_01, expected, "A=1010..., B=0101..."))

    # --- Test 3: impulse at index 0 for both A and B ---
    a_imp0 = [0] * N
    a_imp0[0] = 1
    b_imp0 = [0] * N
    b_imp0[0] = 1
    expected = compute_expected(a_imp0, b_imp0)
    tests.append((a_imp0, b_imp0, expected, "A=imp[0], B=imp[0]"))

    # --- Test 4: impulse at index 1 for both A and B ---
    a_imp1 = [0] * N
    a_imp1[1] = 1
    b_imp1 = [0] * N
    b_imp1[1] = 1
    expected = compute_expected(a_imp1, b_imp1)
    tests.append((a_imp1, b_imp1, expected, "A=imp[1], B=imp[1]"))

    # --- Test 5: impulse pair (0,1) ---
    a_imp01 = [0] * N
    a_imp01[0] = 1
    a_imp01[1] = 1
    b_imp01 = [0] * N
    b_imp01[0] = 1
    b_imp01[1] = 1
    expected = compute_expected(a_imp01, b_imp01)
    tests.append((a_imp01, b_imp01, expected, "A=imp[0,1], B=imp[0,1]"))

    # --- Test 6: all pairs identity ---
    a_id = [1, 0] * (N // 2)
    b_id = [0, 0] * (N // 2)
    b_id[0] = 1  # B = [1,0,0,0,...]
    expected = compute_expected(a_id, b_id)
    tests.append((a_id, b_id, expected, "A=all pairs (1,0), B=imp[0]"))

    # --- Tests 7-10: random vectors ---
    random.seed(0xBEEF)
    for i in range(4):
        a_rand = [random.randrange(0, Q) for _ in range(N)]
        b_rand = [random.randrange(0, Q) for _ in range(N)]
        expected = compute_expected(a_rand, b_rand)
        tests.append((a_rand, b_rand, expected, f"random {i}"))

    # Write input hex file
    with open(hex_path, "w") as f:
        f.write("// poly_mul_sync test vectors\n")
        f.write("// Format: {768 hex chars A coeffs}{768 hex chars B coeffs}\n")
        f.write("// A coeffs: 256 x 12-bit values, A[0] at MSB\n")
        f.write("// B coeffs: 256 x 12-bit values, B[0] after A[255]\n")
        f.write("\n")
        for a_coeffs, b_coeffs, expected, label in tests:
            write_vector(f, a_coeffs, b_coeffs, label)

    print(f"Generated {len(tests)} test vectors → {hex_path}")

    # Print expected results for reference
    print("\nExpected output summary (first 4 coeffs of each test):")
    for a_coeffs, b_coeffs, expected, label in tests:
        first4 = expected[:4]
        print(f"  {label:35s} → [{first4[0]:4d}, {first4[1]:4d}, {first4[2]:4d}, {first4[3]:4d}, ...]")


if __name__ == "__main__":
    generate_vectors()
