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

# ---------------------------------------------------------
# 1. SETUP GRAPH-BASED MULTI-OBJECT STRESS SHADER
# ---------------------------------------------------------
COGWHEEL_STRESS_TEMPLATE = """#version 430 core
layout(local_size_x = 64) in;

struct Node {
    vec3 pos;
    float stress;
    int gear_id; // 0 for Drive Gear, 1 for Driven Gear
};

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

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 propagation_decay; // Material stress dissipation factor
uniform float contact_pressure;   // Force passing between colliding teeth

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

    Node current = nodes_in[i];
    float max_neighbor_stress = current.stress;
    
    // --- PART A: INTRA-GEAR INTERNAL ADJACENCY PROPAGATION ---
    for (int n = 0; n < stencil[i].count; n++) {
        int neighbor_idx = stencil[i].neighbors[n];
        if (neighbor_idx != -1) {
            float nbr_stress = nodes_in[neighbor_idx].stress;
            // Stress concentrates and moves along the rigid structure
            max_neighbor_stress = max(max_neighbor_stress, nbr_stress * propagation_decay);
        }
    }

    // --- PART B: INTER-GEAR INTERSECTING CONTACT FORCE ---
    // If this node is near the other gear's teeth, transmit the shear pressure
    float contact_stress = 0.0;
    if (current.stress > 10.0) { // If stressed, transmit to overlapping gear
        for (uint j = 0; j < uint(num_nodes); j++) {
            if (nodes_in[j].gear_id != current.gear_id) {
                float dist = distance(current.pos, nodes_in[j].pos);
                // Proximity tolerance defining where the teeth mesh together
                if (dist < 0.18) { 
                    contact_stress = current.stress * contact_pressure;
                }
            }
        }
    }

    // Combine internal mechanical shear mapping
    float next_stress = max(max_neighbor_stress, contact_stress);

    // Keep the source boundary load applied at the drive gear core axle
    if (current.gear_id == 0 && length(current.pos.xy - vec2(-0.9, 0.0)) < 0.15) {
        next_stress = 100.0; // Fixed input motor torque load (MPa)
    }

    // Export to output layout buffer
    nodes_out[i].pos = current.pos;
    nodes_out[i].stress = next_stress;
    nodes_out[i].gear_id = current.gear_id;
}
"""

# ---------------------------------------------------------
# 2. GENERATE PROCEDURAL COGWHEEL 3D OBJ DATA
# ---------------------------------------------------------
def make_cogwheel(center_x, center_y, gear_id, num_teeth=12, rings=5, points_per_ring=36):
    """Generates a 3D structural gear manifold with defined profiles."""
    nodes = []
    base_idx_offset = 0
    
    for r in range(rings):
        # Scale radius outwards to forge the gear center vs teeth flanks
        radius_ratio = 0.2 + (r / (rings - 1)) * 0.5
        z = 0.0 # Thin 2D extrusion plane mapping structural plate stress
        
        for p in range(points_per_ring):
            angle = (p / points_per_ring) * 2.0 * np.pi
            
            # Procedural gear teeth profile modulation (sine wave teeth cuts)
            tooth_amplitude = 0.12 if r == (rings - 1) else 0.0
            r_final = radius_ratio + tooth_amplitude * np.sin(num_teeth * angle)
            
            x = center_x + r_final * np.cos(angle)
            y = center_y + r_final * np.sin(angle)
            
            # Initial baseline stress state
            initial_stress = 0.0
            nodes.append(((x, y, z), initial_stress, gear_id))
            
    # Compile graph connection stencils mapping mechanical links
    adjacency = []
    for r in range(rings):
        for p in range(points_per_ring):
            curr = r * points_per_ring + p
            neighbors = []
            
            # Topology loops (Structural ring neighbors)
            left = r * points_per_ring + ((p - 1) % points_per_ring)
            right = r * points_per_ring + ((p + 1) % points_per_ring)
            neighbors.extend([left, right])
            
            # Radial cross-beams 
            if r > 0: neighbors.append((r - 1) * points_per_ring + p)
            if r < rings - 1: neighbors.append((r + 1) * points_per_ring + p)
                
            padded = neighbors + [-1] * (6 - len(neighbors))
            adjacency.append((len(neighbors), padded))
            
    return nodes, adjacency

def generate_meshed_gears_system():
    # Drive Gear centered at left, Driven Gear offset to mesh on right
    nodes_drive, adj_drive = make_cogwheel(center_x=-0.9, center_y=0.0, gear_id=0)
    nodes_driven, adj_driven = make_cogwheel(center_x=0.25, center_y=0.0, gear_id=1)
    
    # Concatenate structures into monolithic GPU payload blocks
    offset = len(nodes_drive)
    adjusted_adj_driven = []
    for count, nbrs in adj_driven:
        new_nbrs = [n + offset if n != -1 else -1 for n in nbrs]
        adjusted_adj_driven.append((count, new_nbrs))
        
    all_nodes = nodes_drive + nodes_driven
    all_adj = adj_drive + adjusted_adj_driven
    
    node_dtype = np.dtype([('pos', 'f4', 3), ('stress', 'f4'), ('gear_id', 'i4')])
    adj_dtype = np.dtype([('count', 'i4'), ('neighbors', 'i4', 6)])
    
    return np.array(all_nodes, dtype=node_dtype), np.array(all_adj, dtype=adj_dtype), len(all_nodes)

# ---------------------------------------------------------
# 3. RUNTIME SOLVER EXECUTION
# ---------------------------------------------------------
def run_gear_stress_simulation():
    ctx = moderngl.create_standalone_context()
    prog = ctx.compute_shader(COGWHEEL_STRESS_TEMPLATE)
    
    node_arr, adj_arr, N = generate_meshed_gears_system()
    
    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['propagation_decay'] = 0.96 # Yield resistance dampening
    prog['contact_pressure'] = 0.92  # Force transfer ratio across tooth intersection
    
    print(f"Propagating stress field vectors across {N} intersecting mesh nodes...")
    t0 = time.time()
    
    # Step simulation iterations to let stress load flow from axle to teeth contact interface
    for step in range(80):
        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_states = np.frombuffer(buf_in.read(), dtype=node_arr.dtype)
    print(f"Stress analysis completed in {time.time() - t0:.4f} seconds.")
    
    # ---- PLOT RECONSTRUCTED COGWHEEL FIELD ----
    fig, ax = plt.subplots(figsize=(10, 8))
    fig.patch.set_facecolor('#111111')
    ax.set_facecolor('#111111')
    
    X = final_states['pos'][:, 0]
    Y = final_states['pos'][:, 1]
    S = final_states['stress']
    
    # Scatter stress heat map coordinates
    scat = ax.scatter(X, Y, c=S, cmap='jet', s=35, edgecolors='none', vmin=0, vmax=100)
    
    cbar = fig.colorbar(scat, ax=ax, shrink=0.7)
    cbar.set_label('Von Mises Effective Stress (MPa)', color='white')
    cbar.ax.yaxis.set_tick_params(color='white', labelcolor='white')
    #cbar.ax.yaxis.set_tick_params(color='white')
    #plt.setp(plt.getp(cbar.ax.flatten(), 'yticklabels'), color='white')
    
    ax.set_title("Intersecting Cogwheels: Mechanical Yield Load Distribution", color='white', fontsize=13)
    ax.axis('off')
    ax.set_aspect('equal')
    
    output_img = "cogwheel_stress_map.png"
    plt.savefig(output_img, dpi=150, facecolor=fig.get_facecolor(), edgecolor='none')
    print(f"✅ Stress distribution plot successfully compiled and saved to '{output_img}'")

if __name__ == "__main__":
    run_gear_stress_simulation()