import gradio as gr
import joblib

# Load the saved models
rf_model = joblib.load('rf_model.pkl')
dt_model = joblib.load('dt_model.pkl')
bagging_model = joblib.load('bagging_model.pkl')
ada_model = joblib.load('ada_model.pkl')

class_labels = {
    0: "normal",
    1: "backdoor",
    2: "ddos",
    3: "dos",
    4: "injection",
    5: "password",
    6: "ransomware",
    7: "scanning",
    8: "xss",
    9: "mitm"
}

def detect_intrusion(features, model_choice="Random Forest"):
    # Convert the input string (comma-separated values) into a list of floats
    features = [list(map(float, features.split(",")))]
    
    # Choose the model based on user selection
    if model_choice == "Random Forest":
        model = rf_model
    elif model_choice == "Decision Tree":
        model = dt_model
    elif model_choice == "Bagging Classifier":
        model = bagging_model
    elif model_choice == "AdaBoost Classifier":
        model = ada_model
    else:
        return "Invalid model choice!"
    
    # Predict the class (multi-class classification)
    prediction = model.predict(features)
    predicted_class = prediction[0]  # Get the predicted class (an integer between 0-8)
    
    # Return the human-readable class description
    if predicted_class == 0:
        return "No Intrusion Detected"
    else:
        return f"Intrusion Detected: {class_labels.get(predicted_class, 'Unknown Attack')}"

# Create a Gradio interface
iface = gr.Interface(fn=detect_intrusion, 
                     inputs=[gr.Textbox(label="Input Features (comma-separated)"), 
                             gr.Dropdown(choices=["Random Forest", "Decision Tree", "Bagging Classifier", "AdaBoost Classifier"], label="Select Model")], 
                     outputs="text", 
                     title="Intrusion Detection System",
                     description="Enter features in the format: feature1, feature2, feature3...")

# Launch the interface locally for testing
iface.launch()