# Import libraries
import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import gradio as gr
import joblib

# Step 1: Generate or Load Sample Data
np.random.seed(42)
data = pd.DataFrame({
    "Pressure": np.random.randint(50, 200, 200),
    "Temperature": np.random.randint(300, 700, 200),
    "Material_Type": np.random.randint(1, 5, 200),
    "Defect_Type": np.random.choice([0, 1, 2], 200, p=[0.6, 0.3, 0.1])
})

# Splitting data
X = data[["Pressure", "Temperature", "Material_Type"]]
y = data["Defect_Type"]
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# Step 2: Train the Model
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train, y_train)

# Save the model
joblib.dump(model, "defect_model.pkl")

# Step 3: Test the Model
y_pred = model.predict(X_test)
print(f"Model Accuracy: {accuracy_score(y_test, y_pred):.2f}")

# Step 4: Create a Gradio Interface
def predict_defect(pressure, temperature, material_type):
    # Load the model
    loaded_model = joblib.load("defect_model.pkl")
    # Predict
    input_data = np.array([[pressure, temperature, material_type]])
    prediction = loaded_model.predict(input_data)[0]
    defect_map = {0: "No Defect", 1: "Crack", 2: "Wrinkle"}
    return defect_map[prediction]

# Gradio Interface
interface = gr.Interface(
    fn=predict_defect,
    inputs=[
        gr.Slider(minimum=50, maximum=200, step=1, label="Pressure"),
        gr.Slider(minimum=300, maximum=700, step=1, label="Temperature"),
        gr.Dropdown(choices=["1", "2", "3", "4"], label="Material Type")
    ],
    outputs=gr.Textbox(label="Predicted Defect"),
    title="Defect Prediction Model",
    description="Predicts defect type based on input features: Pressure, Temperature, and Material Type."
)

# Launch the app
if __name__ == "__main__":
    interface.launch()