import os
import numpy as np
import matplotlib.pyplot as plt
import joblib
import gradio as gr

# Visualization Function
def plot_simulation(pressure, temperature, material_type):
    if material_type <= 0 or temperature <= 0:
        raise ValueError("Temperature and Material Type must be positive.")
    stress = np.linspace(pressure * 0.5, pressure * 1.5, 100)
    strain = stress / (temperature * 0.01 * material_type)
    plt.figure(figsize=(6, 4))
    plt.plot(stress, strain, label="Stress-Strain Curve")
    plt.xlabel("Stress (MPa)")
    plt.ylabel("Strain (%)")
    plt.title("Stress-Strain Simulation")
    plt.legend()
    plt.grid(True)
    output_path = "simulation_plot.png"
    plt.savefig(output_path)
    plt.close()
    return output_path

# Prediction Function
def predict_and_visualize(pressure, temperature, material_type):
    model_path = "defect_model.pkl"
    if not os.path.exists(model_path):
        return "Error: Model file not found.", None
    try:
        model = joblib.load(model_path)
        prediction = model.predict([[pressure, temperature, material_type]])[0]
        defect_map = {0: "No Defect", 1: "Defect"}
        defect_prediction = defect_map.get(prediction, "Unknown")
    except Exception as e:
        return f"Prediction Error: {str(e)}", None
    
    try:
        plot_path = plot_simulation(pressure, temperature, material_type)
    except Exception as e:
        return defect_prediction, f"Visualization Error: {str(e)}"
    
    return defect_prediction, plot_path

# Gradio Interface
interface = gr.Interface(
    fn=predict_and_visualize,
    inputs=[
        gr.Slider(50, 200, step=1, label="Pressure"),
        gr.Slider(300, 700, step=1, label="Temperature"),
        gr.Dropdown(choices=["1", "2", "3", "4"], label="Material Type")
    ],
    outputs=[
        gr.Textbox(label="Prediction"),
        gr.Image(label="Stress-Strain Visualization")
    ],
    title="Defect Prediction & Simulation",
    description="Predict defect type and visualize stress-strain simulation."
)

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