"""
vgpu_cache.py — Production VGPU cache.
Drop this in, replace your ad-hoc shader creation with VGPUCache.
"""
import moderngl
import numpy as np
from typing import Optional


_MATMUL_GLSL = r"""
#version 430 core
layout(local_size_x=16, local_size_y=16) in;
layout(std430, binding=0) readonly buffer A { float _A[]; };
layout(std430, binding=1) readonly buffer B { float _B[]; };
layout(std430, binding=2) writeonly buffer C { float _C[]; };
uniform int _N;
void main() {
    uint r = gl_GlobalInvocationID.y;
    uint c = gl_GlobalInvocationID.x;
    if (r >= uint(_N) || c >= uint(_N)) return;
    float s = 0.0;
    for (uint k = 0; k < uint(_N); k++)
        s += _A[r * uint(_N) + k] * _B[k * uint(_N) + c];
    _C[r * uint(_N) + c] = s;
}
"""


class VGPUCache:
    """Single-instance VGPU kernel cache, lazy-compiling on first use."""

    def __init__(self):
        try:
            self.ctx = moderngl.create_standalone_context()
            self.gpu_available = True
        except Exception:
            self.ctx = None
            self.gpu_available = False
        self._kernels = {}      # (op, N) -> PersistentKernel
        self._stats = {'compiles': 0, 'dispatches': 0, 'cache_hits': 0}

    def matmul(self, A, B):
        """Drop-in replacement for np.matmul / A @ B."""
        N = A.shape[0]
        if self.gpu_available:
            key = ('matmul', N)
            if key not in self._kernels:
                self._kernels[key] = PersistentKernel(self.ctx, N)
                self._stats['compiles'] += 1
            else:
                self._stats['cache_hits'] += 1
            result = self._kernels[key](A.astype(np.float32),
                                         B.astype(np.float32))
            self._stats['dispatches'] += 1
            return result
        else:
            return A.astype(np.float32) @ B.astype(np.float32)

    def report(self):
        return (f"VGPUCache: {self._stats['compiles']} compiles, "
                f"{self._stats['cache_hits']} cache hits, "
                f"{self._stats['dispatches']} dispatches")


# ====================================================================
#  PRODUCTION USAGE
# ====================================================================
if __name__ == "__main__":
    cache = VGPUCache()

    N = 16
    num_calls = 100

    print("Running matmul 100x with cache...")
    for i in range(num_calls):
        A = np.random.randn(N, N).astype(np.float32)
        B = np.random.randn(N, N).astype(np.float32)
        C = cache.matmul(A, B)
        if i == 0 or i == num_calls - 1:
            print(f"  Call {i+1}: C[0,0]={C[0,0]:.4f}")

    print(cache.report())