import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
import trimesh

# Function for Stress Analysis
def stress_analysis(force, die_width, die_height, material_strength):
    try:
        # Calculate parameters
        stress = force / (die_width * die_height)  # Stress = Force / Area
        safety_factor = material_strength / stress  # Safety Factor = Material Strength / Stress

        # Data for plotting
        x = ["Stress", "Material Strength", "Safety Factor"]
        y_stress = [stress, material_strength, safety_factor]

        # Create a combined graph
        fig, ax = plt.subplots()
        ax.plot(x, y_stress, marker='o', label="Simulation Parameters")
        ax.axhline(y=1, color='red', linestyle='--', label="Critical Safety Factor (1)")
        ax.set_ylabel("Value")
        ax.set_title("Stress Analysis Parameters")
        ax.legend()

        # Save the figure to a file
        fig_path = "/tmp/stress_analysis_graph.png"
        fig.savefig(fig_path)

        # Close the figure after saving it
        plt.close(fig)

        return fig_path  # Return the image path
    except Exception as e:
        return f"Error in stress analysis: {str(e)}"

# Function for 3D Visualization with Trimesh
def visualize_die(length, width, thickness):
    try:
        # Create a basic die shape using Trimesh
        box = trimesh.creation.box(extents=(length, width, thickness))
        
        # Save the 3D model as a file
        stl_path = "/tmp/progressive_die.stl"
        box.export(stl_path)

        # Convert the STL file to a PNG image
        screenshot = "/tmp/progressive_die_visualization.png"
        box.show()

        return screenshot  # Return the screenshot path
    except Exception as e:
        return f"Error visualizing die: {str(e)}"

# Combined Interface Function
def interface(force, die_width, die_height, material_strength, length, width, thickness):
    # Stress Analysis
    stress_graph_path = stress_analysis(force, die_width, die_height, material_strength)
    
    # 3D Visualization
    die_visualization_path = visualize_die(length, width, thickness)
    
    return stress_graph_path, die_visualization_path

# Gradio Interface
inputs = [
    gr.Slider(minimum=100, maximum=500, value=200, label="Material Yield Strength (MPa)"),
    gr.Slider(minimum=100, maximum=500, value=200, label="Die Width (mm)"),
    gr.Slider(minimum=100, maximum=500, value=200, label="Die Height (mm)"),
    gr.Slider(minimum=5000, maximum=20000, value=12000, label="Force (N)"),
    gr.Slider(minimum=100, maximum=500, value=200, label="Length (mm)"),
    gr.Slider(minimum=100, maximum=500, value=200, label="Width (mm)"),
    gr.Slider(minimum=10, maximum=50, value=10, label="Thickness (mm)"),
]

outputs = [
    gr.Image(type="filepath", label="Stress Analysis Graph"),
    gr.Image(type="filepath", label="3D Visualization"),
]

app = gr.Interface(fn=interface, inputs=inputs, outputs=outputs)
app.launch()