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:
        # Try using professional pymeshlab first
        import pymeshlab
        ms = pymeshlab.MeshSet()
        ms.load_new_mesh(input_path)
        
        target_faces = int(ms.current_mesh().face_number() * keep_ratio)
        
        # Optimize using quadric edge collapse to preserve 95% visual effect
        try:
            ms.apply_filter('meshing_decimation_quadric_edge_collapse', 
                            targetfacenum=target_faces, 
                            preserveboundary=True, 
                            preservenormal=True, 
                            preservetopology=True)
        except AttributeError:
            ms.meshing_decimation_quadric_edge_collapse(
                targetfacenum=target_faces,
                preserveboundary=True,
                preservenormal=True,
                preservetopology=True
            )
            
        ms.save_current_mesh(output_path)
        return True
    except ImportError:
        pass
    except Exception as e:
        print(f"Pymeshlab simplification failed: {e}. Falling back to Open3D...")
        
    try:
        # Try using open3d as second fallback
        import open3d as o3d
        o3d_mesh = o3d.io.read_triangle_mesh(input_path)
        if len(o3d_mesh.triangles) > 0:
            target_faces = max(1, int(len(o3d_mesh.triangles) * keep_ratio))
            smp_mesh = o3d_mesh.simplify_quadric_decimation(target_number_of_triangles=target_faces)
            smp_mesh.compute_triangle_normals()
            o3d.io.write_triangle_mesh(output_path, smp_mesh)
            return True
    except ImportError:
        pass
    except Exception as e:
        print(f"Open3D simplification failed: {e}. Falling back to PyFQMR...")
        
    try:
        # Try using pyfqmr as third fallback
        import pyfqmr
        import trimesh
        mesh_data = trimesh.load(input_path, file_type='stl')
        target_faces = max(1, int(len(mesh_data.faces) * keep_ratio))
        simplifier = pyfqmr.Simplify()
        simplifier.setMesh(mesh_data.vertices, mesh_data.faces)
        simplifier.simplify_mesh(target_count=target_faces, aggressiveness=7, preserve_border=True, verbose=False)
        
        mesh_parts = simplifier.getMesh()
        smp_mesh = trimesh.Trimesh(vertices=mesh_parts[0], faces=mesh_parts[1], process=False)
        smp_mesh.export(output_path, file_type='stl')
        return True
    except ImportError:
        pass
    except Exception as e:
        print(f"PyFQMR simplification failed: {e}. Falling back to custom algorithm...")

    try:
        try:
            import trimesh
            mesh_data = trimesh.load(input_path, file_type='stl')
            if hasattr(mesh_data, 'triangles'):
                vertices = mesh_data.triangles.reshape(-1, 3)
            else:
                vertices = mesh_data.vertices[mesh_data.faces].reshape(-1, 3)
            use_trimesh = True
        except ImportError:
            # Load mesh using numpy-stl fallback
            m = mesh.Mesh.from_file(input_path)
            vertices = m.vectors.reshape(-1, 3)
            use_trimesh = False
            
        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:
            if use_trimesh:
                mesh_data.export(output_path, file_type='stl')
            else:
                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.
        # Function to simplify given a grid size
        def do_simplify(g_size):
            v_idx = np.round((vertices - min_v) / g_size).astype(np.int64)
            
            # Fast 1D hash to avoid extremely slow np.unique(axis=0) on 2D arrays
            max_idx = v_idx.max(axis=0) + 1
            v_1d = v_idx[:, 0] + v_idx[:, 1] * max_idx[0] + v_idx[:, 2] * max_idx[0] * max_idx[1]
            
            # Find unique grid cells and map old vertices to them
            _, unique_idx, inv_idx = np.unique(v_1d, 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
            
        target_faces = max(1, int((len(vertices) // 3) * keep_ratio))
        low_g = max_dim * 0.0005
        high_g = max_dim * 0.2
        best_verts = vertices
        best_faces = np.arange(len(vertices)).reshape(-1, 3)
        
        # Binary search for the right grid size
        for _ in range(8):
            g_size = (low_g + high_g) / 2
            v, f = do_simplify(g_size)
            best_verts, best_faces = v, f
            
            if len(f) > target_faces:
                # too many faces, make grid coarser (larger)
                low_g = g_size
            else:
                # too few faces, make grid finer (smaller)
                high_g = g_size
                
            if abs(len(f) - target_faces) < target_faces * 0.05:
                break
                
        new_vertices, valid_faces = best_verts, best_faces
        
        if use_trimesh:
            simplified = trimesh.Trimesh(vertices=new_vertices, faces=valid_faces, process=False)
            simplified.export(output_path, file_type='stl')
            return True

        # Build the simplified mesh using fallback
        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])