import numpy as np
import time
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import moderngl

# ---------------------------------------------------------
# 1. SETUP RAW GLSL COMPUTE SHADER FOR 3D ADJACENCY MESHES
# ---------------------------------------------------------
MESH_HEAT_TEMPLATE = """#version 430 core
layout(local_size_x = 64) in;

struct Node {
    vec3 pos;
    float temp;
};

// Neighbor links mapped out from the OBJ edge connections
struct Adjacency {
    int count;
    int neighbors[6]; // Max 6 neighboring edge connections per vertex node
};

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_alpha;

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;

    // Loop over the node connections defined in the OBJ topology
    for (int n = 0; n < stencil[i].count; n++) {
        int neighbor_idx = stencil[i].neighbors[n];
        Node nbr = nodes_in[neighbor_idx];

        // Geometric Laplacian using 3D physical distance weights
        float dist = distance(current.pos, nbr.pos);
        if (dist > 0.0001) {
            float weight = 1.0 / dist;
            laplacian += (nbr.temp - current.temp) * weight;
            total_weight += weight;
        }
    }

    // Apply normalized thermal flux step over 3D manifold surface
    float next_temp = current.temp;
    if (total_weight > 0.0) {
        next_temp += dt_alpha * (laplacian / total_weight);
    }

    // Write state back out to VRAM
    nodes_out[i].pos = current.pos;
    nodes_out[i].temp = next_temp;
}
"""

# ---------------------------------------------------------
# 2. HELPER TO CREATE ARBITRARY 3D GEOMETRY DATA
# ---------------------------------------------------------
def generate_3d_tube_mesh(segments=40, rings=40):
    """Generates structural node arrays mirroring standard OBJ data maps."""
    nodes = []
    # Build 3D Cylindrical Tube
    for r in range(rings):
        z = (r / (rings - 1)) * 4.0 - 2.0
        for s in range(segments):
            theta = (s / segments) * 2.0 * np.pi
            x = np.cos(theta)
            y = np.sin(theta)
            # Create a localized high thermal hot spot at the center base
            temp = 100.0 if (abs(z) < 0.3 and x > 0.7) else 0.0
            nodes.append(((x, y, z), temp))
            
    num_nodes = len(nodes)
    node_data = np.array(nodes, dtype=[('pos', 'f4', 3), ('temp', 'f4')])
    
    # Pre-compute structural vertex adjacency loops matching OBJ faces
    adjacency_list = []
    for r in range(rings):
        for s in range(segments):
            idx = r * segments + s
            neighbors = []
            
            # Find natural structural neighbors (Left, Right, Up, Down ring paths)
            left = r * segments + ((s - 1) % segments)
            right = r * segments + ((s + 1) % segments)
            neighbors.extend([left, right])
            
            if r > 0: neighbors.append((r - 1) * segments + s)
            if r < rings - 1: neighbors.append((r + 1) * segments + s)
                
            # Pad adjacency structures to match static GLSL struct sizes
            count = len(neighbors)
            padded = neighbors + [-1] * (6 - count)
            adjacency_list.append((count, padded))
            
    adj_dtype = np.dtype([('count', 'i4'), ('neighbors', 'i4', 6)])
    adj_data = np.array(adjacency_list, dtype=adj_dtype)
    
    return node_data, adj_data, num_nodes

# ---------------------------------------------------------
# 3. RUNTIME SIMULATION DISPATCH
# ---------------------------------------------------------
def run_3d_pde_mesh():
    ctx = moderngl.create_standalone_context()
    
    # Compile 3D Adjacency Shader
    prog = ctx.compute_shader(MESH_HEAT_TEMPLATE)
    
    # Generate 3D Data
    node_arr, adj_arr, N = generate_3d_tube_mesh()
    
    # Build GPU VRAM Buffers
    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_alpha'] = 0.15 # Time step thermal conductivity factor
    
    print(f"Solving 3D Heat Flow across {N} nodes using structural adjacency buffers...")
    
    # Time-stepping simulation execution loops completely on GPU
    for step in range(120):
        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)
        
        # Double buffer ping-pong: copy next frame out to current frame
        ctx.copy_buffer(buf_in, buf_out)
        
    # Read finished frame back to host CPU memory
    final_nodes = np.frombuffer(buf_in.read(), dtype=node_arr.dtype)
    
    # ---- 3D PLOT RESULTS ----
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111, projection='3d')
    
    # Unpack geometry vectors
    X = final_nodes['pos'][:, 0]
    Y = final_nodes['pos'][:, 1]
    Z = final_nodes['pos'][:, 2]
    T = final_nodes['temp']
    
    # Render scattered nodes colored by their thermal field state
    scat = ax.scatter(X, Y, Z, c=T, cmap='inferno', s=15, edgecolors='none')
    fig.colorbar(scat, ax=ax, label='Temperature (°C)')
    ax.set_title("3D Complex Surface Heat Dissipation")
    ax.axis('off')
    
    plt.savefig("3d_mesh_heat.png", dpi=150)
    print("✅ Success! 3D Geometry thermal solution saved to '3d_mesh_heat.png'")

if __name__ == "__main__":
    run_3d_pde_mesh()