import os
import gradio as gr
import joblib
import numpy as np
from sklearn.datasets import load_iris

# Tải thông tin nhãn lớp Iris
iris = load_iris()
target_names = iris.target_names

# Kiểm tra sự tồn tại của mô hình
model_path = "model.pkl"
if os.path.exists(model_path):
    model = joblib.load(model_path)
else:
    model = None


def predict(sepal_length, sepal_width, petal_length, petal_width):
    if model is None:
        return "Không tìm thấy tệp model.pkl. Vui lòng huấn luyện mô hình trước!", {}

    # Tạo mảng đặc trưng cho đầu vào
    features = np.array(
        [[sepal_length, sepal_width, petal_length, petal_width]]
    )

    # Dự đoán lớp và tính toán xác suất
    prediction = model.predict(features)[0]
    probabilities = model.predict_proba(features)[0]

    class_name = target_names[prediction].upper()

    # Tạo từ điển chứa xác suất của từng loài hoa
    prob_dict = {
        target_names[i].capitalize(): float(probabilities[i]) for i in range(3)
    }

    return class_name, prob_dict


# Xây dựng giao diện Gradio
demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Slider(
            minimum=float(iris.data[:, 0].min()),
            maximum=float(iris.data[:, 0].max()),
            value=5.1,
            label="Sepal Length (cm)",
        ),
        gr.Slider(
            minimum=float(iris.data[:, 1].min()),
            maximum=float(iris.data[:, 1].max()),
            value=3.5,
            label="Sepal Width (cm)",
        ),
        gr.Slider(
            minimum=float(iris.data[:, 2].min()),
            maximum=float(iris.data[:, 2].max()),
            value=1.4,
            label="Petal Length (cm)",
        ),
        gr.Slider(
            minimum=float(iris.data[:, 3].min()),
            maximum=float(iris.data[:, 3].max()),
            value=0.2,
            label="Petal Width (cm)",
        ),
    ],
    outputs=[
        gr.Textbox(label="Loài hoa dự đoán (Predicted Species)"),
        gr.Label(num_top_classes=3, label="Xác suất chi tiết (Probabilities)"),
    ],
    title="🌸 Iris Species Classification Predictor",
    description="Ứng dụng ML nhỏ dự đoán loài hoa Iris, được huấn luyện và triển khai tự động qua quy trình CI/CD GitHub Actions lên Gradio SDK của Hugging Face Spaces.",
    theme="soft",
)

# Chỉ chạy launch khi gọi trực tiếp file này
if __name__ == "__main__":
    demo.launch()