import gradio as gr
from design import generate_die, visualize_die
from analysis import stress_analysis
from optimization import optimize_tool

# Create Gradio App
with gr.Blocks() as app:
    gr.Markdown("## Press Tool AI Suite")
    gr.Markdown("Select a tool below to get started:")

    with gr.Tabs():
        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="STEP File Location")
            visualization_output = gr.Image(label="3D Visualization")
            die_button = gr.Button("Generate Die")
            die_button.click(
                lambda l, w, t: (generate_die(l, w, t), visualize_die(l, w, t)),
                inputs=[length, width, thickness],
                outputs=[die_output, visualization_output],
            )

        with gr.Tab("Stress Analysis"):
            gr.Markdown("### Select Simulation Tool and 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)
            temperature_change = gr.Number(label="Temperature Change (°C)", value=50)
            alpha = gr.Number(label="Thermal Expansion Coefficient (1/°C)", value=1e-5)
            elastic_modulus = gr.Number(label="Elastic Modulus (MPa)", value=200000)
            fatigue_strength = gr.Number(label="Fatigue Strength (MPa)", value=150)

            safety_factor_output = gr.Textbox(label="Safety Factor")
            stress_chart = gr.Plot()
            stress_button = gr.Button("Analyze Stress")
            stress_button.click(
                lambda f, dw, dh, ms, tc, a, em, fs: stress_analysis(f, dw, dh, ms, tc, a, em, fs),
                inputs=[force, die_width, die_height, material_strength, temperature_change, alpha, elastic_modulus, fatigue_strength],
                outputs=[safety_factor_output, stress_chart],
            )

        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(
                lambda s, fr, dc, m: optimize_tool(s, fr, dc, m),
                inputs=[speed, feed_rate, depth_of_cut, material],
                outputs=optimization_results,
            )

# Launch the app
app.launch()