app.py

# app.py

import gradio as gr
import pyvista as pv
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import numpy as np
import trimesh
import plotly.graph_objects as go
from PIL import Image
import io
import re
import os
import tempfile
import warnings
from shapely.geometry import box, Point
from shapely.affinity import scale

# --- General Setup ---
warnings.filterwarnings('ignore')
plt.style.use('seaborn-v0_8-darkgrid')
# Ensure a temporary directory exists for file outputs
TEMP_DIR = "temp_outputs"
if not os.path.exists(TEMP_DIR):
    os.makedirs(TEMP_DIR)


# ###########################################################################
# ## 🛠️ TOOL 1: CAD to 2D Orthographic Views
# ###########################################################################

def generate_ortho_views(input_file):
    """
    Takes an uploaded 3D model file and generates a single image
    containing the front, top, and side orthographic views with dimensions.
    """
    if input_file is None:
        raise gr.Error("Please upload a 3D model file (e.g., .stl, .obj).")
    
    try:
        mesh = pv.read(input_file.name)
    except Exception as e:
        raise gr.Error(f"Error loading the 3D model: {e}")

    # --- Generate Base Views with PyVista ---
    views = ['xy', 'xz', 'yz']
    temp_image_files = []
    
    # Configure plotter for headless environment
    plotter = pv.Plotter(off_screen=True, window_size=[1000, 1000])
    plotter.background_color = 'white'
    plotter.enable_parallel_projection()
    
    outline_mesh = mesh.extract_feature_edges()
    outline_kwargs = {'color': 'black', 'line_width': 2}
    
    for view in views:
        plotter.clear()
        plotter.add_mesh(outline_mesh, **outline_kwargs)
        plotter.camera_position = view
        plotter.reset_camera()
        plotter.camera.zoom(1.2)
        
        # Save to a temporary file
        temp_fd, temp_path = tempfile.mkstemp(suffix=".png", dir=TEMP_DIR)
        os.close(temp_fd)
        plotter.screenshot(temp_path)
        temp_image_files.append(temp_path)

    # --- Combine and Add Dimensions with Matplotlib ---
    fig, axes = plt.subplots(1, 3, figsize=(18, 6))
    fig.patch.set_facecolor('white')
    
    bounds = mesh.bounds
    x_dim, y_dim, z_dim = (bounds[1] - bounds[0]), (bounds[3] - bounds[2]), (bounds[5] - bounds[4])
    
    dims_data = [(x_dim, y_dim), (x_dim, z_dim), (y_dim, z_dim)]
    view_titles = ["Front View (XY)", "Top View (XZ)", "Side View (YZ)"]
    
    for i, ax in enumerate(axes):
        img = plt.imread(temp_image_files[i])
        ax.imshow(img)
        ax.set_title(view_titles[i], fontsize=12, pad=10)
        ax.axis('off')
        
        img_w, img_h = img.shape[1], img.shape[0]
        dim1, dim2 = dims_data[i]
        
        # Horizontal Dimension
        ax.annotate(f"{dim1:.2f}", xy=(img_w * 0.2, img_h * 0.9), xytext=(img_w * 0.8, img_h * 0.9),
                    arrowprops=dict(arrowstyle='<->', color='dimgray', lw=1.5),
                    va='bottom', ha='center', fontsize=10, color='black')

        # Vertical Dimension
        ax.annotate(f"{dim2:.2f}", xy=(img_w * 0.1, img_h * 0.2), xytext=(img_w * 0.1, img_h * 0.8),
                    arrowprops=dict(arrowstyle='<->', color='dimgray', lw=1.5),
                    rotation=90, va='center', ha='right', fontsize=10, color='black')

    plt.tight_layout(rect=[0, 0, 1, 0.96])
    
    # Save final combined image
    final_fd, final_path = tempfile.mkstemp(suffix=".png", dir=TEMP_DIR)
    os.close(final_fd)
    plt.savefig(final_path, dpi=200, facecolor=fig.get_facecolor())
    plt.close(fig)
    
    # Clean up intermediate files
    for temp_file in temp_image_files:
        os.remove(temp_file)
        
    return final_path


# ###########################################################################
# ## 💨 TOOL 2: 2D CFD Simulator
# ###########################################################################

def run_cfd_simulation(Lx, Ly, Nx, Ny, u_in, rho, mu, ox1, oy1, ox2, oy2, max_iter, p_iter, tol, progress=gr.Progress()):
    """
    Runs the 2D CFD simulation and returns the result plot and a summary.
    """
    # --- Setup ---
    dx, dy = Lx / (Nx - 1), Ly / (Ny - 1)
    x, y = np.linspace(0, Lx, Nx), np.linspace(0, Ly, Ny)
    X, Y = np.meshgrid(x, y)
    dt = 0.001
    nu = mu / rho
    alpha_u, alpha_v, alpha_p = 0.4, 0.4, 0.1 # Relaxation factors

    # --- Obstacle Mask ---
    obstacle_mask = np.zeros((Ny, Nx), dtype=bool)
    if not (ox1 == 0 and oy1 == 0 and ox2 == 0 and oy2 == 0):
        i_ox1, j_oy1 = int(round(ox1 / dx)), int(round(oy1 / dy))
        i_ox2, j_oy2 = int(round(ox2 / dx)), int(round(oy2 / dy))
        obstacle_mask[j_oy1:j_oy2, i_ox1:i_ox2] = True

    # --- Initialization ---
    u, v, p = np.zeros((Ny, Nx)), np.zeros((Ny, Nx)), np.ones((Ny, Nx))
    reynolds = u_in * Lx / nu
    
    # --- Main Loop ---
    progress(0, desc="Starting Simulation...")
    for step in range(max_iter):
        un, vn = u.copy(), v.copy()

        # Boundary Conditions
        u[:, 0], v[:, 0], p[:, 0] = u_in, 0.0, p[:, 1] # Inlet
        u[:, -1], v[:, -1], p[:, -1] = u[:, -2], v[:, -2], 0.0 # Outlet
        u[0, :], v[0, :], p[0, :] = 0.0, 0.0, p[1, :] # Bottom Wall
        u[-1, :], v[-1, :], p[-1, :] = 0.0, 0.0, p[-2, :] # Top Wall
        u[obstacle_mask], v[obstacle_mask] = 0.0, 0.0

        # Pressure Source Term (b)
        b = rho * (1/dt * ((u[1:-1, 2:] - u[1:-1, 0:-2]) / (2*dx) + (v[2:, 1:-1] - v[0:-2, 1:-1]) / (2*dy)))
        
        # Pressure Poisson Solver
        for _ in range(p_iter):
            pn = p.copy()
            p[1:-1, 1:-1] = alpha_p * (((pn[1:-1, 2:] + pn[1:-1, 0:-2]) * dy**2 + (pn[2:, 1:-1] + pn[0:-2, 1:-1]) * dx**2 - b * dx**2 * dy**2) / (2 * (dx**2 + dy**2))) + (1 - alpha_p) * p[1:-1, 1:-1]
            p[obstacle_mask] = 0 # Dummy value inside obstacle
            
        # Momentum Update
        u[1:-1, 1:-1] = alpha_u * (un[1:-1, 1:-1] - un[1:-1, 1:-1] * dt / dx * (un[1:-1, 1:-1] - un[1:-1, 0:-2]) \
                          - vn[1:-1, 1:-1] * dt / dy * (un[1:-1, 1:-1] - un[0:-2, 1:-1]) \
                          - dt / (2 * rho * dx) * (p[1:-1, 2:] - p[1:-1, 0:-2]) \
                          + nu * (dt / dx**2 * (un[1:-1, 2:] - 2 * un[1:-1, 1:-1] + un[1:-1, 0:-2]) \
                          + dt / dy**2 * (un[2:, 1:-1] - 2 * un[1:-1, 1:-1] + un[0:-2, 1:-1]))) + (1 - alpha_u) * u[1:-1, 1:-1]

        v[1:-1, 1:-1] = alpha_v * (vn[1:-1, 1:-1] - un[1:-1, 1:-1] * dt / dx * (vn[1:-1, 1:-1] - vn[1:-1, 0:-2]) \
                          - vn[1:-1, 1:-1] * dt / dy * (vn[1:-1, 1:-1] - vn[0:-2, 1:-1]) \
                          - dt / (2 * rho * dy) * (p[2:, 1:-1] - p[0:-2, 1:-1]) \
                          + nu * (dt / dx**2 * (vn[1:-1, 2:] - 2 * vn[1:-1, 1:-1] + vn[1:-1, 0:-2]) \
                          + dt / dy**2 * (vn[2:, 1:-1] - 2 * vn[1:-1, 1:-1] + vn[0:-2, 1:-1]))) + (1 - alpha_v) * v[1:-1, 1:-1]

        diff = np.max(np.abs(u - un))
        if diff < tol:
            break
        
        if step % 100 == 0:
            progress(step / max_iter, desc=f"Iteration {step}, Diff: {diff:.2e}")

    # --- Plotting ---
    u[obstacle_mask], v[obstacle_mask], p[obstacle_mask] = np.nan, np.nan, np.nan
    velocity_mag = np.sqrt(u**2 + v**2)
    fig, axes = plt.subplots(1, 2, figsize=(16, 6))

    # Velocity Plot
    im = axes[0].pcolormesh(X, Y, velocity_mag, cmap=cm.jet, shading='auto')
    fig.colorbar(im, ax=axes[0], label='Velocity (m/s)')
    axes[0].streamplot(X, Y, u, v, color='white', linewidth=0.8, density=1.5)
    axes[0].set_title('Velocity Field and Streamlines')
    axes[0].set_aspect('equal', adjustable='box')

    # Pressure Plot
    im = axes[1].pcolormesh(X, Y, p, cmap=cm.coolwarm, shading='auto')
    fig.colorbar(im, ax=axes[1], label='Pressure (Pa)')
    axes[1].set_title('Pressure Field')
    axes[1].set_aspect('equal', adjustable='box')

    plt.tight_layout()
    
    summary = (f"**Simulation Summary**\n"
               f"- Convergence reached at iteration: {step}\n"
               f"- Reynolds Number: {reynolds:.2f}\n"
               f"- Max Velocity: {np.nanmax(velocity_mag):.4f} m/s\n"
               f"- Max Pressure: {np.nanmax(p):.4f} Pa")
               
    return fig, summary


# ###########################################################################
# ## 📝 TOOL 3: Text to CAD Model
# ###########################################################################

class TextToCADGenerator:
    def __init__(self):
        self.shapes_library = {
            'cube': self._create_cube, 'box': self._create_cube, 'plate': self._create_plate,
            'sphere': self._create_sphere, 'ball': self._create_sphere,
            'cylinder': self._create_cylinder, 'tube': self._create_cylinder, 'rod': self._create_rod, 'pipe': self._create_pipe,
            'cone': self._create_cone, 'pyramid': self._create_pyramid,
            'torus': self._create_torus, 'ring': self._create_torus, 'washer': self._create_washer, 'bearing': self._create_bearing,
            'gear': self._create_gear, 'bracket': self._create_bracket,
            'screw': self._create_screw, 'bolt': self._create_screw,
            'nut': self._create_nut, 'flange': self._create_flange
        }

    def _extract_dimensions(self, prompt):
        dims = {'length': 10, 'width': 10, 'height': 10, 'radius': 5, 'thickness': 2}
        patterns = {
            'length': r'length.*?(\d+\.?\d*)', 'width': r'width.*?(\d+\.?\d*)',
            'height': r'height.*?(\d+\.?\d*)', 'radius': r'radius.*?(\d+\.?\d*)',
            'diameter': r'diameter.*?(\d+\.?\d*)', 'thickness': r'thick.*?\s(\d+\.?\d*)'
        }
        for key, pattern in patterns.items():
            match = re.search(pattern, prompt, re.IGNORECASE)
            if match:
                val = float(match.group(1))
                if key == 'diameter': dims['radius'] = val / 2
                else: dims[key] = val
        return dims

    def run(self, prompt):
        prompt = prompt.lower()
        params = self._extract_dimensions(prompt)
        shape_type = 'cube'
        for shape_key in self.shapes_library.keys():
            if shape_key in prompt:
                shape_type = shape_key
                break
        
        mesh = self.shapes_library.get(shape_type, self._create_cube)(params)
        drawing = self._generate_2d_drawing(mesh, shape_type.title(), params)
        plotly_fig = self._generate_3d_viz(mesh)
        
        # Save mesh to a temporary STL file
        fd, stl_path = tempfile.mkstemp(suffix=".stl", dir=TEMP_DIR)
        os.close(fd)
        mesh.export(stl_path)
        
        return drawing, plotly_fig, stl_path

    def _generate_2d_drawing(self, mesh, title, params):
        fig, axes = plt.subplots(1, 3, figsize=(15, 5))
        fig.suptitle(f"2D Drawing: {title}", fontsize=16)
        bounds = mesh.bounds
        dims = {'width': bounds[1]-bounds[0], 'depth': bounds[3]-bounds[2], 'height': bounds[5]-bounds[4]}
        
        views = ['Front', 'Top', 'Side']
        extents = [
            (dims['width'], dims['height']),
            (dims['width'], dims['depth']),
            (dims['depth'], dims['height'])
        ]
        
        for i, ax in enumerate(axes):
            ax.set_title(f"{views[i]} View")
            rect = plt.Rectangle((-extents[i][0]/2, -extents[i][1]/2), extents[i][0], extents[i][1], fill=False, edgecolor='black', linewidth=2)
            ax.add_patch(rect)
            ax.set_aspect('equal', 'box')
            ax.set_xlim(-extents[i][0], extents[i][0])
            ax.set_ylim(-extents[i][1], extents[i][1])
            ax.grid(True, linestyle='--', alpha=0.6)
            ax.text(0, -extents[i][1]*0.8, f"W: {extents[i][0]:.2f}", ha='center')
            ax.text(-extents[i][0]*0.8, 0, f"H: {extents[i][1]:.2f}", va='center', rotation=90)
            
        plt.tight_layout(rect=[0, 0.03, 1, 0.95])
        return fig

    def _generate_3d_viz(self, mesh):
        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],
            color='lightblue', opacity=0.9
        )])
        fig.update_layout(title="Interactive 3D Model", scene=dict(aspectmode='data'))
        return fig

    # --- Shape Creation Methods ---
    def _create_cube(self, d): return trimesh.creation.box(extents=[d['length'], d['width'], d['height']])
    def _create_plate(self, d): return trimesh.creation.box(extents=[d['length'], d['width'], d['thickness']])
    def _create_sphere(self, d): return trimesh.creation.icosphere(radius=d['radius'])
    def _create_cylinder(self, d): return trimesh.creation.cylinder(radius=d['radius'], height=d.get('height', d.get('length')))
    def _create_rod(self, d): return trimesh.creation.cylinder(radius=d.get('radius', 1), height=d.get('length', 50))
    def _create_cone(self, d): return trimesh.creation.cone(radius=d['radius'], height=d['height'])
    def _create_pyramid(self, d):
        size = d.get('width', 10)
        vertices = np.array([[0,0,d['height']], [-size/2,-size/2,0], [size/2,-size/2,0], [size/2,size/2,0], [-size/2,size/2,0]])
        faces = np.array([[0,1,2], [0,2,3], [0,3,4], [0,4,1], [1,4,3], [1,3,2]])
        return trimesh.Trimesh(vertices=vertices, faces=faces)
    def _create_torus(self, d): return trimesh.creation.torus(major_radius=d['radius'], minor_radius=d['thickness'])
    def _create_washer(self, d):
        outer = trimesh.creation.cylinder(radius=d['radius'], height=d['thickness'])
        inner = trimesh.creation.cylinder(radius=d['radius']*0.6, height=d['thickness']*1.2)
        return outer.difference(inner)
    def _create_bearing(self, d): return self._create_washer(d)
    def _create_pipe(self, d):
        outer = trimesh.creation.cylinder(radius=d['radius'], height=d.get('length', 100))
        inner = trimesh.creation.cylinder(radius=d['radius']-d['thickness'], height=d.get('length', 100)*1.2)
        return outer.difference(inner)
    def _create_gear(self, d): return trimesh.creation.gear(tooth_number=d.get('teeth', 12), radial_pitch=d['radius'], tooth_width=d['thickness'])
    def _create_bracket(self, d):
        base = trimesh.creation.box(extents=[d['length'], d['width'], d['thickness']])
        wall = trimesh.creation.box(extents=[d['thickness'], d['width'], d['height']])
        wall.apply_translation([d['length']/2 - d['thickness']/2, 0, d['height']/2 - d['thickness']/2])
        return base + wall
    def _create_screw(self, d):
        shaft = trimesh.creation.cylinder(radius=d['radius'], height=d['length'])
        head = trimesh.creation.cylinder(radius=d['radius']*2, height=d['thickness'])
        head.apply_translation([0, 0, d['length']/2])
        return shaft + head
    def _create_nut(self, d):
        hex_prism = trimesh.creation.cylinder(radius=d['radius'], height=d['height'], sections=6)
        hole = trimesh.creation.cylinder(radius=d['radius']*0.5, height=d['height']*1.2)
        return hex_prism.difference(hole)
    def _create_flange(self, d):
        outer = trimesh.creation.cylinder(radius=d['radius'], height=d['height'])
        inner = trimesh.creation.cylinder(radius=d['radius']*0.4, height=d['height']*1.2)
        return outer.difference(inner)

text_to_cad_generator = TextToCADGenerator()


# ###########################################################################
# ## 🔡 TOOL 4: Text to G-Code
# ###########################################################################

def parse_gcode_description(description):
    """Parses text to extract parameters for G-code generation."""
    parsed = {"width": 100, "height": 50, "holes": [], "slots": []}
    
    size_match = re.search(r'(\d+)\s*[xX]\s*(\d+)', description)
    if size_match:
        parsed["width"] = float(size_match.group(1))
        parsed["height"] = float(size_match.group(2))
    
    for hole_match in re.finditer(r'(\d+)\s*mm\s+hole\s+at\s+\((\d+),\s*(\d+)\)', description):
        parsed["holes"].append({'d': float(hole_match.group(1)), 'x': float(hole_match.group(2)), 'y': float(hole_match.group(3))})
    
    return parsed

def generate_text_to_gcode(description):
    """Generates a 3-view drawing, G-code, and an SVG from text."""
    params = parse_gcode_description(description)
    w, h = params['width'], params['height']

    # --- Generate Drawing ---
    fig, ax = plt.subplots(figsize=(8, 5))
    ax.set_aspect('equal')
    shape = box(0, 0, w, h)
    x, y = shape.exterior.xy
    ax.plot(x, y, color='black', label='Outline')
    for hole in params['holes']:
        r = hole['d'] / 2
        circle = Point(hole['x'], hole['y']).buffer(r)
        hx, hy = circle.exterior.xy
        ax.plot(hx, hy, color='blue', label=f'Hole D={hole["d"]}')
    ax.set_xlim(-10, w + 10)
    ax.set_ylim(-10, h + 10)
    ax.grid(True)
    ax.set_title("2D Part Drawing")
    
    # Save drawing
    fd, drawing_path = tempfile.mkstemp(suffix=".png", dir=TEMP_DIR)
    os.close(fd)
    fig.savefig(drawing_path)
    plt.close(fig)
    
    # Save SVG
    fd, svg_path = tempfile.mkstemp(suffix=".svg", dir=TEMP_DIR)
    os.close(fd)
    fig.savefig(svg_path)
    plt.close(fig)

    # --- Generate G-Code ---
    gcode = ["G21 ; Units: mm", "G90 ; Absolute Positioning", "G0 Z5 ; Lift Z"]
    gcode.append("; Cut outline")
    gcode.extend([f"G0 X0 Y0", "G1 Z-1 F100", f"G1 X{w} F500", f"G1 Y{h}", "G1 X0", "G1 Y0", "G0 Z5"])
    for hole in params['holes']:
        x, y, r = hole['x'], hole['y'], hole['d'] / 2
        gcode.append(f"; Drill hole at ({x}, {y})")
        gcode.extend([f"G0 X{x - r} Y{y}", "G1 Z-1 F100", f"G2 I{r} J0 F300 ; Circular interpolation", "G0 Z5"])
    gcode.append("M30 ; End Program")
    
    return drawing_path, "\n".join(gcode), svg_path


# ###########################################################################
# ## 🖼️ GRADIO INTERFACE
# ###########################################################################

with gr.Blocks(theme=gr.themes.Soft(), title="CAD-Wizard") as demo:
    gr.Markdown("# 🧙‍♂️ CAD-Wizard: Your All-in-One CAD Assistant")
    gr.Markdown("Select a tool from the tabs below to get started.")

    with gr.Tabs():
        # --- TAB 1: CAD to 2D ---
        with gr.TabItem("CAD to 2D Views", id=0):
            with gr.Row():
                with gr.Column(scale=1):
                    cad_file_input = gr.File(label="Upload 3D Model (.stl, .obj, .ply)")
                    cad_2d_button = gr.Button("Generate 2D Views", variant="primary")
                with gr.Column(scale=2):
                    cad_2d_output = gr.Image(label="Orthographic Views with Dimensions")
            gr.Examples(
                [["examples/bracket.stl"], ["examples/gear.stl"]],
                inputs=cad_file_input,
                outputs=cad_2d_output,
                fn=generate_ortho_views,
                cache_examples=True
            )

        # --- TAB 2: 2D CFD Simulator ---
        with gr.TabItem("2D CFD Simulator", id=1):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### Fluid & Domain Properties")
                    cfd_u_in = gr.Slider(0.001, 0.1, value=0.01, step=0.001, label="Inlet Velocity (m/s)")
                    cfd_rho = gr.Number(value=1.0, label="Density (kg/m³)")
                    cfd_mu = gr.Number(value=0.02, label="Viscosity (Pa.s)")
                    cfd_Lx = gr.Slider(1.0, 5.0, value=2.0, label="Channel Length (m)")
                    cfd_Ly = gr.Slider(0.5, 2.0, value=1.0, label="Channel Height (m)")
                    
                    gr.Markdown("### Obstacle (set all to 0 for none)")
                    cfd_ox1 = gr.Slider(0.0, 5.0, value=0.5, label="Obstacle X1")
                    cfd_oy1 = gr.Slider(0.0, 2.0, value=0.4, label="Obstacle Y1")
                    cfd_ox2 = gr.Slider(0.0, 5.0, value=0