import gradio as gr
import plotly.graph_objs as go
import trimesh
import numpy as np
from PIL import Image
import torch
from diffusers import StableDiffusionPipeline
import os
import matplotlib.pyplot as plt

# Load the Stable Diffusion model for text-to-image generation
device = "cuda" if torch.cuda.is_available() else "cpu"
pipeline = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4").to(device)

# Get the current directory
current_dir = os.getcwd()
# Default object file path
DEFAULT_OBJ_FILE = os.path.join(current_dir, "female.obj")
# Temporary texture file path
TEMP_TEXTURE_FILE = os.path.join(current_dir, "generated_texture.png")
# File path to save the 2D image
OUTPUT_IMAGE_FILE = os.path.join(current_dir, "output_image.png")
DEFAULT_GLB_FILE= os.path.join(current_dir, "vroid_girl1.glb")

def apply_texture(mesh, texture_file):
    texture_image = Image.open(texture_file)
    uv_coords = mesh.visual.uv
    uv_coords = np.clip(uv_coords, 0, 1)
    texture_colors = np.array([
        texture_image.getpixel((
            int(u * (texture_image.width - 1)),
            int(v * (texture_image.height - 1))
        )) for u, v in uv_coords
    ])
    texture_colors = texture_colors / 255.0
    return texture_colors

def display_3d_object(obj_file, texture_file, light_intensity, ambient_intensity, color):
    file_extension = obj_file.split('.')[-1].lower()
    if file_extension == 'obj':
        mesh = trimesh.load(obj_file)
    elif file_extension == 'glb':
        mesh = load_glb_file(obj_file)
    else:
        raise ValueError("Unsupported file format. Please upload a .obj or .glb file.")

    if texture_file:
        colors = apply_texture(mesh, texture_file)
    else:
        colors = color

    ambient_intensity = max(0, min(ambient_intensity, 1))
    fig = go.Figure(data=[
        go.Mesh3d(
            x=mesh.vertices[:, 0],
            y=mesh.vertices[:, 1],
            z=mesh.vertices[:, 2],
            i=mesh.faces[:, 0],
            j=mesh.faces[:, 1],
            k=mesh.faces[:, 2],
            facecolor=colors if texture_file else None,
            color=color if not texture_file else None,
            opacity=0.50,
            lighting=dict(
                ambient=ambient_intensity,
                diffuse=light_intensity,
                specular=0.5,
                roughness=0.1,
                fresnel=0.2
            ),
            lightposition=dict(
                x=100,
                y=200,
                z=300
            )
        )
    ])
    fig.update_layout(scene=dict(aspectmode='data'))
    return fig

def load_glb_file(filename):
    trimesh_scene = trimesh.load(filename)
    if isinstance(trimesh_scene, trimesh.Scene):
        mesh = trimesh_scene.dump(concatenate=True)
    else:
        mesh = trimesh_scene
    return mesh

def extract_uv_texture(obj_file):
    mesh = trimesh.load(obj_file)
    if mesh.visual.uv is None or len(mesh.visual.uv) == 0:
        raise ValueError("The mesh does not have UV mapping information.")
    
    texture_image = None
    if mesh.visual.material.image is not None:
        texture_image = mesh.visual.material.image
    else:
        texture_size = 1024
        texture_image = Image.new('RGB', (texture_size, texture_size), color=(255, 255, 255))
    
    return texture_image

def generate_clothing_image(prompt):
    image = pipeline(prompt).images[0]
    image.save(TEMP_TEXTURE_FILE)
    return TEMP_TEXTURE_FILE, image

def update_texture_display(prompt, texture_file):
    if prompt:
        texture_path, image = generate_clothing_image(prompt)
        return image
    elif texture_file:
        return Image.open(texture_file)
    return None

with gr.Blocks() as demo:
    gr.Markdown("## 3D Object Viewer with Custom Texture, Color, and Adjustable Lighting")

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Texture Options")
            prompt_input = gr.Textbox(label="Enter a Prompt to Generate Texture", placeholder="Type a prompt...")
            generate_button = gr.Button("Generate Texture")
            texture_file = gr.File(label="Upload Texture file (PNG or JPG, optional)", type="filepath")
            texture_preview = gr.Image(label="Texture Preview", visible=True)

            gr.Markdown("### Lighting & Color Settings")
            light_intensity_slider = gr.Slider(minimum=0, maximum=2, step=0.1, value=0.8, label="Light Intensity")
            ambient_intensity_slider = gr.Slider(minimum=0, maximum=1, step=0.1, value=0.5, label="Ambient Intensity")
            color_picker = gr.ColorPicker(value="#D3D3D3", label="Object Color")
            submit_button = gr.Button("Submit")
            obj_file = gr.File(label="Upload OBJ or GLB file", value=DEFAULT_OBJ_FILE, type='filepath')

        with gr.Column(scale=2):
            display = gr.Plot(label="3D Viewer")
            extract_button = gr.Button("Extract UV Texture")
            output_image = gr.Image(label="Extracted UV Texture", visible=True)

    def update_display(file, texture, light_intensity, ambient_intensity, color):
        texture_to_use = TEMP_TEXTURE_FILE if os.path.exists(TEMP_TEXTURE_FILE) else texture
        return display_3d_object(file, texture_to_use, light_intensity, ambient_intensity, color)

    def extract_and_display_uv_texture(file):
        uv_texture_image = extract_uv_texture(file)
        uv_texture_image.save(OUTPUT_IMAGE_FILE)
        return uv_texture_image

    submit_button.click(fn=update_display, inputs=[obj_file, texture_file, light_intensity_slider, ambient_intensity_slider, color_picker], outputs=display)
    generate_button.click(fn=update_texture_display, inputs=[prompt_input, texture_file], outputs=texture_preview)
    texture_file.change(fn=update_texture_display, inputs=[prompt_input, texture_file], outputs=texture_preview)
    extract_button.click(fn=extract_and_display_uv_texture, inputs=[obj_file], outputs=output_image)

    demo.load(fn=update_display, inputs=[obj_file, texture_file, light_intensity_slider, ambient_intensity_slider, color_picker], outputs=display)

    gr.Examples(
        examples=[[DEFAULT_OBJ_FILE, None],[DEFAULT_GLB_FILE, None]],
        inputs=[obj_file, texture_file],
        label="Example Files"
    )

demo.launch(debug=True)