import numpy as np
import os
import sys


def reverse_endian(data):
    result = bytearray()
    for i in range(len(data), 0, -1):
        result.append(data[i - 1])
    return bytes(result)


def generate_numbers(cases, dtype=np.float16):
    if dtype == np.float16:
        dsize = 2
    elif dtype == np.float32:
        dsize = 4
    else:
        raise ValueError(f"Unknown dtype {dtype}")
    x = np.frombuffer(os.urandom(cases * dsize), dtype=dtype)
    y = np.frombuffer(os.urandom(cases * dsize), dtype=dtype)
    return x, y, dsize


def generate_fp_vector(cases, filename, dtype=np.float16, big_endian=False, comp_file=None):
    x, y, dsize = generate_numbers(cases, dtype)
    np.seterr(all='ignore')
    sum = x + y
    mul = x * y
    with open(filename, 'w') as f:
        for i in range(cases):
            t = lambda v: reverse_endian(v) if big_endian else v
            f.write(' '.join([
                t(x.tobytes()[i * dsize:i * dsize + dsize]).hex(),
                t(y.tobytes()[i * dsize:i * dsize + dsize]).hex(),
                t(sum.tobytes()[i * dsize:i * dsize + dsize]).hex(),
                t(mul.tobytes()[i * dsize:i * dsize + dsize]).hex(),
            ]) + '\n')
    if comp_file is not None:
        gt = x > y
        lt = x < y
        ge = x >= y
        le = x <= y
        eq = x == y
        with open(comp_file, 'w') as f:
            for i in range(cases):
                f.write(''.join([
                    '1' if gt[i] else '0',
                    '1' if lt[i] else '0',
                    '1' if ge[i] else '0',
                    '1' if le[i] else '0',
                    '1' if eq[i] else '0'
                ]) + f'  // {x[i]:10.3e} {y[i]:10.3e}\n')


if __name__ == '__main__':
    generate = 50
    if len(sys.argv) == 2 and sys.argv[1].isdigit():
        generate = int(sys.argv[1])
    else:
        print(f"Usage: {sys.argv[0]} [number of generated tests]")

    generate_fp_vector(generate, 'fp16_test.hex', dtype=np.float16, big_endian=True)
    generate_fp_vector(generate, 'fp32_test.hex', dtype=np.float32, big_endian=True, comp_file='fp32_test_comp.hex')