import gradio as gr
import xgboost as xgb
import numpy as np
import joblib
import os


# Load your model (adjust path as needed)
def load_model():
    if os.path.exists("best_model.json"):
        model = xgb.Booster()
        model.load_model("best_model.json")
        print("✅ Model loaded using XGBoost's native method.")
        return model
    elif os.path.exists("best_model.pkl"):
        model = joblib.load("best_model.pkl")
        print("✅ Model loaded using Joblib.")
        return model
    else:
        print("❌ No model file found.")
        return None


model = load_model()


# Prediction function
def predict_employee_status(satisfaction_level, last_evaluation, number_project,
                            average_monthly_hours, time_spend_company,
                            work_accident, promotion_last_5years, salary):
    input_data = np.array([[satisfaction_level, last_evaluation, number_project,
                            average_monthly_hours, time_spend_company,
                            work_accident, promotion_last_5years, salary]])

    if model is None:
        return "❌ No model found. Please upload the model file."

    if isinstance(model, xgb.Booster):
        dmatrix = xgb.DMatrix(input_data)
        prediction = model.predict(dmatrix)[0]
        result = "✅ The employee is likely to Quit." if prediction > 0.5 else "✅ The employee is likely to Stay."
    else:
        prediction = model.predict(input_data)[0]
        result = "✅ The employee is likely to Quit." if prediction == 1 else "✅ The employee is likely to Stay."

    return result


# Gradio interface
interface = gr.Interface(
    fn=predict_employee_status,
    inputs=[
        gr.inputs.Number(label="Satisfaction Level"),
        gr.inputs.Number(label="Last Evaluation"),
        gr.inputs.Number(label="Number of Projects"),
        gr.inputs.Number(label="Average Monthly Hours"),
        gr.inputs.Number(label="Time Spent at Company (Years)"),
        gr.inputs.Number(label="Work Accident (0 = No, 1 = Yes)"),
        gr.inputs.Number(label="Promotion in Last 5 Years (0 = No, 1 = Yes)"),
        gr.inputs.Dropdown(choices=[0, 1, 2], label="Salary Level (0 = Low, 1 = Medium, 2 = High)")
    ],
    outputs="text",
    title="Employee Retention Prediction",
    description="Predict whether an employee will stay or quit based on their profile.",
    live=False
)

# Launch Gradio app
interface.launch()