import gradio as gr
import trimesh
import numpy as np
from PIL import Image
import tempfile

def apply_rust_texture(section):
    # Load the original mesh
    mesh = trimesh.load('train.glb', force='mesh')
    rust_texture = Image.open('rust_steel.png').convert('RGB')

    # Train body selection logic
    bounds = mesh.bounds
    min_bounds, max_bounds = bounds[0], bounds[1]
    body_height = max_bounds[2] - min_bounds[2]
    body_midpoint = (max_bounds + min_bounds) / 2

    # Define more sophisticated criteria for the body selection
    if section == 'body':
        # Assuming the body is within certain height and centered horizontally
        selected_indices = np.where((mesh.vertices[:, 2] >= min_bounds[2]) & 
                                    (mesh.vertices[:, 2] <= min_bounds[2] + 0.75 * body_height) &
                                    (mesh.vertices[:, 0] >= min_bounds[0]) &
                                    (mesh.vertices[:, 0] <= max_bounds[0]))[0]

    elif section == 'roof':
        # Roof is the top quarter of the height
        selected_indices = np.where(mesh.vertices[:, 2] > min_bounds[2] + 0.75 * body_height)[0]

    # Ensure indices are within bounds
    valid_indices = selected_indices[selected_indices < len(mesh.vertices)]

    # Create or initialize UV mapping
    uv = np.random.rand(len(mesh.vertices), 2)

    # Apply the rust texture only to the selected indices
    new_uv = np.copy(uv)
    new_uv[valid_indices, :] = np.random.rand(len(valid_indices), 2)

    # Create a new material with the rust texture image
    material = trimesh.visual.texture.SimpleMaterial(image=rust_texture)
    new_visuals = trimesh.visual.TextureVisuals(uv=new_uv, material=material)

    # Create a new mesh with the updated visuals
    textured_mesh = trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.faces, visual=new_visuals, validate=True, process=False)

    # Save the modified mesh to a temporary file for Gradio visualization
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.glb')
    textured_mesh.export(temp_file.name, file_type='glb')
    temp_file.close()
    return temp_file.name

with gr.Blocks() as app:
    gr.Markdown("### Apply Rust Texture to Train Body Section")
    original_model = gr.Model3D('train.glb', label="Original Model")
    modified_model = gr.Model3D(label="Textured Model")

    section_dropdown = gr.Dropdown(choices=['body', 'roof'], label="Select Section")
    button = gr.Button("Apply Rust Texture to Section")
    button.click(apply_rust_texture, inputs=section_dropdown, outputs=modified_model)

app.launch()