#!/usr/bin/env python3
"""Generate test vectors for mod_add_sync module.

Produces mod_add_input.hex: one 24-bit hex word per line.
Each word encodes {a[11:0], b[11:0]} as a single hex value
suitable for $readmemh.

Also prints expected results to stdout for reference.
"""

import os

Q = 3329

def mod_add(a, b):
    """Compute (a + b) mod Q."""
    s = a + b
    if s >= Q:
        s -= Q
    return s

# Test cases: (a, b, description)
test_cases = [
    # Normal operations
    (0, 0, "zero + zero"),
    (100, 200, "small + small (no overflow)"),
    (500, 1000, "medium + medium (no overflow)"),
    (2000, 1000, "medium pair (no overflow)"),
    (1500, 1500, "medium pair (no overflow)"),
    (3328, 0, "max + zero"),
    (0, 3328, "zero + max"),

    # Near-Q values (still < Q each)
    (3000, 300, "near max + small (no overflow)"),
    (300, 3000, "small + near max (no overflow)"),
    (2000, 1500, "medium pair (overflow by small)"),
    (2000, 2000, "medium pair (overflow)"),
    (2500, 2500, "large pair (overflow)"),
    (3000, 1000, "large + medium (overflow)"),

    # Exact Q boundary
    (3328, 1, "max + 1 (exact overflow to zero)"),
    (1, 3328, "1 + max (exact overflow to zero)"),

    # Max values (both Q-1)
    (3328, 3328, "max + max (double overflow)"),

    # Various overflow amounts
    (2000, 3328, "medium + max (overflow)"),
    (3328, 2000, "max + medium (overflow)"),
    (2500, 1000, "large + medium (overflow)"),
    (1000, 2500, "medium + large (overflow)"),
    (1234, 2100, "random pair (overflow)"),
    (456, 2900, "small + near max (overflow)"),

    # No overflow but close
    (1500, 1828, "pair sums exactly to Q-1"),
    (1829, 1500, "pair sums exactly to Q (overflow -> 0)"),

    # Back-to-back stress (all different values)
    (1111, 2222, "stress pair"),
    (333, 444, "stress pair"),
    (777, 1888, "stress pair"),
    (2999, 333, "stress pair near overflow"),
    (500, 2800, "stress pair near max"),
]

# Restore previously removed concise test cases
additional = [
    (42, 42),
    (777, 1337),
    (1664, 1664),   # exactly Q//2 + Q//2 = Q-1? 1664+1664=3328 = Q-1
    (200, 3100),    # overflow
    (0, 1),
    (1, 0),
    (3327, 2),      # exact overflow: 3329 → 0
    (1000, 2329),   # overflow: 3329 → 0
    (50, 3279),     # overflow: 3329 → 0
    (1600, 1729),   # overflow by 0: 3329 → 0
]
test_cases.extend([(a, b) for a, b in additional])

def format_hex(val, bits):
    """Format integer as zero-padded hex string."""
    nibbles = (bits + 3) // 4
    return f"{val:0{nibbles}x}"

vectors_dir = os.path.join(os.path.dirname(__file__), "vectors")
os.makedirs(vectors_dir, exist_ok=True)

output_path = os.path.join(vectors_dir, "mod_add_input.hex")

with open(output_path, "w") as f:
    for idx, item in enumerate(test_cases):
        if len(item) == 3:
            a, b, desc = item
        else:
            a, b = item
            desc = ""
        # Ensure inputs are in valid range
        assert 0 <= a < Q, f"a={a} out of range"
        assert 0 <= b < Q, f"b={b} out of range"
        # Pack: upper 12 bits = a, lower 12 bits = b
        word = (a << 12) | b
        expected = mod_add(a, b)
        f.write(format_hex(word, 24) + "\n")
        if desc:
            print(f"[{idx:3d}] a={a:4d} b={b:4d} -> expected={expected:4d}  ({desc})")
        else:
            print(f"[{idx:3d}] a={a:4d} b={b:4d} -> expected={expected:4d}")

print(f"\nWrote {len(test_cases)} vectors to {output_path}")
