# Vector generator for src/vectors/**/*.zig using numpy/scipy

import os
import os.path
import random

import io
import numpy
import scipy


# Floating point precision to round and serialize to
PRECISION = 8


################################################################################
# Helper functions for generating random types
################################################################################


def random_complex64(n):
    return numpy.around(numpy.array([complex(2 * random.random() - 1.0, 2 * random.random() - 1.0) for _ in range(n)]).astype(numpy.complex64), PRECISION)


def random_float32(n):
    return numpy.around(numpy.array([2 * random.random() - 1.0 for _ in range(n)]).astype(numpy.float32), PRECISION)


def random_bit(n):
    return numpy.array([random.randint(0, 1) for _ in range(n)]).astype(numpy.bool_)


################################################################################
# Vector Serialization
################################################################################


def serialize(vector):
    """Serialize a vector into a Zig array."""

    if isinstance(vector, numpy.ndarray):
        if vector.dtype == "float32":
            return f"[{len(vector)}]f32{{" + ("", " ")[len(vector) > 1] + ", ".join([f"{e:.{PRECISION}f}" for e in vector]) + ("", " ")[len(vector) > 1] + "}"
        elif vector.dtype == "complex64":
            return f"[{len(vector)}]std.math.Complex(f32){{" + ("", " ")[len(vector) > 1] + ", ".join([f".{{ .re = {e.real:.{PRECISION}f}, .im = {e.imag:.{PRECISION}f} }}" for e in vector]) + ("", " ")[len(vector) > 1] + "}"
        elif vector.dtype == "bool":
            return f"[{len(vector)}]u1{{" + ("", " ")[len(vector) > 1] + ", ".join([str(int(e)) for e in vector]) + ("", " ")[len(vector) > 1] + "}"
        else:
            raise NotImplementedError(f"Unsupported ndarray data type: {vector.dtype}")
    elif isinstance(vector, bytes):
        return f"[{len(vector)}]u8{{" + ("", " ")[len(vector) > 1] + ", ".join([f'0x{e:02x}' for e in vector]) + ("", " ")[len(vector) > 1] + "}"
    else:
        raise NotImplementedError(f"Unsupported vector type: {type(vector)}")


################################################################################
# Generate
################################################################################


def generate(filename):
    """Generate test vectors in `filename`."""

    # Initialize random seed for determinism
    random.seed(1)

    python_buffer = ""

    with open(filename, "r+") as f:
        # Extract python code from python tag and stop at autogenerated tag
        in_python_tag = False
        for line in iter(f.readline, None):
            if not line or "@autogenerated" in line:
                break
            elif "@python" in line:
                in_python_tag = not in_python_tag
            elif in_python_tag:
                python_buffer += line[3:]

        # Initialize vectors
        vectors = {}

        def vector(name, v):
            print(f"Generating vector {filename}:{name}")
            vectors[name] = v

        # Run code
        exec(python_buffer)

        # Truncate rest of file
        f.truncate(f.tell())
        f.seek(f.tell())

        # Serialize vectors
        f.write("\n")
        for name in vectors:
            f.write(f"pub const {name} = {serialize(vectors[name])};\n")


if __name__ == "__main__":
    # Walk vectors directory
    for root, _, files in os.walk("src/vectors"):
        for file in files:
            if file.endswith(".zig"):
                generate(os.path.join(root, file))