# app.py
import gradio as gr
from ai_models import load_models
from ansys_simulation import run_ansys_simulation  # Import this function only once
from visualization import generate_graph
import os

# Load AI models
stress_model, deformation_model, pass_fail_model = load_models()

def ai_and_simulation_workflow(thickness, hole_diameter, force):
    # AI Pre-Screening
    try:
        pre_screen = pass_fail_model.predict([[thickness, hole_diameter, force]])[0]
    except Exception as e:
        return f"Error with AI model prediction: {str(e)}", None

    if pre_screen == 0:
        return "AI predicts failure. Adjust parameters.", None
     
    # ANSYS Simulation
    try:
        max_stress, max_deformation = run_ansys_simulation(thickness, hole_diameter, force)
    except Exception as e:
        return f"Error with ANSYS simulation: {str(e)}", None

    # Generate Visualization
    try:
        graph_path = generate_graph(thickness, hole_diameter, force, max_stress, max_deformation)
        if not os.path.exists(graph_path):
            return f"Error: Graph path '{graph_path}' does not exist.", None
    except Exception as e:
        return f"Error with graph generation: {str(e)}", None

    return f"Pass: Stress={max_stress} Pa, Deformation={max_deformation} mm", graph_path

# Gradio UI
interface = gr.Interface(
    fn=ai_and_simulation_workflow,
    inputs=[
        gr.Slider(10, 50, step=1, label="Thickness (mm)"),
        gr.Slider(5, 25, step=1, label="Hole Diameter (mm)"),
        gr.Slider(1000, 15000, step=500, label="Force (N)")
    ],
    outputs=[
        gr.Textbox(label="Simulation Status"),
        gr.Image(label="Dynamic Results Visualization")
    ],
    title="AI-Driven ANSYS Design Validator",
    description="Dynamic tool for optimizing and validating press tool designs.",
    theme="default",
    live=True
)

if __name__ == "__main__":
    interface.launch()