import gradio as gr
import os
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
import numpy as np
from fenics import *

# Function to process CAD file and extract geometric properties
def process_cad_file(cad_file_path):
    file_extension = os.path.splitext(cad_file_path)[1].lower()
    if file_extension in ['.step', '.stp']:
        geometry_info = {"type": "block", "dimensions": (150, 75, 25)}
    elif file_extension in ['.iges', '.igs']:
        geometry_info = {"type": "block", "dimensions": (120, 60, 20)}
    elif file_extension == '.stl':
        geometry_info = {"type": "block", "dimensions": (200, 100, 50)}
    else:
        raise ValueError(f"Unsupported CAD file format: {file_extension}")
    return geometry_info

# Generate 2D visualization
def generate_2d_image(dimensions):
    length, width, _ = dimensions
    plt.figure(figsize=(5, 5))
    plt.plot([0, length, length, 0, 0], [0, 0, width, width, 0], 'b-')
    plt.fill_between([0, length], [0, 0], [width, width], color='lightblue', alpha=0.3)
    plt.title("2D Top View")
    plt.xlabel("Length")
    plt.ylabel("Width")
    plt.grid()
    image_path = "2d_image.png"
    plt.savefig(image_path)
    plt.close()
    return image_path

# Generate 3D visualization
def generate_3d_image(dimensions):
    length, width, height = dimensions
    vertices = np.array([
        [0, 0, 0], [length, 0, 0], [length, width, 0], [0, width, 0],
        [0, 0, height], [length, 0, height], [length, width, height], [0, width, height]
    ])
    faces = [
        [vertices[j] for j in [0, 1, 2, 3]],  # Bottom
        [vertices[j] for j in [4, 5, 6, 7]],  # Top
        [vertices[j] for j in [0, 1, 5, 4]],  # Front
        [vertices[j] for j in [2, 3, 7, 6]],  # Back
        [vertices[j] for j in [1, 2, 6, 5]],  # Right
        [vertices[j] for j in [0, 3, 7, 4]]   # Left
    ]
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    for face in faces:
        poly = Poly3DCollection([face], alpha=0.5, edgecolor='r')
        ax.add_collection3d(poly)
    ax.set_xlabel("Length")
    ax.set_ylabel("Width")
    ax.set_zlabel("Height")
    ax.set_xlim([0, length])
    ax.set_ylim([0, width])
    ax.set_zlim([0, height])
    plt.title("3D View")
    image_path = "3d_image.png"
    plt.savefig(image_path)
    plt.close()
    return image_path

# Simulate using FEniCS
def simulate_with_fenics(dimensions):
    length, width, height = dimensions

    # Material properties
    E = 200e9  # Young's modulus (Pa)
    nu = 0.3   # Poisson's ratio
    mu = E / (2 * (1 + nu))
    lmbda = E * nu / ((1 + nu) * (1 - 2 * nu))

    # Create mesh
    mesh = BoxMesh(Point(0, 0, 0), Point(length, width, height), 10, 5, 3)

    # Define function space for displacement
    V = VectorFunctionSpace(mesh, "P", 1)

    # Boundary conditions
    def left_boundary(x, on_boundary):
        return on_boundary and near(x[0], 0)

    bc = DirichletBC(V, Constant((0, 0, 0)), left_boundary)

    # Define strain and stress
    def epsilon(u):
        return sym(grad(u))

    def sigma(u):
        return lmbda * div(u) * Identity(3) + 2 * mu * epsilon(u)

    # Define variational problem
    u = TrialFunction(V)
    v = TestFunction(V)
    f = Constant((0, 0, -1000))  # Uniform load
    a = inner(sigma(u), epsilon(v)) * dx
    L = dot(f, v) * dx

    # Solve the problem
    u = Function(V)
    solve(a == L, u, bc)

    # Post-process results
    V_magnitude = FunctionSpace(mesh, "P", 1)
    u_magnitude = sqrt(dot(u, u))
    u_magnitude = project(u_magnitude, V_magnitude)
    max_displacement = u_magnitude.vector().max()

    # Save results
    file = File("fenics_displacement.pvd")
    file << u

    return f"Maximum displacement: {max_displacement:.6f} mm"

# Main function
def process_and_simulate(cad_file):
    try:
        # Process CAD file
        geometry_info = process_cad_file(cad_file.name)
        dimensions = geometry_info["dimensions"]

        # Generate visualizations
        image_2d = generate_2d_image(dimensions)
        image_3d = generate_3d_image(dimensions)

        # Simulate using FEniCS
        simulation_output = simulate_with_fenics(dimensions)

        return "fenics_displacement.pvd", image_2d, image_3d, simulation_output
    except Exception as e:
        return f"Error: {str(e)}", None, None, None

# Gradio Interface
iface = gr.Interface(
    fn=process_and_simulate,
    inputs=gr.File(label="Upload CAD File (STEP, IGES, STL)"),
    outputs=[
        gr.File(label="Displacement Results (PVD)"),
        gr.Image(label="2D Visualization"),
        gr.Image(label="3D Visualization"),
        gr.Textbox(label="Simulation Output Log"),
    ],
    title="FEniCS-Based Simulation Processor",
    description=(
        "Upload a CAD file (STEP, IGES, STL) to extract geometry, generate "
        "2D/3D visualizations, and simulate using FEniCS."
    ),
)

# Launch the Interface
iface.launch()