import numpy as np
import time
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import moderngl

# ---------------------------------------------------------
# 1. THE GEOMETRIC FINITE-VOLUME LAPLACIAN SHADER
# ---------------------------------------------------------
OBJ_MESH_SHADER = """#version 430 core
layout(local_size_x = 64) in;

struct Node {
    vec3 pos;
    float val;
};

struct Adjacency {
    int count;
    int neighbors[8]; // Supports up to 8 connected edges per vertex
};

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;

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;

    // Propagate over true adjacency links parsed from Blender faces
    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);
    }

    // Maintain a constant boundary source (e.g., heat/load applied to the bottom of the object)
    if (current.pos.z < -0.9) {
        next_val = 100.0; 
    }

    nodes_out[i].pos = current.pos;
    nodes_out[i].val = next_val;
}
"""

# ---------------------------------------------------------
# 2. ROBUST BLENDER .OBJ PARSER & GRAPH BUILDER
# ---------------------------------------------------------
def load_blender_obj(filepath):
    """Parses a standard wavefront .obj file into structured 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()
                # Extract indices handling face structural variants like 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)
    
    # Construct structural adjacency map from face loops
    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)
            
    # Pack dictionary into padded GPU layout blocks (Max 8 connections)
    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 records
    node_data = np.zeros(num_nodes, dtype=[('pos', 'f4', 3), ('val', 'f4')])
    node_data['pos'] = verts
    # Apply initial conditions (all cold except the heat source boundary mapped in the shader)
    node_data['val'] = 0.0 
    
    adj_data = np.array(adjacency_list, dtype=[('count', 'i4'), ('neighbors', 'i4', 8)])
    
    return node_data, adj_data, faces

# ---------------------------------------------------------
# 3. RUNTIME DISPATCH & VISUALIZATION
# ---------------------------------------------------------
def run_blender_pde(obj_filename):
    ctx = moderngl.create_standalone_context()
    prog = ctx.compute_shader(OBJ_MESH_SHADER)
    
    # Load your exported model
    try:
        node_arr, adj_arr, faces = load_blender_obj(obj_filename)
    except FileNotFoundError:
        print(f"Creating a placeholder dummy '{obj_filename}' for testing...")
        create_mock_obj(obj_filename)
        node_arr, adj_arr, faces = load_blender_obj(obj_filename)

    N = len(node_arr)
    
    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.1  # Stabilization stepping multiplier
    
    print(f"Processing PDE over {N} vertices entirely on the GPU...")
    t0 = time.time()
    
    for step in range(200):
        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"Done in {time.time() - t0:.4f}s. Generating realistic 3D surface plot...")
    
    # ---- 3D TRISURF PLOT (Displays continuous filled polygons, not dots!) ----
    fig = plt.figure(figsize=(10, 8))
    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 the parsed face geometry to draw solid filled faces
    surf = ax.plot_trisurf(X, Y, Z, triangles=faces, cmap='inferno', linewidth=0.1, edgecolors='#222222')
    
    # Map array values onto surface face color rendering context
    surf.set_array(V)
    surf.set_clim(0, 100)
    
    cbar = fig.colorbar(surf, ax=ax, shrink=0.6, aspect=12)
    cbar.ax.yaxis.set_tick_params(color='white', labelcolor='white')
    cbar.set_label('Scalar Distribution Field', color='white')
    
    ax.set_title(f"Blender Model PDE Simulation: {obj_filename}", color='white', fontsize=12)
    ax.axis('off')
    
    output_img = "blender_pde_render.png"
    plt.savefig(output_img, dpi=150, facecolor=fig.get_facecolor(), edgecolor='none')
    print(f"✅ Render complete! Open '{output_img}' to view your results.")

def create_mock_obj(filename):
    """Generates a basic fallback cylinder mesh file if you don't have one handy."""
    with open(filename, 'w') as f:
        f.write("# Fallback Mesh\n")
        idx = 1
        for z in np.linspace(-1, 1, 15):
            for theta in np.linspace(0, 2*np.pi, 16, endpoint=False):
                f.write(f"v {np.cos(theta):.4f} {np.sin(theta):.4f} {z:.4f}\n")
        for r in range(14):
            for i in range(16):
                p1 = r * 16 + i + 1
                p2 = r * 16 + ((i + 1) % 16) + 1
                p3 = (r + 1) * 16 + i + 1
                p4 = (r + 1) * 16 + ((i + 1) % 16) + 1
                f.write(f"f {p1} {2} {p3}\n")
                f.write(f"f {p2} {p4} {p3}\n")

if __name__ == "__main__":
    # Point this to any .obj file you export from Blender!
    run_blender_pde("untitled.obj")