"""gen_vectors.py - Test vector generator for comp_decomp module.

Generates (mode, d, coeff) input tuples and computes expected
compress/decompress results matching the Python reference implementation.
Uses round-half-up (round(2.5)=3, round(3.5)=4).
"""

import os
import random
import sys

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', '..', 'lib'))

from vector_gen import VectorGenerator

Q = 3329  # ML-KEM prime modulus


def round_customize(a: float) -> int:
    """Round-half-up: round(2.5)=3, round(3.5)=4.
    
    Matches the reference implementation in de_compress.py.
    """
    tmp = int(a) + 0.5
    if a >= tmp:
        return int(a) + 1
    else:
        return int(a)


def compress(z_q: int, q: int, d: int) -> int:
    """Compress coefficient: round((2^d / q) * z_q) mod 2^d."""
    _2d = 2 ** d
    return round_customize((_2d / q) * z_q) % _2d


def decompress(z_d: int, q: int, d: int) -> int:
    """Decompress coefficient: round((q / 2^d) * z_d)."""
    _2d = 2 ** d
    return round_customize((q / _2d) * z_d)


class CompDecompVectorGenerator(VectorGenerator):
    """Generates test vectors for the comp_decomp_sync module."""

    def generate_one(self, params: dict) -> dict:
        """Generate a single (mode, d, coeff) test vector.

        Args:
            params: dict with 'mode' ('compress'/'decompress') and 'd' (int).

        Returns:
            dict with 'input' and 'expected' keys.
        """
        mode = params.get('mode', 'compress')
        d = params.get('d', 10)

        if mode == 'compress':
            coeff = random.randint(0, Q - 1)
            result = compress(coeff, Q, d)
        else:
            coeff = random.randint(0, (2 ** d) - 1)
            result = decompress(coeff, Q, d)

        return {
            'input': {'mode': mode, 'd': d, 'coeff': coeff},
            'expected': {'result': result}
        }

    def write_hex_file(self, vectors: list[dict], filepath: str) -> None:
        """Write input vectors as 'MODE D COEFF' format.

        MODE: 'C' for compress, 'D' for decompress.
        D and COEFF are hex-encoded.

        Example: 'C A 3FF' for compress with d=10, coeff=0x3FF.
        """
        os.makedirs(os.path.dirname(filepath), exist_ok=True)
        with open(filepath, 'w') as f:
            for v in vectors:
                inp = v['input']
                mode_char = 'C' if inp['mode'] == 'compress' else 'D'
                d_val = inp['d']
                coeff_val = inp['coeff']
                f.write(f'{mode_char} {d_val:X} {coeff_val:03X}\n')

    def write_expected_file(self, vectors: list[dict], filepath: str) -> None:
        """Write expected output as zero-padded 3-digit hex values.

        One value per line matching the %03X format in the testbench.
        """
        os.makedirs(os.path.dirname(filepath), exist_ok=True)
        with open(filepath, 'w') as f:
            for v in vectors:
                result = v['expected']['result']
                f.write(f'{result:03X}\n')

    def compare_results(self, got: list[str], expected_file: str) -> bool:
        """Compare RTL output against expected values.

        Args:
            got: List of hex result strings from simulation.
            expected_file: Path to expected hex file.

        Returns:
            bool: True if all results match.
        """
        with open(expected_file, 'r') as f:
            expected = [line.strip() for line in f
                        if line.strip() and not line.startswith('#')]

        if len(got) != len(expected):
            return False

        for i, (g, e) in enumerate(zip(got, expected)):
            if g.upper() != e.upper():
                return False

        return True