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

def visualize_dynamic_texture(x_min, x_max, y_min, y_max, z_min, z_max):
    # Load the original mesh
    mesh = trimesh.load('train.glb', force='mesh')
    rust_texture = Image.open('rust_steel.png').convert('RGB')

    # Identify vertices within the given bounding box
    selected_indices = np.where((mesh.vertices[:, 0] >= x_min) & (mesh.vertices[:, 0] <= x_max) &
                                (mesh.vertices[:, 1] >= y_min) & (mesh.vertices[:, 1] <= y_max) &
                                (mesh.vertices[:, 2] >= z_min) & (mesh.vertices[:, 2] <= z_max))[0]

    # Initialize UV coordinates
    uv = np.random.rand(len(mesh.vertices), 2)
    new_uv = np.zeros_like(uv)
    new_uv[selected_indices, :] = uv[selected_indices, :]

    # Create material and apply texture
    material = trimesh.visual.texture.SimpleMaterial(image=rust_texture)
    color_visuals = trimesh.visual.TextureVisuals(uv=new_uv, image=rust_texture, material=material)
    textured_mesh = trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.faces, visual=color_visuals, validate=True, process=False)

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

# Get bounding box values from the original mesh to define slider ranges
mesh = trimesh.load('train.glb', force='mesh')
bounds = mesh.bounds
x_min_range, x_max_range = bounds[0][0], bounds[1][0]
y_min_range, y_max_range = bounds[0][1], bounds[1][1]
z_min_range, z_max_range = bounds[0][2], bounds[1][2]

with gr.Blocks() as app:
    gr.Markdown("### Select Model Portion with Sliders")
    original_model = gr.Model3D('train.glb', label="Original Model")
    modified_model = gr.Model3D(label="Textured Model")

    # Add sliders for bounding box selection
    with gr.Row():
        x_min_slider = gr.Slider(minimum=x_min_range, maximum=x_max_range, step=0.01, label="X Min", value=x_min_range)
        x_max_slider = gr.Slider(minimum=x_min_range, maximum=x_max_range, step=0.01, label="X Max", value=x_max_range)

    with gr.Row():
        y_min_slider = gr.Slider(minimum=y_min_range, maximum=y_max_range, step=0.01, label="Y Min", value=y_min_range)
        y_max_slider = gr.Slider(minimum=y_min_range, maximum=y_max_range, step=0.01, label="Y Max", value=y_max_range)

    with gr.Row():
        z_min_slider = gr.Slider(minimum=z_min_range, maximum=z_max_range, step=0.01, label="Z Min", value=z_min_range)
        z_max_slider = gr.Slider(minimum=z_min_range, maximum=z_max_range, step=0.01, label="Z Max", value=z_max_range)

    button = gr.Button("Apply Texture")
    button.click(visualize_dynamic_texture, inputs=[x_min_slider, x_max_slider, y_min_slider, y_max_slider, z_min_slider, z_max_slider], outputs=modified_model)

app.launch()