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

def apply_masked_texture(section):
    # Load the original mesh
    mesh = trimesh.load('train.glb', force='mesh')

    # Load the new texture image
    im = Image.open('rust_steel.png').convert('RGB')

    # Calculate bounding box coordinates
    bounds = mesh.bounds
    min_bounds, max_bounds = bounds[0], bounds[1]
    mid_x = (max_bounds[0] + min_bounds[0]) / 2
    mid_y = (max_bounds[1] + min_bounds[1]) / 2
    mid_z = (max_bounds[2] + min_bounds[2]) / 2

    # Define sections
    sections = {
        'upper': np.where(mesh.vertices[:, 2] > mid_z)[0],
        'lower': np.where(mesh.vertices[:, 2] <= mid_z)[0],
        'middle': np.where((mesh.vertices[:, 0] >= min_bounds[0]) & (mesh.vertices[:, 0] <= max_bounds[0]) &
                           (mesh.vertices[:, 1] >= min_bounds[1]) & (mesh.vertices[:, 1] <= max_bounds[1]) &
                           (mesh.vertices[:, 2] >= min_bounds[2]) & (mesh.vertices[:, 2] <= max_bounds[2]))[0],
        'top': np.where(mesh.vertices[:, 1] > mid_y)[0],
        'right': np.where(mesh.vertices[:, 0] > mid_x)[0]
    }

    # Get indices for the selected section
    selected_indices = sections[section]

    # Copy the original visual properties (colors, UVs)
    if mesh.visual.vertex_colors is not None:
        new_vertex_colors = np.copy(mesh.visual.vertex_colors)
    else:
        new_vertex_colors = np.ones((len(mesh.vertices), 4)) * 255  # Assuming original is white or neutral

    # Apply the new texture to only the selected section
    for idx in selected_indices:
        new_vertex_colors[idx, :] = [255, 255, 0, 255]  # Example yellow RGBA color

    # Apply the new visual properties
    textured_mesh = trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.faces, visual=trimesh.visual.ColorVisuals(vertex_colors=new_vertex_colors), validate=True, process=False)

    # Save 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("### 3D Model Texture Application")
    original_model = gr.Model3D('train.glb', label="Original Model")
    modified_model = gr.Model3D(label="Textured Model")

    section_dropdown = gr.Dropdown(choices=['upper', 'lower', 'middle', 'top', 'right'], label="Select Section")
    button = gr.Button("Apply Texture to Section")
    button.click(apply_masked_texture, inputs=section_dropdown, outputs=modified_model)

app.launch()