import numpy as np
import matplotlib
matplotlib.use('Agg') # Safe headless mode for Linux/ModernGL interop
import matplotlib.pyplot as plt
from vgpu_cache import VGPUCache

def compute_gpu_von_mises(sxx, syy, sxy):
    """
    Computes von Mises effective stress for Plane Stress using VGPUCache.
    Accepts 2D NumPy arrays or PyTorch Tensors.
    """
    c = VGPUCache()
    
    # 1. Compute sxx^2 and syy^2
    sxx_sq = c.multiply(sxx, sxx)
    syy_sq = c.multiply(syy, syy)
    
    # 2. Compute sxx * syy
    sxx_syy = c.multiply(sxx, syy)
    
    # 3. Compute 3 * sxy^2
    sxy_sq = c.multiply(sxy, sxy)
    three = np.float32(3.0)
    three_sxy_sq = c.multiply(three, sxy_sq)
    
    # 4. Combine the elements under the radical: (sxx_sq + syy_sq - sxx_syy + three_sxy_sq)
    total = c.add(sxx_sq, syy_sq)
    total = c.subtract(total, sxx_syy)
    total = c.add(total, three_sxy_sq)
    
    # 5. Take the elementwise square root
    von_mises = c.sqrt(total)
    
    return von_mises

if __name__ == "__main__":
    # Simulate a 200x200 stress grid field for a mechanical plate under bending/shear load
    # Coordinates for mesh
    x = np.linspace(-1, 1, 200)
    y = np.linspace(-1, 1, 200)
    X, Y = np.meshgrid(x, y)
    
    # Mock normal and shear stress fields (in Megapascals / MPa)
    sxx = (X**2 * 150.0).astype(np.float32)        # Tensile bending stress
    syy = (Y**2 * 50.0).astype(np.float32)         # Transverse stress
    sxy = ((X * Y) * 80.0).astype(np.float32)      # Torsional shear stress
    
    print("Dispatched to GPU Cache...")
    vm_stress = compute_gpu_von_mises(sxx, syy, sxy)
    
    # Use your patched reduction 'max' to find the peak critical stress area
    max_stress = vm_stress.max().item()
    print(f"Max von Mises effective stress: {max_stress:.2f} MPa")
    
    # ---- Render & Save Contour Plot ----
    plt.figure(figsize=(8, 6))
    contour = plt.contourf(X, Y, vm_stress, levels=30, cmap='jet')
    plt.colorbar(contour, label='Effective Stress $\sigma_v$ (MPa)')
    plt.title("Von Mises Stress Distribution (Plane Stress Model)")
    plt.xlabel("X coordinate")
    plt.ylabel("Y coordinate")
    
    output_img = "von_mises_plot.png"
    plt.savefig(output_img, dpi=150)
    print(f"✅ Stress profile successfully generated and saved to '{output_img}'")