import pandas as pd
import shap
import matplotlib.pyplot as plt
import numpy as np
import joblib
import gradio as gr

# Load the model
try:
    loaded_model = joblib.load('machine_failure_prediction_model.joblib')
    print("Model loaded successfully for Gradio.")
except FileNotFoundError:
    print("Error: 'machine_failure_prediction_model.joblib' not found. Ensure the model file is in the correct directory.")
    loaded_model = None

def predict_failure_with_shap(Type, air_temp, process_temp, rotational_speed, torque, tool_wear):
    if loaded_model is None:
        return "Error: Model not loaded.", "", None

    input_data = pd.DataFrame({
        'Type': [Type],
        'Air temperature [K]': [air_temp],
        'Process temperature [K]': [process_temp],
        'Rotational speed [rpm]': [rotational_speed],
        'Torque [Nm]': [torque],
        'Tool wear [min]': [tool_wear]
    })

    predicted_failure = loaded_model.predict(input_data)[0]
    predicted_proba = loaded_model.predict_proba(input_data)[0]

    prediction_label = f"Predicted Failure: {'Failure' if predicted_failure == 1 else 'No Failure'}"
    probabilities_label = f"Probabilities (No Failure, Failure): {predicted_proba[0]:.4f}, {predicted_proba[1]:.4f}"

    preprocessor = loaded_model.named_steps['preprocessor']
    classifier = loaded_model.named_steps['classifier']
    X_transformed = preprocessor.transform(input_data)

    explainer = shap.TreeExplainer(classifier)
    shap_values = explainer.shap_values(X_transformed)
    feature_names = preprocessor.get_feature_names_out()

    if isinstance(shap_values, list):
        shap_val = shap_values[1][0]
        base_val = explainer.expected_value[1]
    else:
        shap_val = shap_values[0, :] if shap_values.ndim == 2 else shap_values[0, :, 1]
        base_val = explainer.expected_value if not isinstance(explainer.expected_value, (list, tuple, np.ndarray)) else explainer.expected_value[1]

    fig = plt.figure()
    shap.waterfall_plot(
        shap.Explanation(
            values=shap_val,
            base_values=base_val,
            data=X_transformed[0],
            feature_names=feature_names
        ), show=False
    )
    plt.title("SHAP Waterfall Plot for Failure Prediction")
    plt.tight_layout()

    return prediction_label, probabilities_label, fig


# Define Gradio interface OUTSIDE the function
iface = gr.Interface(
    fn=predict_failure_with_shap,
    inputs=[
        gr.Dropdown(choices=['L', 'M', 'H'], label='Machine Type'),
        gr.Number(label='Air temperature [K]', step=0.1),
        gr.Number(label='Process temperature [K]', step=0.1),
        gr.Number(label='Rotational speed [rpm]', step=1),
        gr.Number(label='Torque [Nm]', step=0.1),
        gr.Number(label='Tool wear [min]', step=1)
    ],
    outputs=[
        gr.Label(label='Predicted Failure (0=No Failure, 1=Failure)'),
        gr.Label(label='Prediction Probabilities'),
        gr.Plot(label='SHAP Waterfall Plot')
    ],
    title="Machine Failure Prediction with SHAP Analysis",
    description="Enter the machine parameters to predict failure and see the SHAP analysis.",
    flagging_mode='never'
)

if __name__ == "__main__":
    iface.launch(server_name="0.0.0.0", server_port=7860)