import numpy as np
from stl import mesh
import struct
import sys
import os

def simplify_stl(input_path, output_path, keep_ratio=0.1):
    try:
        # Load mesh using numpy-stl
        m = mesh.Mesh.from_file(input_path)
        vertices = m.vectors.reshape(-1, 3)
        
        min_v = vertices.min(axis=0)
        max_v = vertices.max(axis=0)
        bbox_size = max_v - min_v
        max_dim = np.max(bbox_size)
        
        if max_dim == 0:
            m.save(output_path)
            return True
            
        # Target roughly a resolution that gives us keep_ratio faces.
        # This is a heuristic approach to grid-based vertex clustering.
        # We start with a baseline resolution (e.g. 2% of max dimension) 
        # and adjust if needed.
        grid_resolution = max_dim * 0.02
        
        # Function to simplify given a grid size
        def do_simplify(g_size):
            v_idx = np.round((vertices - min_v) / g_size).astype(np.int32)
            # Find unique grid cells and map old vertices to them
            _, unique_idx, inv_idx = np.unique(v_idx, axis=0, return_index=True, return_inverse=True)
            new_verts = vertices[unique_idx]
            
            # Map faces to new vertices
            faces = inv_idx.reshape(-1, 3)
            
            # Remove degenerate faces (faces where at least two vertices resolve to the same cell)
            valid = (faces[:,0] != faces[:,1]) & (faces[:,1] != faces[:,2]) & (faces[:,0] != faces[:,2])
            valid_faces = faces[valid]
            
            return new_verts, valid_faces
            
        new_vertices, valid_faces = do_simplify(grid_resolution)
        
        # Build the simplified mesh
        new_m = mesh.Mesh(np.zeros(valid_faces.shape[0], dtype=mesh.Mesh.dtype))
        
        # Vectorized assignment
        new_m.vectors[:, 0, :] = new_vertices[valid_faces[:, 0]]
        new_m.vectors[:, 1, :] = new_vertices[valid_faces[:, 1]]
        new_m.vectors[:, 2, :] = new_vertices[valid_faces[:, 2]]
        
        # Calculate normals correctly
        new_m.update_normals()
        new_m.save(output_path)
        return True
        
    except Exception as e:
        print(f"Error simplifying STL: {e}")
        return False

if __name__ == "__main__":
    if len(sys.argv) > 2:
        simplify_stl(sys.argv[1], sys.argv[2])