import ctypes
import numpy as np
import time
import os

def benchmark():
    lib_path = os.path.abspath('./librsr.so')
    if not os.path.exists(lib_path):
        print("Library not found. Please run 'make' first.")
        return

    # Load Library
    try:
        lib = ctypes.CDLL(lib_path)
    except OSError as e:
        print(f"Failed to load library: {e}")
        return

    # Define signature
    # void rsr_gemm_dispatch(const int8_t* A, const int8_t* B, int32_t* C, const int32_t* bias, int N, int K, int M, int block_size)
    try:
        lib.rsr_gemm_dispatch.argtypes = [
            np.ctypeslib.ndpointer(dtype=np.int8, ndim=2, flags='C_CONTIGUOUS'),
            np.ctypeslib.ndpointer(dtype=np.int8, ndim=2, flags='C_CONTIGUOUS'),
            np.ctypeslib.ndpointer(dtype=np.int32, ndim=2, flags='C_CONTIGUOUS'),
            ctypes.POINTER(ctypes.c_int32), # Bias
            ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int
        ]
        lib.rsr_gemm_dispatch.restype = None
    except AttributeError:
         print("Error: rsr_gemm_dispatch not found.")
         return

    # void rsr_gemm_avx2(const int8_t* A, const int8_t* B, int32_t* C, int N, int K, int M)
    try:
        lib.rsr_gemm_avx2.argtypes = [
            np.ctypeslib.ndpointer(dtype=np.int8, ndim=2, flags='C_CONTIGUOUS'),
            np.ctypeslib.ndpointer(dtype=np.int8, ndim=2, flags='C_CONTIGUOUS'),
            np.ctypeslib.ndpointer(dtype=np.int32, ndim=2, flags='C_CONTIGUOUS'),
            ctypes.c_int, ctypes.c_int, ctypes.c_int
        ]
        lib.rsr_gemm_avx2.restype = None
        has_avx2 = True
    except AttributeError:
        print("Warning: rsr_gemm_avx2 not found in library.")
        has_avx2 = False

    # Benchmark Sizes
    sizes = [256, 512, 1024, 2048]
    block_size = 4 # Optimal

    print(f"{'N':<8} {'Numpy':<10} {'RSR(Templ)':<12} {'RSR+AVX2':<10} {'Spd(Tmpl)':<10} {'Spd(AVX)':<10}")
    print("-" * 65)

    for N in sizes:
        K = N
        M = N

        # Generate A in range [0, 3] (2-bit)
        A = np.random.randint(0, 4, size=(N, K), dtype=np.int8)
        # Generate B in range [-5, 5]
        B = np.random.randint(-5, 6, size=(K, M), dtype=np.int8)

        C_rsr = np.zeros((N, M), dtype=np.int32)
        C_avx = np.zeros((N, M), dtype=np.int32)

        # 1. Numpy (BLAS)
        start = time.perf_counter()
        C_ref = np.dot(A.astype(np.int32), B.astype(np.int32))
        end = time.perf_counter()
        t_numpy = end - start

        # 2. RSR Templated
        start = time.perf_counter()
        null_ptr = ctypes.POINTER(ctypes.c_int32)()
        lib.rsr_gemm_dispatch(A, B, C_rsr, null_ptr, N, K, M, block_size)
        end = time.perf_counter()
        t_rsr = end - start

        # 3. RSR AVX2
        t_avx = 0
        if has_avx2:
            start = time.perf_counter()
            lib.rsr_gemm_avx2(A, B, C_avx, N, K, M)
            end = time.perf_counter()
            t_avx = end - start

        # Validation
        # diff = np.abs(C_ref - C_rsr).max()
        # diff_avx = np.abs(C_ref - C_avx).max()
        # print(f"Diff: {diff}, {diff_avx}")

        speedup_rsr = t_numpy / t_rsr if t_rsr > 0 else 0
        speedup_avx = t_numpy / t_avx if t_avx > 0 else 0

        print(f"{N:<8} {t_numpy:<10.4f} {t_rsr:<12.4f} {t_avx:<10.4f} {speedup_rsr:<10.1f} {speedup_avx:<10.1f}")

if __name__ == "__main__":
    benchmark()
