import numpy as np
import time
from vgpu_cache import VGPUCache

def solve_2d_heat(nx=128, ny=128, steps=200, gamma=0.2):
    """
    Solves the 2D Heat Equation PDE using VGPUCache vector ops.
    gamma = alpha * dt / (dx**2) must be <= 0.25 for numerical stability.
    """
    c = VGPUCache()
    
    # 1. Initialize temperature grid with a hot spot in the center
    u_init = np.zeros((nx, ny), dtype=np.float32)
    u_init[nx//4:3*nx//4, ny//4:3*ny//4] = 100.0  # Hot square
    
    # We will maintain boundary conditions manually or via padding
    u = u_init.copy()
    
    print(f"Starting PDE simulation on a {nx}x{ny} grid for {steps} steps...")
    t0 = time.time()
    
    for step in range(steps):
        # 2. Extract shifted arrays to construct the 5-point stencil neighbors.
        # To avoid custom shader modifications, we do fast slicing on CPU / metadata pointers, 
        # but execute all heavy bulk elementwise arithmetic on the GPU.
        
        # Internal domain is [1:-1, 1:-1]
        center = u[1:-1, 1:-1]
        
        # Shifted neighbors
        left  = u[1:-1, 0:-2]
        right = u[1:-1, 2:]
        up    = u[0:-2, 1:-1]
        down  = u[2:,   1:-1]
        
        # 3. Compute stencil arithmetic strictly using VGPUCache methods:
        # Step A: sum the 4 neighbors -> (left + right + up + down)
        neighbors_sum = c.add(left, right)
        neighbors_sum = c.add(neighbors_sum, up)
        neighbors_sum = c.add(neighbors_sum, down)
        
        # Step B: calculate 4 * center
        four = np.float32(4.0)
        four_center = c.multiply(four, center)
        
        # Step C: calculate Laplacian -> (neighbors_sum - 4 * center)
        laplacian = c.subtract(neighbors_sum, four_center)
        
        # Step D: Scale by step size -> gamma * Laplacian
        g = np.float32(gamma)
        flux = c.multiply(g, laplacian)
        
        # Step E: Update the active internal region -> u_next = center + flux
        u_next_internal = c.add(center, flux)
        
        # 4. Write internal updates back to the primary grid (preserving Dirichlet boundaries)
        u[1:-1, 1:-1] = u_next_internal
        
        if step % 50 == 0 or step == steps - 1:
            # Check maximum temperature remaining in the system via cached reduction
            max_temp = c.max(u)
            print(f"  Step {step} | Max Temp: {max_temp}°C")
            
    total_time = time.time() - t0
    print("-" * 50)
    print(f"Simulation completed in {total_time:.4f} seconds.")
    print(c.report())
    return u

if __name__ == "__main__":
    # Run simulation
    final_grid = solve_2d_heat(nx=256, ny=256, steps=100, gamma=0.2)