"""
vgpu_cache.py — Universal VGPU persistent kernel cache.
Works on all platforms: RPi5 (numpy), Linux/macOS/Windows (OpenGL/Metal).

Author: Per Lindholm
License: MIT

Run: python vgpu_cache.py
"""
import numpy as np
import os
import sys
import platform

# ====================================================================
#  PLATFORM DETECTION
# ====================================================================
IS_RASPBERRY_PI = platform.machine().startswith('arm') and 'linux' in sys.platform
IS_MACOS = platform.system() == 'Darwin'
IS_WINDOWS = platform.system() == 'Windows'

print(f"[VGPU] Platform: {platform.system()} {platform.machine()}")
print(f"[VGPU] Raspberry Pi: {IS_RASPBERRY_PI}")

# ====================================================================
#  GPU BACKEND (optional)
# ====================================================================
GPU_AVAILABLE = False
GPU_BACKEND = None

# Try OpenGL (Linux/macOS/Windows)
try:
    import moderngl
    
    if IS_RASPBERRY_PI:
        # RPi5: use simple numpy (OpenGL compute shaders unreliable)
        print("[VGPU] RPi5 detected: using CPU (numpy) backend")
        GPU_AVAILABLE = False
    else:
        # Try GPU for desktop platforms
        try:
            ctx = moderngl.create_standalone_context()
            if hasattr(ctx, 'compute_shader') and ctx.compute_shader:
                GPU_AVAILABLE = True
                GPU_BACKEND = 'moderngl'
                print("[VGPU] ✅ GPU available (moderngl)")
            else:
                print("[VGPU] ⚠️  No compute shaders; using CPU")
        except Exception as e:
            print(f"[VGPU] ❌ GPU init failed: {e}; using CPU")
            
except ImportError:
    print("[VGPU] moderngl not installed; using CPU (numpy)")

# For macOS, try Metal (not implemented yet - future work)
if IS_MACOS and not GPU_AVAILABLE:
    print("[VGPU] macOS detected: consider installing Metal support (future)")

# ====================================================================
#  CPU Matmul with Caching
# ====================================================================
class VGPUCache:
    def __init__(self):
        self._kernels = {}
        self._stats = {
            'compiles': 0,
            'cache_hits': 0,
            'dispatches': 0,
            'cpu_calls': 0,
            'gpu_calls': 0,
        }
        
        self.gpu_available = GPU_AVAILABLE
        self.gpu_backend = GPU_BACKEND
        
        if self.gpu_available and self.gpu_backend == 'moderngl':
            try:
                self.ctx = moderngl.create_standalone_context()
            except:
                self.gpu_available = False
                self.gpu_backend = None
                print("[VGPU] GPU context creation failed; using CPU")
        else:
            self.ctx = None
            
        print(f"[VGPU] Using {'GPU' if self.gpu_available else 'CPU'} backend")

    def matmul(self, A, B) -> np.ndarray:
        """A @ B. Cached on (M, K, N) shape."""
        M, K = A.shape
        K2, N = B.shape
        assert K == K2, f"Inner dimensions must match: {K} vs {K2}"
        
        A = A.astype(np.float32)
        B = B.astype(np.float32)
        key = (M, K, N)
        
        # Check cache
        if key in self._kernels:
            self._stats['cache_hits'] += 1
        else:
            self._kernels[key] = None  # placeholder
            self._stats['compiles'] += 1
        
        # Try GPU if available
        if self.gpu_available and self.gpu_backend == 'moderngl':
            try:
                result = self._matmul_gpu(A, B, key)
                self._stats['gpu_calls'] += 1
                self._stats['dispatches'] += 1
                return result
            except Exception as e:
                print(f"[VGPU] GPU matmul failed: {e}; falling back to CPU")
                self.gpu_available = False
        
        # CPU fallback
        self._stats['cpu_calls'] += 1
        self._stats['dispatches'] += 1
        return A @ B

    def _matmul_gpu(self, A, B, key):
        """GPU matmul using moderngl (desktop only)."""
        # Lazy compile on first use
        if self._kernels[key] is None:
            self._kernels[key] = self._compile_kernel(key)
        
        kernel = self._kernels[key]
        return kernel(A, B)

    def _compile_kernel(self, key):
        """Compile GLSL compute shader for given shape."""
        M, K, N = key
        
        # Build GLSL shader
        shader_src = f"""
#version 430 core
layout(local_size_x = 16, local_size_y = 16) in;
layout(std430, binding = 0) readonly buffer A_buf {{ float A[]; }};
layout(std430, binding = 1) readonly buffer B_buf {{ float B[]; }};
layout(std430, binding = 2) writeonly buffer C_buf {{ float C[]; }};
uniform int M = {M};
uniform int K = {K};
uniform int N = {N};
void main() {{
    uint row = gl_GlobalInvocationID.y;
    uint col = gl_GlobalInvocationID.x;
    if (row >= uint(M) || col >= uint(N)) return;
    float sum = 0.0;
    for (uint k = 0u; k < uint(K); k++) {{
        sum += A[row * uint(K) + k] * B[k * uint(N) + col];
    }}
    C[row * uint(N) + col] = sum;
}}
"""
        # Compile and setup buffers
        program = self.ctx.compute_shader(shader_src)
        A_buf = self.ctx.buffer(reserve=M * K * 4)
        B_buf = self.ctx.buffer(reserve=K * N * 4)
        C_buf = self.ctx.buffer(reserve=M * N * 4)
        
        gx = max(1, (N + 15) // 16)
        gy = max(1, (M + 15) // 16)
        
        class GPUKernel:
            def __init__(self, ctx, program, A_buf, B_buf, C_buf, gx, gy):
                self.ctx = ctx
                self.program = program
                self.A_buf = A_buf
                self.B_buf = B_buf
                self.C_buf = C_buf
                self.gx = gx
                self.gy = gy
            
            def __call__(self, A, B):
                self.A_buf.write(A.tobytes())
                self.B_buf.write(B.tobytes())
                self.A_buf.bind_to_storage_buffer(0)
                self.B_buf.bind_to_storage_buffer(1)
                self.C_buf.bind_to_storage_buffer(2)
                self.program.run(self.gx, self.gy, 1)
                return np.frombuffer(self.C_buf.read(), dtype=np.float32).reshape(A.shape[0], B.shape[1]).copy()
        
        print(f"[VGPU] ✅ Compiled kernel for {M}x{K}x{N}")
        return GPUKernel(self.ctx, program, A_buf, B_buf, C_buf, gx, gy)

    def report(self) -> str:
        backend = "GPU" if self.gpu_available else "CPU"
        return (f"VGPUCache [{backend}]: "
                f"{self._stats['compiles']} compiles, "
                f"{self._stats['cache_hits']} cache hits, "
                f"{self._stats['dispatches']} dispatches, "
                f"{self._stats['cpu_calls']} CPU, "
                f"{self._stats['gpu_calls']} GPU")

    def clear(self):
        """Clear all cached kernels."""
        self._kernels.clear()


# ====================================================================
#  SELF-TEST
# ====================================================================
if __name__ == "__main__":
    print("=" * 70)
    print("VGPU Universal Cache Test")
    print("=" * 70)
    
    cache = VGPUCache()
    
    # Test small matmul
    print("\n[1] Testing small matmul (16x16):")
    np.random.seed(42)
    A = np.random.randn(16, 16).astype(np.float32)
    B = np.random.randn(16, 16).astype(np.float32)
    expected = A @ B
    
    C = cache.matmul(A, B)
    diff = np.max(np.abs(C - expected))
    print(f"  Max diff: {diff:.2e}")
    assert diff < 1e-4, f"Matmul failed: diff={diff}"
    print("  ✓ Verified")
    
    # Test rectangular
    print("\n[2] Testing rectangular matmul (32x64 @ 64x16):")
    A = np.random.randn(32, 64).astype(np.float32)
    B = np.random.randn(64, 16).astype(np.float32)
    expected = A @ B
    
    C = cache.matmul(A, B)
    diff = np.max(np.abs(C - expected))
    print(f"  Max diff: {diff:.2e}")
    assert diff < 1e-4, f"Matmul failed: diff={diff}"
    print("  ✓ Verified")
    
    # Test cache hit
    print("\n[3] Testing cache (should hit):")
    C2 = cache.matmul(A, B)
    diff2 = np.max(np.abs(C2 - expected))
    print(f"  Max diff: {diff2:.2e}")
    print(f"  Cache stats: {cache.report()}")
    
    # Benchmark
    print("\n[4] Benchmarking 100 matmuls:")
    import time
    start = time.time()
    for _ in range(100):
        A = np.random.randn(32, 32).astype(np.float32)
        B = np.random.randn(32, 32).astype(np.float32)
        C = cache.matmul(A, B)
    elapsed = time.time() - start
    print(f"  100 matmuls in {elapsed:.3f}s ({elapsed/100*1000:.1f}ms each)")
    print(f"  Final cache stats: {cache.report()}")
    
    print("\n" + "=" * 70)
    print("✅ All tests passed!")
    print("=" * 70)