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

# Function to create the CAD model
def create_bird_ear_shell(height, radius, thickness):
    outer_shell = (
        cq.Workplane("XY")
        .lineTo(0, height)                # Vertical height
        .radiusArc((radius, 0), radius)  # Curve to base
        .close()
        .revolve()                       # Revolve to create 3D shape
    )
    inner_shell = (
        cq.Workplane("XY")
        .lineTo(0, height - thickness)
        .radiusArc((radius - thickness, 0), radius - thickness)
        .close()
        .revolve()
    )
    shell = outer_shell.cut(inner_shell)
    return shell

# Function to generate a 2D visualization
def generate_2d_image(height, radius):
    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"

# Function to generate a 3D STL file
def generate_3d_stl(height, radius, thickness):
    bird_shell_model = create_bird_ear_shell(height, radius, thickness)
    stl_path = "ear_bird_shell.stl"
    with open(stl_path, "wb") as f:
        exportShape(bird_shell_model.val(), "STL", f)
    return stl_path

# Gradio interface functions
def generate_visuals(height, radius, thickness):
    # Generate 2D image
    two_d_image = generate_2d_image(height, radius)
    # Generate 3D STL file
    three_d_model = generate_3d_stl(height, radius, thickness)
    return two_d_image, three_d_model

# 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)"),
    ],
    title="Ear-Like Bird Shell Simulation",
    description="Input the parameters to create an ear-like bird shell model. You can visualize it in 2D and download the 3D STL file for further use."
)

# Launch the app
interface.launch()