import numpy as np
import time
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from vgpu_cache import VGPUCache

def solve_2d_heat(nx=256, ny=256, steps=100, gamma=0.2):
    """
    Solves the 2D Heat Equation PDE using VGPUCache vector ops.
    """
    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
    
    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):
        # Internal domain
        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]
        
        # Compute stencil arithmetic via VGPUCache
        neighbors_sum = c.add(left, right)
        neighbors_sum = c.add(neighbors_sum, up)
        neighbors_sum = c.add(neighbors_sum, down)
        
        four_center = c.multiply(np.float32(4.0), center)
        laplacian = c.subtract(neighbors_sum, four_center)
        flux = c.multiply(np.float32(gamma), laplacian)
        
        # Update active internal region
        u[1:-1, 1:-1] = c.add(center, flux)
        
        # ... (Inside your simulation loop) ...
        if step % 20 == 0 or step == steps - 1:
            # FIX: Use .item() to safely get the Python scalar and avoid the DeprecationWarning
            max_temp = c.max(u).item()
            print(f"  Step {step:03d} | Max Temp: {max_temp:.2f}°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 the PDE solver
    final_grid = solve_2d_heat(nx=250, ny=250, steps=2500, gamma=0.2)
    
    # ---- FIX: Force non-interactive backend to prevent OpenGL/GDK context crash ----
    import matplotlib
    matplotlib.use('Agg') 
    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d import Axes3D
    # ---------------------------------------------------------------------------------

    print("\nGenerating 3D surface plot...")
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111, projection='3d')
    
    x = np.arange(0, final_grid.shape[0])
    y = np.arange(0, final_grid.shape[1])
    X, Y = np.meshgrid(x, y)
    
    surf = ax.plot_surface(X, Y, final_grid, cmap='inferno', edgecolor='none')
    
    ax.set_title("2D Heat Equation Simulation (VGPU Cache)", fontsize=14)
    ax.set_xlabel("X Grid")
    ax.set_ylabel("Y Grid")
    ax.set_zlabel("Temperature (°C)")
    ax.set_zlim(0, 100)
    fig.colorbar(surf, ax=ax, shrink=0.5, aspect=10, label="Temperature")
    
    # FIX: Save to disk instead of using plt.show() which clashes with ModernGL
    output_filename = "heat_plot.png"
    plt.savefig(output_filename, dpi=150)
    print(f"✅ Success! 3D Surface plot saved safely to '{output_filename}'")