import numpy as np
import ctypes, time, os

# compile the optimised version
c_code_opt = r'''
#include <stdint.h>

void hyper_dot_opt(const uint8_t *X, const uint8_t *W_T,
                   uint32_t *out,
                   const uint16_t *lut,
                   int M, int K, int N) {
    #pragma omp parallel for collapse(2)
    for (int i = 0; i < M; i++) {
        for (int j = 0; j < N; j++) {
            const uint8_t *row_x = X + i * K;
            const uint8_t *row_w = W_T + j * K;
            uint32_t sum = 0;
            #pragma omp simd reduction(+:sum)
            for (int k = 0; k < K; k++) {
                sum += lut[row_x[k] * 256 + row_w[k]];
            }
            out[i * N + j] = sum;
        }
    }
}
'''

with open("hyper_dot_opt.c", "w") as f:
    f.write(c_code_opt)
os.system("gcc -shared -fPIC -O3 -march=native -fopenmp hyper_dot_opt.c -o libhyper_opt.so")

lib = ctypes.CDLL("./libhyper_opt.so")
lib.hyper_dot_opt.argtypes = [
    ctypes.POINTER(ctypes.c_uint8),  # X
    ctypes.POINTER(ctypes.c_uint8),  # W_T (N x K)
    ctypes.POINTER(ctypes.c_uint32), # out
    ctypes.POINTER(ctypes.c_uint16), # lut
    ctypes.c_int, ctypes.c_int, ctypes.c_int
]

# ---------- data preparation ----------
M, K, N = 100, 784, 100
np.random.seed(42)
X = np.random.randint(0, 256, (M, K), dtype=np.uint8)
W = np.random.randint(0, 256, (K, N), dtype=np.uint8)
W_T = np.ascontiguousarray(W.T)          # zero‑op coordinate change

lut = np.zeros((256, 256), dtype=np.uint16)
for i in range(256):
    lut[i, :] = i * np.arange(256, dtype=np.uint16)
lut = lut.flatten()

# ---------- benchmark ----------
out_hyper = np.zeros((M, N), dtype=np.uint32)
pX = X.ctypes.data_as(ctypes.POINTER(ctypes.c_uint8))
pWT = W_T.ctypes.data_as(ctypes.POINTER(ctypes.c_uint8))
pOut = out_hyper.ctypes.data_as(ctypes.POINTER(ctypes.c_uint32))
pLut = lut.ctypes.data_as(ctypes.POINTER(ctypes.c_uint16))

# warm‑up
lib.hyper_dot_opt(pX, pWT, pOut, pLut, M, K, N)

start = time.perf_counter()
for _ in range(1000):
    lib.hyper_dot_opt(pX, pWT, pOut, pLut, M, K, N)
hyper_opt_time = (time.perf_counter() - start) / 100

# NumPy comparison
Xf = X.astype(np.float32); Wf = W.astype(np.float32)
_ = np.dot(Xf, Wf)
start = time.perf_counter()
for _ in range(1000):
    out_numpy = np.dot(Xf, Wf)
numpy_time = (time.perf_counter() - start) / 100

assert np.allclose(out_hyper, out_numpy), "Mismatch"
print(f"Optimized XYFLOW time : {hyper_opt_time*1e6:.2f} µs")
print(f"NumPy dot time        : {numpy_time*1e6:.2f} µs")
print(f"Speed‑up factor       : {numpy_time/hyper_opt_time:.2f}x")