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

def visualize_texture(section):
    mesh = trimesh.load('train.glb', force='mesh')
    im = Image.open('defect.jpg').convert('RGB')

    # Calculate bounds
    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 selected section
    selected_indices = sections[section]

    # Create 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=im)
    color_visuals = trimesh.visual.TextureVisuals(uv=new_uv, image=im, material=material)
    textured_mesh = trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.faces, visual=color_visuals, validate=True, process=False)

    return textured_mesh.export(file_type='glb')

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")
    
    with gr.Row():
        button_upper = gr.Button("Texture Upper Part", data="upper")
        button_lower = gr.Button("Texture Lower Part", data="lower")
        button_middle = gr.Button("Texture Middle", data="middle")
        button_top = gr.Button("Texture Top", data="top")
        button_right = gr.Button("Texture Right", data="right")

    buttons = [button_upper, button_lower, button_middle, button_top, button_right]
    for button in buttons:
        button.click(visualize_texture, inputs=button, outputs=modified_model)

app.launch()