import cadquery as cq
import numpy as np
import matplotlib.pyplot as plt
import gradio as gr
import os

# Ensure there's a directory for saving files
if not os.path.exists("generated_files"):
    os.makedirs("generated_files")

# Function for Progressive Die Design
def generate_die(length, width, thickness):
    try:
        # Generate the die design
        plate = cq.Workplane("XY").box(length, width, thickness)
        punch = cq.Workplane("XY").rect(10, 10).extrude(5).translate((length/4, width/4, thickness/2))
        die = plate.cut(punch)
        
        # Save the die to a file
        filename = "generated_files/progressive_die.step"
        cq.exporters.export(die, filename)
        
        # Return success message and file path for download
        return f"Progressive die design saved. [Download the file](/{filename})"
    except Exception as e:
        return f"Error generating die: {str(e)}"

# Function for Stress Analysis
def stress_analysis(force, die_width, die_height, material_strength):
    try:
        # Validate inputs to avoid invalid calculations
        if force <= 0 or die_width <= 0 or die_height <= 0 or material_strength <= 0:
            return "All values must be positive."

        # Calculate stress and safety factor
        stress = force / (die_width * die_height)
        safety_factor = material_strength / stress
        
        # Create stress analysis plot
        fig, ax = plt.subplots()
        ax.bar(["Stress", "Material Strength"], [stress, material_strength])
        ax.set_ylabel("Stress (MPa)")
        ax.set_title("Stress Analysis")
        plt.tight_layout()  # Ensure the plot is displayed correctly
        plt.close(fig)  # Close the plot so it doesn't render twice
        
        # Return the results and plot
        return f"Safety Factor: {round(safety_factor, 2)}", fig
    except Exception as e:
        return f"Error in stress analysis: {str(e)}", None

# Function for Tool Optimization
def optimize_tool(speed, feed_rate, depth_of_cut, material):
    try:
        # Adjust tool life estimate based on material
        material_factor = {"Steel": 1.0, "Aluminum": 0.8, "Titanium": 1.2}
        tool_life = 1000 / (speed * feed_rate * depth_of_cut * material_factor.get(material, 1.0))
        
        recommended_speed = 0.8 * speed
        recommended_feed_rate = 0.9 * feed_rate
        
        # Return optimization results
        return {
            "Estimated Tool Life (hrs)": round(tool_life, 2),
            "Recommended Speed (m/min)": round(recommended_speed, 2),
            "Recommended Feed Rate (mm/rev)": round(recommended_feed_rate, 2)
        }
    except Exception as e:
        return {"Error": str(e)}

# Gradio interface functions
def progressive_die_interface(length, width, thickness):
    return generate_die(length, width, thickness)

def stress_analysis_interface(force, width, height, material_strength):
    return stress_analysis(force, width, height, material_strength)

def tool_optimization_interface(speed, feed_rate, depth_of_cut, material):
    return optimize_tool(speed, feed_rate, depth_of_cut, material)

# Create the Gradio Interface
with gr.Blocks() as app:
    gr.Markdown("## Press Tool AI Suite")
    gr.Markdown("Select a tool below to get started:")
    
    with gr.Tabs():
        # Tab for Progressive Die Design
        with gr.Tab("Progressive Die Design"):
            gr.Markdown("### Enter Dimensions for Progressive Die")
            length = gr.Number(label="Length (mm)", value=100)
            width = gr.Number(label="Width (mm)", value=50)
            thickness = gr.Number(label="Thickness (mm)", value=10)
            die_output = gr.Textbox(label="Output")
            die_button = gr.Button("Generate Die")
            die_button.click(progressive_die_interface, inputs=[length, width, thickness], outputs=die_output)
        
        # Tab for Stress Analysis
        with gr.Tab("Stress Analysis"):
            gr.Markdown("### Enter Parameters for Stress Analysis")
            force = gr.Number(label="Force (N)", value=10000)
            die_width = gr.Number(label="Width (m)", value=0.05)
            die_height = gr.Number(label="Height (m)", value=0.01)
            material_strength = gr.Number(label="Material Strength (MPa)", value=250)
            safety_factor_output = gr.Textbox(label="Safety Factor")
            stress_chart = gr.Plot()
            stress_button = gr.Button("Analyze Stress")
            stress_button.click(stress_analysis_interface, inputs=[force, die_width, die_height, material_strength], outputs=[safety_factor_output, stress_chart])
        
        # Tab for Tool Optimization
        with gr.Tab("Tool Optimization"):
            gr.Markdown("### Enter Machining Parameters for Tool Optimization")
            speed = gr.Number(label="Cutting Speed (m/min)", value=100)
            feed_rate = gr.Number(label="Feed Rate (mm/rev)", value=0.2)
            depth_of_cut = gr.Number(label="Depth of Cut (mm)", value=1.0)
            material = gr.Dropdown(choices=["Steel", "Aluminum", "Titanium"], label="Material", value="Steel")
            optimization_results = gr.JSON(label="Optimization Results")
            optimize_button = gr.Button("Optimize Tool")
            optimize_button.click(tool_optimization_interface, inputs=[speed, feed_rate, depth_of_cut, material], outputs=optimization_results)

# Launch the app
app.launch()