import gradio as gr
import open3d as o3d
import matplotlib.pyplot as plt
import numpy as np
import os
from diffusers import DiffusionPipeline
import torch
from datasets import load_dataset
import transformers
import accelerate

# Load Stable Diffusion XL model
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")
pipe.to("cuda" if torch.cuda.is_available() else "cpu")

# Ensure directories exist
os.makedirs("outputs/2d_drawings", exist_ok=True)
os.makedirs("outputs/3d_models", exist_ok=True)

# Parse plot size
def parse_plot_size(plot_size):
    try:
        dimensions = plot_size.lower().replace("feet", "").strip().split("x")
        if len(dimensions) != 2:
            raise ValueError("Invalid plot size format. Use '30x40 feet'.")
        width, length = map(float, dimensions)
        return width, length
    except Exception as e:
        raise ValueError(f"Error parsing plot size: {e}")

# Generate 2D drawing with Stable Diffusion
def generate_2d_drawing_stable_diffusion(plot_size, num_floors, storey_height, additional_reqs):
    try:
        # Create a descriptive prompt
        width, length = parse_plot_size(plot_size)
        height = num_floors * storey_height
        prompt = (
            f"Architectural floor plan of a {num_floors}-storey building with a plot size of "
            f"{width}x{length} feet, each storey height {storey_height} feet. "
            f"Additional details: {additional_reqs}"
        )

        # Generate image with Stable Diffusion
        result = pipe(prompt)
        generated_image = result.images[0]

        # Save the generated image
        output_path = os.path.join("outputs/2d_drawings", "2d_drawing_stable_diffusion.png")
        generated_image.save(output_path)
        return output_path
    except Exception as e:
        print(f"Error in generate_2d_drawing_stable_diffusion: {e}")
        raise e

# Generate 3D model
def generate_3d_model(plot_size, num_floors, storey_height):
    try:
        width, length = parse_plot_size(plot_size)
        height = num_floors * storey_height
        mesh = o3d.geometry.TriangleMesh.create_box(width, length, height)
        mesh.compute_vertex_normals()
        output_path = os.path.join("outputs/3d_models", "3d_model.obj")
        o3d.io.write_triangle_mesh(output_path, mesh)
        return output_path
    except Exception as e:
        print(f"Error in generate_3d_model: {e}")
        raise e

# Main function
def generate_building(plot_size, storey_height, num_floors, additional_reqs):
    try:
        # Generate 2D drawing using Stable Diffusion
        drawing_path = generate_2d_drawing_stable_diffusion(plot_size, num_floors, storey_height, additional_reqs)

    except Exception as e:
        print(f"Stable Diffusion failed: {e}. Falling back to Matplotlib.")
        drawing_path = generate_2d_drawing_matplotlib(plot_size, num_floors, storey_height)

    # Generate 3D model
    model_path = generate_3d_model(plot_size, num_floors, storey_height)
    return drawing_path, model_path

# Gradio interface
iface = gr.Interface(
    fn=generate_building,
    inputs=[
        gr.Textbox(label="Plot Size (e.g., 30x40 feet)", placeholder="Enter plot dimensions"),
        gr.Number(label="Storey Height (ft)", value=10),
        gr.Number(label="Number of Floors", value=3),
        gr.Textbox(label="Additional Requirements", placeholder="Optional")
    ],
    outputs=[
        gr.Image(label="2D Drawing"),
        gr.File(label="3D Model (Download)")
    ],
    title="Building Design Generator with AI",
    description="Generates realistic 2D drawings using Stable Diffusion XL and 3D models based on user inputs."
)

if __name__ == "__main__":
    iface.launch()