import ctypes
import numpy as np
import time
import os

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

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

    # void rsr_gemm_dispatch(const int8_t* A, const int8_t* B, int32_t* C, int N, int K, int M, int block_size)
    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.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int
    ]
    lib.rsr_gemm_dispatch.restype = None

    print("🔧 TRME Autotuner: Finding optimal Block Size...")
    print("-" * 50)

    # Tuning parameters
    block_sizes = [4, 6, 8]
    N = 1024 # Tuning size
    K = 1024
    M = 1024

    A = np.random.randint(0, 4, size=(N, K), dtype=np.int8)
    B = np.random.randint(-5, 6, size=(K, M), dtype=np.int8)
    C = np.zeros((N, M), dtype=np.int32)

    best_time = float('inf')
    best_block = 4

    for bs in block_sizes:
        # Warmup
        lib.rsr_gemm_dispatch(A, B, C, N, K, M, bs)

        # Measure
        start = time.perf_counter()
        for _ in range(5):
            lib.rsr_gemm_dispatch(A, B, C, N, K, M, bs)
        end = time.perf_counter()
        avg_time = (end - start) / 5

        print(f"Block Size {bs}: {avg_time:.6f} s")

        if avg_time < best_time:
            best_time = avg_time
            best_block = bs

    print("-" * 50)
    print(f"✅ Optimal Block Size: {best_block}")
    return best_block

if __name__ == "__main__":
    autotune()
