import cadquery as cq
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
import gradio as gr
from cadquery.occ_impl.exporters import exportShape


# Function to create 3D ear-like shell using CadQuery
def create_bird_ear_shell(height, radius, thickness):
    try:
        # Create outer shell
        outer_shell = (
            cq.Workplane("XY")
            .lineTo(0, height)
            .radiusArc((radius, 0), radius)
            .close()
            .revolve()
        )
        # Create inner shell
        inner_shell = (
            cq.Workplane("XY")
            .lineTo(0, height - thickness)
            .radiusArc((radius - thickness, 0), radius - thickness)
            .close()
            .revolve()
        )
        # Subtract inner shell from outer to create a hollow shell
        shell = outer_shell.cut(inner_shell)
        return shell
    except Exception as e:
        print(f"Error in shell creation: {e}")
        return None


# Function to generate a 2D visualization using Matplotlib
def generate_2d_image(height, radius):
    try:
        theta = np.linspace(0, 2 * np.pi, 100)
        x_outer = radius * np.cos(theta)
        y_outer = height * np.sin(theta) / max(np.sin(theta))

        plt.figure(figsize=(6, 6))
        plt.plot(x_outer, y_outer, label="Outer Shell", color="blue")
        plt.fill_between(x_outer, y_outer, color="lightblue", alpha=0.5)
        plt.axhline(0, color="black", linewidth=0.5, linestyle="--")
        plt.axvline(0, color="black", linewidth=0.5, linestyle="--")
        plt.title("2D Cross-Section of Ear-Like Bird Shell")
        plt.xlabel("Radius (mm)")
        plt.ylabel("Height (mm)")
        plt.grid()
        plt.legend()
        plt.savefig("2d_visualization.png")
        return "2d_visualization.png"
    except Exception as e:
        print(f"Error in generating 2D image: {e}")
        return None


# Function to generate a 3D STL file using CadQuery
def generate_3d_stl(height, radius, thickness):
    try:
        bird_shell_model = create_bird_ear_shell(height, radius, thickness)
        if bird_shell_model is None:
            raise ValueError("Invalid CAD geometry.")
        stl_path = "ear_bird_shell.stl"
        with open(stl_path, "wb") as f:
            exportShape(bird_shell_model.val(), "STL", f)
        return stl_path
    except Exception as e:
        print(f"Error in generating 3D STL: {e}")
        return None


# Function to create a 3D surface plot using Matplotlib
def plot_3d_surface(height, radius):
    try:
        fig = plt.figure(figsize=(8, 8))
        ax = fig.add_subplot(111, projection="3d")

        # Generate dummy data for the 3D surface
        theta = np.linspace(0, 2 * np.pi, 100)
        z = np.linspace(0, height, 50)
        theta, z = np.meshgrid(theta, z)
        x = (radius + 5 * np.cos(z)) * np.cos(theta)
        y = (radius + 5 * np.cos(z)) * np.sin(theta)

        # Plot the surface
        ax.plot_surface(x, y, z, cmap="viridis", edgecolor="k", alpha=0.8)
        ax.set_title("3D Surface of Shell")
        ax.set_xlabel("X")
        ax.set_ylabel("Y")
        ax.set_zlabel("Height")
        plt.savefig("3d_surface_plot.png")
        return "3d_surface_plot.png"
    except Exception as e:
        print(f"Error in generating 3D plot: {e}")
        return None


# Main function to generate visuals
def generate_visuals(height, radius, thickness):
    two_d_image = generate_2d_image(height, radius)
    three_d_model = generate_3d_stl(height, radius, thickness)
    three_d_plot = plot_3d_surface(height, radius)
    return two_d_image, three_d_model, three_d_plot


# Gradio Interface
interface = gr.Interface(
    fn=generate_visuals,
    inputs=[
        gr.Number(label="Height (mm)", value=50.0),
        gr.Number(label="Radius (mm)", value=25.0),
        gr.Number(label="Wall Thickness (mm)", value=3.0),
    ],
    outputs=[
        gr.Image(label="2D Visualization"),
        gr.File(label="Download 3D Model (STL)"),
        gr.Image(label="3D Surface Visualization"),
    ],
    title="Ear-Like Bird Shell Simulation",
    description="Input parameters to create an ear-like bird shell model. Visualize it in 2D and 3D, and download the 3D STL file."
)

interface.launch()