import gradio as gr
import trimesh
import numpy as np
import tempfile
import os
import traceback

def fit_mesh(mesh_file):
    """
    SuperFit Algorithm:
    Replaces each isolated part of the input mesh with a best-fit geometric primitive 
    (an oriented bounding box).
    """
    if mesh_file is None:
        return None
        
    try:
        # Load the mesh/scene
        scene = trimesh.load(mesh_file.name, force='scene')
        primitives = []
        
        def process_geometry(geom):
            # Generate the Oriented Bounding Box (OBB) for the geometry
            obb = geom.bounding_box_oriented
            
            # Transfer the color if available
            try:
                color = None
                if hasattr(geom.visual, 'face_colors') and len(geom.visual.face_colors) > 0:
                    color = geom.visual.face_colors[0]
                elif hasattr(geom.visual, 'main_color'):
                    color = geom.visual.main_color
                    
                if color is not None:
                    obb.visual.face_colors = color
            except Exception:
                pass
                
            return obb

        # Process each geometry in the scene
        if isinstance(scene, trimesh.Scene):
            for name, geom in scene.geometry.items():
                if isinstance(geom, trimesh.Trimesh):
                    # Sometimes a single geometry node contains multiple disconnected parts
                    components = geom.split(only_watertight=False)
                    for comp in components:
                        if len(comp.vertices) > 4: # Ignore tiny dust/noise
                            primitives.append(process_geometry(comp))
        elif isinstance(scene, trimesh.Trimesh):
            components = scene.split(only_watertight=False)
            for comp in components:
                if len(comp.vertices) > 4:
                    primitives.append(process_geometry(comp))
                
        # If we failed to extract any primitives, return the original
        if not primitives:
            return mesh_file.name
            
        # Create a new scene with the bounding box primitives
        out_scene = trimesh.Scene(primitives)
        
        # Export to a temporary GLB file
        out_path = tempfile.mktemp(suffix=".glb")
        out_scene.export(out_path)
        
        return out_path
        
    except Exception as e:
        print("Error in SuperFit algorithm:")
        traceback.print_exc()
        # Fallback to the original mesh if the algorithm fails
        return mesh_file.name

def segment_mesh(mesh_file, num_clusters=5):
    """
    Lightweight geometric segmentation algorithm to replace Stirnix/PartField.
    Uses K-Means clustering on face centroids to split the mesh into distinct parts.
    """
    if mesh_file is None:
        return None
        
    try:
        from sklearn.cluster import KMeans
        
        # Force load as a single mesh
        mesh = trimesh.load(mesh_file.name, force='mesh')
        
        # Calculate face centroids
        centroids = mesh.triangles.mean(axis=1)
        
        # Cluster faces
        num_clusters = int(num_clusters)
        clustering = KMeans(n_clusters=num_clusters, random_state=42, n_init=10).fit(centroids)
        labels = clustering.labels_
        
        # Split mesh into submeshes
        submeshes = []
        for i in range(num_clusters):
            face_mask = (labels == i)
            sub_faces = mesh.faces[face_mask]
            if len(sub_faces) > 0:
                submesh = trimesh.Trimesh(vertices=mesh.vertices, faces=sub_faces, process=True)
                submeshes.append(submesh)
                
        out_scene = trimesh.Scene(submeshes)
        out_path = tempfile.mktemp(suffix=".glb")
        out_scene.export(out_path)
        return out_path
        
    except Exception as e:
        print("Error in segmentation algorithm:")
        traceback.print_exc()
        return mesh_file.name

with gr.Blocks() as demo:
    gr.Markdown("# SuperFit API (All-in-One)")
    gr.Markdown("Replaces complex meshes with fitted primitive shapes AND provides geometric segmentation.")
    
    with gr.Tab("SuperFit (Fit Primitives)"):
        with gr.Row():
            mesh_input = gr.File(label="Input Mesh (GLB)")
            mesh_output = gr.File(label="Fitted Primitives (GLB)")
        btn = gr.Button("Fit Primitives")
        btn.click(fn=fit_mesh, inputs=[mesh_input], outputs=[mesh_output], api_name="fit_mesh")
        
    with gr.Tab("Segmentation (Stirnix Replacement)"):
        with gr.Row():
            seg_input = gr.File(label="Input Mesh (GLB)")
            num_clusters = gr.Slider(minimum=1, maximum=20, value=5, step=1, label="Number of Parts")
        seg_output = gr.File(label="Segmented Mesh (GLB)")
        seg_btn = gr.Button("Segment Mesh")
        seg_btn.click(fn=segment_mesh, inputs=[seg_input, num_clusters], outputs=[seg_output], api_name="segment")

if __name__ == "__main__":
    demo.launch()