import numpy as np
import time
import matplotlib
matplotlib.use('Agg')  # Safe headless mode for Linux/ModernGL GUI interop
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import moderngl

# ---------------------------------------------------------
# 1. ADVANCED HEAT CONDUCTION + CONVECTION COMPUTE SHADER
# ---------------------------------------------------------
OBJ_STOVE_SHADER = """#version 430 core
layout(local_size_x = 64) in;

struct Node {
    vec3 pos;
    float val;
};

struct Adjacency {
    int count;
    int neighbors[8];
};

layout(std430, binding = 0) readonly buffer InputMesh  { Node nodes_in[]; };
layout(std430, binding = 1) writeonly buffer OutputMesh { Node nodes_out[]; };
layout(std430, binding = 2) readonly buffer GraphLayout { Adjacency stencil[]; };

uniform int num_nodes;
uniform float dt_coeff;
uniform float z_bottom; 
uniform float cooling_rate; // Simulates heat escaping into ambient air

void main() {
    uint i = gl_GlobalInvocationID.x;
    if (i >= uint(num_nodes)) return;

    Node current = nodes_in[i];
    float laplacian = 0.0;
    float total_weight = 0.0;

    // Heat propagation over mesh edges
    for (int n = 0; n < stencil[i].count; n++) {
        int neighbor_idx = stencil[i].neighbors[n];
        if (neighbor_idx != -1) {
            Node nbr = nodes_in[neighbor_idx];
            float dist = distance(current.pos, nbr.pos);
            
            if (dist > 0.0001) {
                float weight = 1.0 / dist;
                laplacian += (nbr.val - current.val) * weight;
                total_weight += weight;
            }
        }
    }

    float next_val = current.val;
    if (total_weight > 0.0) {
        next_val += dt_coeff * (laplacian / total_weight);
    }

    // Environmental Air Dissipation (Newton's law of cooling)
    // Pulls temperature back down toward ambient room temp (0°C base)
    next_val -= cooling_rate * next_val;

    // STRICT BASE LAYER THRESHOLD:
    // Tightened boundary tolerance to isolate ONLY the absolute flat underside faces
    if (current.pos.z <= z_bottom + 0.005) {
        next_val = 100.0; 
    }

    // Keep temperatures within physically safe margins
    nodes_out[i].pos = current.pos;
    nodes_out[i].val = clamp(next_val, 0.0, 100.0);
}
"""

# ---------------------------------------------------------
# 2. WAVEFRONT .OBJ PARSER & GRAPH LAYOUT BUILDER
# ---------------------------------------------------------
def load_blender_obj(filepath):
    """Parses .obj vertex strings ('v') and polygon links ('f') into matrices."""
    vertices = []
    faces = []
    
    print(f"Reading topology from '{filepath}'...")
    with open(filepath, 'r') as f:
        for line in f:
            if line.startswith('v '):
                parts = line.split()
                vertices.append([float(parts[1]), float(parts[2]), float(parts[3])])
            elif line.startswith('f '):
                parts = line.split()
                # Parse indices supporting variant styles (e.g., v/vt/vn)
                face = [int(p.split('/')[0]) - 1 for p in parts[1:]]
                faces.append(face)
                
    verts = np.array(vertices, dtype=np.float32)
    num_nodes = len(verts)
    
    # Generate an adjacency edge lookup map from triangles
    adj_dict = {i: set() for i in range(num_nodes)}
    for face in faces:
        for i in range(len(face)):
            v1 = face[i]
            v2 = face[(i + 1) % len(face)]
            adj_dict[v1].add(v2)
            adj_dict[v2].add(v1)
            
    # Format edge lists to fit static buffer structures (Max 8 links per vertex)
    adjacency_list = []
    for i in range(num_nodes):
        neighbors = list(adj_dict[i])[:8]
        count = len(neighbors)
        padded = neighbors + [-1] * (8 - count)
        adjacency_list.append((count, padded))
        
    # Build structured binary blocks for VRAM
    node_data = np.zeros(num_nodes, dtype=[('pos', 'f4', 3), ('val', 'f4')])
    node_data['pos'] = verts
    node_data['val'] = 0.0  # Cold start condition
    
    adj_data = np.array(adjacency_list, dtype=[('count', 'i4'), ('neighbors', 'i4', 8)])
    
    return node_data, adj_data, faces

# ---------------------------------------------------------
# 3. CORE PROCESSING PIPELINE
# ---------------------------------------------------------
def run_stove_simulation(obj_filename, steps=140):
    ctx = moderngl.create_standalone_context()
    prog = ctx.compute_shader(OBJ_STOVE_SHADER)
    
    node_arr, adj_arr, faces = load_blender_obj(obj_filename)
    N = len(node_arr)
    
    # AUTOMATIC BOUNDING BOX MINIMUM EXTRACTOR
    z_min = float(np.min(node_arr['pos'][:, 2]))
    print(f"-> Base boundary detected at Z = {z_min:.4f}. Stove element aligned here.")
    
    # Upload parameters and arrays to GPU
    buf_in = ctx.buffer(node_arr.tobytes())
    buf_out = ctx.buffer(reserve=node_arr.nbytes)
    buf_adj = ctx.buffer(adj_arr.tobytes())
    
    prog['num_nodes'] = N
    prog['dt_coeff'] = 0.04       # Stable thermal conduction rate step
    prog['cooling_rate'] = 0.008   # Convective cooling coefficient to break saturation
    prog['z_bottom'] = z_min      # Tight baseline limit passed to GPU
    
    print(f"Simulating thermal conduction loop for {steps} iterations entirely on GPU...")
    t0 = time.time()
    
    for step in range(steps):
        buf_in.bind_to_storage_buffer(0)
        buf_out.bind_to_storage_buffer(1)
        buf_adj.bind_to_storage_buffer(2)
        
        prog.run(max(1, (N + 63) // 64), 1, 1)
        ctx.copy_buffer(buf_in, buf_out)
        
    final_mesh = np.frombuffer(buf_in.read(), dtype=node_arr.dtype)
    print(f"Simulation completed in {time.time() - t0:.4f}s.")
    
    # ---- PLOT 3D SURFACE GRADIENT ----
    print("Generating 3D surface mesh plot...")
    fig = plt.figure(figsize=(11, 9))
    ax = fig.add_subplot(111, projection='3d')
    fig.patch.set_facecolor('#111111')
    ax.set_facecolor('#111111')
    
    X = final_mesh['pos'][:, 0]
    Y = final_mesh['pos'][:, 1]
    Z = final_mesh['pos'][:, 2]
    V = final_mesh['val']
    
    # plot_trisurf uses parsed faces to draw solid surfaces instead of loose dots
    surf = ax.plot_trisurf(X, Y, Z, triangles=faces, cmap='inferno', linewidth=0.2, edgecolors='#222222')
    surf.set_array(V)
    surf.set_clim(0, 100)
    
    cbar = fig.colorbar(surf, ax=ax, shrink=0.5, aspect=15)
    cbar.ax.yaxis.set_tick_params(color='white', labelcolor='white')
    cbar.set_label('Temperature Field (°C)', color='white')
    
    ax.set_title(f"Stove Convection Model: {obj_filename} ({steps} steps)", color='white', fontsize=13)
    ax.axis('off')
    
    output_img = "stove_simulation_render.png"
    plt.savefig(output_img, dpi=150, facecolor=fig.get_facecolor(), edgecolor='none')
    print(f"✅ Render complete! Open '{output_img}' to view your heat gradient profile.")

if __name__ == "__main__":
    # Runs the calculation over your exported Blender cube
    run_stove_simulation("untitled.obj", steps=140)