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

# Function to visualize texture based on selection criteria
def visualize_dynamic_texture(predefined_section, 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')

    # Predefined sections
    if predefined_section == 'right compartments':
        selected_indices = np.where(mesh.vertices[:, 0] > (train_bounds[0][0] + train_bounds[1][0]) / 2)[0]
    elif predefined_section == 'left compartments':
        selected_indices = np.where(mesh.vertices[:, 0] <= (train_bounds[0][0] + train_bounds[1][0]) / 2)[0]
    elif predefined_section == 'freight_body':
        selected_indices = np.where((mesh.vertices[:, 0] >= train_bounds[0][0]) & (mesh.vertices[:, 0] <= train_bounds[1][0]) &
                                    (mesh.vertices[:, 2] <= (train_bounds[0][2] + train_bounds[1][2]) / 2))[0]
    elif predefined_section == 'custom':
        # Use custom sliders for custom section
        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]
    else:
        selected_indices = np.array([])

    # Initialize UV mapping
    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 the new 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

# Load train model to establish bounding box ranges
train_model = trimesh.load('train.glb', force='mesh')
train_bounds = train_model.bounds

# Get slider ranges based on train model bounds
x_min_range, x_max_range = train_bounds[0][0], train_bounds[1][0]
y_min_range, y_max_range = train_bounds[0][1], train_bounds[1][1]
z_min_range, z_max_range = train_bounds[0][2], train_bounds[1][2]

# Gradio UI with a single window, dynamic updates, and real-time changes
with gr.Blocks() as app:
    gr.Markdown("### 3D Model Texture Application with Predefined and Custom Sections")
    original_model = gr.Model3D('train.glb', label="Original Model")
    modified_model = gr.Model3D(label="Model with Applied Texture")

    # Dropdown for predefined and custom selection
    section_dropdown = gr.Radio(choices=['right compartments', 'left compartments', 'freight_body', 'custom'], label="Select Section", value='custom')

    # Custom sliders for the bounding box selection
    with gr.Row(visible=True) as custom_controls:
        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)

        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)

        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)

    # Toggle visibility of custom controls
    def toggle_custom_controls(predefined_section):
        return gr.update(visible=(predefined_section == 'custom'))

    section_dropdown.change(fn=toggle_custom_controls, inputs=section_dropdown, outputs=custom_controls)

    # Update model dynamically
    def update_model(predefined_section, x_min, x_max, y_min, y_max, z_min, z_max):
        return visualize_dynamic_texture(predefined_section, x_min, x_max, y_min, y_max, z_min, z_max)

    # Add event listeners for real-time updates when sliders or dropdown change
    inputs = [section_dropdown, x_min_slider, x_max_slider, y_min_slider, y_max_slider, z_min_slider, z_max_slider]
    for input_component in inputs:
        input_component.change(fn=update_model, inputs=inputs, outputs=modified_model)

app.launch()