import gradio as gr
import pandas as pd
import numpy as np
import joblib
from xgboost import XGBClassifier

# Load các mô hình và encoder
model = joblib.load("xgboost_model.pkl")
scaler = joblib.load("scaler.pkl")
label_encoder = joblib.load("label_encoder.pkl")

def predict_attack(file):
    # Đọc file CSV
    df = pd.read_csv(file.name)

    # Tiền xử lý dữ liệu
    X = scaler.transform(df)

    # Dự đoán
    predictions = model.predict(X)
    prediction_probs = model.predict_proba(X)
    confidences = np.max(prediction_probs, axis=1)
    predicted_labels = label_encoder.inverse_transform(predictions)

    # Thêm cột dự đoán và độ tin cậy
    df["Predicted Label"] = predicted_labels
    df["Confidence"] = confidences.round(4)

    # Tóm tắt kết quả
    summary = df["Predicted Label"].value_counts().reset_index()
    summary.columns = ["Class", "Count"]
    summary["Percentage"] = (summary["Count"] / len(df) * 100).round(2).astype(str) + "%"
    summary["Avg Confidence"] = df.groupby("Predicted Label")["Confidence"].mean().round(4).values

    # Confusion matrix giả định: không có nhãn thật → dùng cột dự đoán chính
    labels = label_encoder.classes_
    confusion_matrix_html = "<table border='1'><tr><th>Actual \\ Predicted</th>" + "".join(
        f"<th>{label}</th>" for label in labels
    ) + "</tr>"

    for actual in labels:
        confusion_matrix_html += f"<tr><td>{actual}</td>"
        for predicted in labels:
            count = len(df[(df["Predicted Label"] == predicted) & (df["Predicted Label"] == actual)])
            cell_color = " style='background-color: #ffcccc'" if count > 0 and actual != predicted else ""
            confusion_matrix_html += f"<td{cell_color}>{count}</td>"
        confusion_matrix_html += "</tr>"
    confusion_matrix_html += "</table>"

    # Hiển thị kết quả dạng bảng HTML
    summary_html = summary.to_html(index=False)
    prediction_sample = df.head(10).to_html(index=False)

    html_output = f"""
    <h3>Kết quả phân tích:</h3>
    <h4>Summary of Detections:</h4>
    {summary_html}
    <h4>Actual vs Predicted Labels Confusion Matrix:</h4>
    {confusion_matrix_html}
    <h4>Prediction Results (Total: {len(df)} records)</h4>
    <p>Top 10 records shown:</p>
    {prediction_sample}
    """

    return html_output

# Tạo giao diện
iface = gr.Interface(
    fn=predict_attack,
    inputs=gr.File(label="Tải file CSV đã xử lý"),
    outputs=gr.HTML(),
    title="🔍 Dự đoán tấn công mạng bằng mô hình XGBoost",
    description="Tải file CSV để hệ thống phân tích và dự đoán loại tấn công (nếu có).",
)

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