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

Generates input vectors for ntt_core module in hex format compatible
with $readmemh. Each line encodes {1-bit mode, 256x12-bit coefficients}
packed as a single 3073-bit hex number (769 hex chars).

Q = 3329, N = 256, primitive root zeta = 17.
NTT uses Cooley-Tukey DIT with bit-reversed zeta ROM.
INTT uses Gentleman-Sande DIF with zeta ROM in reverse.

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

import sys
import os
import random

# ============================================================
# Constants
# ============================================================
Q = 3329          # ML-KEM modulus
N = 256           # NTT size
LAYERS = 7        # log2(N) - 1 for Cooley-Tukey

# zeta_rom values: bit-reversed powers of the primitive root zeta=17
# These match sync_rtl/ntt/zeta_rom.v exactly
ZETA_ROM = [
    1,    1729, 2580, 3289, 2642, 630,  1897, 848,
    1062, 1919, 193,  797,  2786, 3260, 569,  1746,
    296,  2447, 1339, 1476, 3046, 56,   2240, 1333,
    1426, 2094, 535,  2882, 2393, 2879, 1974, 821,
    289,  331,  3253, 1756, 1197, 2304, 2277, 2055,
    650,  1977, 2513, 632,  2865, 33,   1320, 1915,
    2319, 1435, 807,  452,  1438, 2868, 1534, 2402,
    2647, 2617, 1481, 648,  2474, 3110, 1227, 910,
    17,   2761, 583,  2649, 1637, 723,  2288, 1100,
    1409, 2662, 3281, 233,  756,  2156, 3015, 3050,
    1703, 1651, 2789, 1789, 1847, 952,  1461, 2687,
    939,  2308, 2437, 2388, 733,  2337, 268,  641,
    1584, 2298, 2037, 3220, 375,  2549, 2090, 1645,
    1063, 319,  2773, 757,  2099, 561,  2466, 2594,
    2804, 1092, 403,  1026, 1143, 2150, 2775, 886,
    1722, 1212, 1874, 1029, 2110, 2935, 885,  2154,
]


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


def ntt_forward(coeffs: list) -> list:
    """Forward NTT: Cooley-Tukey DIT, matches ntt_core RTL exactly.

    Processes 7 layers. At each layer:
      - Pairs are (j, j+layer_len)
      - Zeta for each block comes from zeta_rom (index increments per block)
      - Butterfly: t = zeta*b; a_out = a+t; b_out = a-t (all mod Q)

    Input:  normal-order coefficients (index 0..255)
    Output: bit-reversed NTT result
    """
    a = list(coeffs)
    layer_len = 128
    zeta_idx = 1  # Forward NTT starts at zeta_rom[1]

    for layer in range(LAYERS):
        for start in range(0, N, 2 * layer_len):
            zeta = ZETA_ROM[zeta_idx]
            for j in range(start, start + layer_len):
                t = barrett_mul(zeta, a[j + layer_len])
                a_j_plus_len = (a[j] - t) % Q
                a[j] = (a[j] + t) % Q
                a[j + layer_len] = a_j_plus_len
            zeta_idx += 1
        layer_len >>= 1

    return a


def intt_inverse(coeffs: list) -> list:
    """Inverse NTT: Gentleman-Sande DIF, matches ntt_core RTL exactly.

    Processes 7 layers in reverse order (len=2,4,8,...,128).
      - Zeta for each block comes from zeta_rom (index decrements per block)
      - Butterfly: a_out = a+b; diff = b-a; b_out = zeta*diff (all mod Q)

    After all layers, output is scaled by N (no 1/2 factors in GS butterfly).
    The RTL then scales output by 3303 in the OUTPUT state for mode=1.
    """
    a = list(coeffs)
    layer_len = 2
    zeta_idx = 127  # Inverse NTT starts at zeta_rom[127]

    for layer in range(LAYERS):
        for start in range(0, N, 2 * layer_len):
            zeta = ZETA_ROM[zeta_idx]
            for j in range(start, start + layer_len):
                a_sum = (a[j] + a[j + layer_len]) % Q
                diff = (a[j + layer_len] - a[j]) % Q
                a[j] = a_sum
                a[j + layer_len] = barrett_mul(zeta, diff)
            zeta_idx -= 1
        layer_len <<= 1

    # Apply output scaling (multiply by 3303) as the RTL does in mode=1
    for i in range(N):
        a[i] = barrett_mul(a[i], 3303)

    return a


def coeffs_to_hex(coeffs: list) -> str:
    """Convert 256 12-bit coefficients to a 768-char hex string.

    coeffs[0] is the MSB of the hex output, coeffs[255] is the LSB.
    Each coefficient is 3 hex chars (12 bits).
    """
    result = 0
    for c in coeffs:
        result = (result << 12) | (c & 0xFFF)
    return f"{result:0768X}"


def write_vector(f, mode: int, coeffs: list, label: str):
    """Write a single test vector to the hex file.

    Format: {mode_hex_digit}{768 hex chars for 256 coeffs}
    Total: 769 hex chars per line.
    mode=0 -> hex digit '0', mode=1 -> hex digit '1'.
    """
    mode_hex = "0" if mode == 0 else "1"
    coeffs_hex = coeffs_to_hex(coeffs)
    line = mode_hex + coeffs_hex
    f.write(f"// {label}\n")
    f.write(line + "\n")


def hex_char_to_int(c: str) -> int:
    """Convert single hex char to integer."""
    return int(c, 16)


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

    # All tests are listed here with labels
    tests = []

    # --- Test 0: Forward NTT on all zeros ---
    zeros = [0] * N
    expected = ntt_forward(zeros)
    tests.append((0, zeros, expected, "FWD: all zeros"))

    # --- Test 1: Forward NTT on impulse at index 0 ---
    imp0 = [0] * N
    imp0[0] = 1
    expected = ntt_forward(imp0)
    tests.append((0, imp0, expected, "FWD: impulse at [0]"))

    # --- Test 2: Forward NTT on impulse at index 1 ---
    imp1 = [0] * N
    imp1[1] = 1
    expected = ntt_forward(imp1)
    tests.append((0, imp1, expected, "FWD: impulse at [1]"))

    # --- Test 3: Forward NTT on ramp [0,1,2,...,255] ---
    ramp = [i % Q for i in range(N)]
    expected = ntt_forward(ramp)
    tests.append((0, ramp, expected, "FWD: ramp 0..255"))

    # --- Test 4: Forward NTT on all ones ---
    ones = [1] * N
    expected = ntt_forward(ones)
    tests.append((0, ones, expected, "FWD: all ones"))

    # --- Test 5: Inverse NTT on all zeros ---
    expected = intt_inverse(zeros)
    tests.append((1, zeros, expected, "INV: all zeros"))

    # --- Test 6: Inverse NTT on impulse at index 0 ---
    expected = intt_inverse(imp0)
    tests.append((1, imp0, expected, "INV: impulse at [0]"))

    # --- Test 7: Inverse NTT on NTT(impulse) → should recover impulse*256 ---
    ntt_imp0 = ntt_forward(imp0)
    expected = intt_inverse(ntt_imp0)
    tests.append((1, ntt_imp0, expected, "INV(NTT(imp[0])) → imp[0]*256"))

    # --- Test 8: Inverse NTT on NTT(ramp) → should recover ramp*256 ---
    ntt_ramp = ntt_forward(ramp)
    expected = intt_inverse(ntt_ramp)
    tests.append((1, ntt_ramp, expected, "INV(NTT(ramp)) → ramp*256"))

    # --- Tests 9-12: Forward NTT on random vectors ---
    random.seed(0x5EED)
    for i in range(4):
        rand_vec = [random.randrange(0, Q) for _ in range(N)]
        expected = ntt_forward(rand_vec)
        tests.append((0, rand_vec, expected, f"FWD: random {i}"))

    # --- Tests 13-16: Inverse NTT on random vectors ---
    for i in range(4):
        rand_vec = [random.randrange(0, Q) for _ in range(N)]
        expected = intt_inverse(rand_vec)
        tests.append((1, rand_vec, expected, f"INV: random {i}"))

    # --- Tests 17-18: Inverse NTT on NTT(random) → roundtrip ---
    for i in range(2):
        rand_vec = [random.randrange(0, Q) for _ in range(N)]
        ntt_vec = ntt_forward(rand_vec)
        recovered = intt_inverse(ntt_vec)
        tests.append((1, ntt_vec, recovered, f"INV(NTT(random {i})) → random*256"))

    # Write input hex file
    with open(hex_path, "w") as f:
        f.write("// ntt_core test vectors\n")
        f.write("// Format: {1 hex digit mode}{768 hex chars coeffs}\n")
        f.write("// mode: 0=forward NTT, 1=inverse NTT\n")
        f.write("// coeffs: 256 x 12-bit values, coeff[0] at MSB position\n")
        f.write("\n")
        for mode, coeffs, expected, label in tests:
            write_vector(f, mode, coeffs, label)

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

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


if __name__ == "__main__":
    generate_vectors()