import streamlit as st
import pandas as pd
import numpy as np
import joblib
from xgboost import XGBClassifier

# Load all models and preprocessors
model = joblib.load("best_xgboost_model.pkl")
label_encoder = joblib.load("label_encoder.pkl")
scaler = joblib.load("standard_scaler.pkl")
var_thresh = joblib.load("variance_threshold.pkl")
mi_features = joblib.load("selected_mi_features.pkl")
roa_features = joblib.load("selected_roa_features.pkl")
class_names = joblib.load("class_names.pkl")

FEATURE_COLUMNS = [
    "flow_duration", "header_length", "protocol_type", "duration", "rate", "srate", "drate",
    "fin_flag_number", "syn_flag_number", "rst_flag_number", "psh_flag_number", "ack_flag_number",
    "ece_flag_number", "cwr_flag_number", "ack_count", "syn_count", "fin_count", "urg_count",
    "rst_count", "http", "https", "dns", "telnet", "smtp", "ssh", "irc", "tcp", "udp", "dhcp",
    "arp", "icmp", "ipv", "llc", "tot_sum", "min", "max", "avg", "std", "tot_size", "iat",
    "number", "magnitue", "radius", "covariance", "variance", "weight"
]

# Streamlit UI
st.title("🔐 Detection of DDoS Attack in IoT Networks")
st.markdown("Paste a single row of IoT traffic data for DDoS detection (comma **or** tab-separated).")

default_row = (
    "0,54,6,64,0.329807153,0.329807153,0,1,0,1,0,0,0,0,1,0,1,0,0,"
    "0,0,0,0,0,0,0,1,0,0,0,0,1,1,567,54,54,54,0,54,83343831.92,"
    "9.5,10.39230485,0,0,0,141.55"
)

user_input = st.text_area("Paste the feature row here (label column optional)", default_row, height=150)

if st.button("🔍 Predict Attack Type"):
    try:
        delimiter = '\t' if '\t' in user_input else ','

        input_parts = user_input.strip().split(delimiter)
        
        # Handle optional label at the end
        try:
            _ = float(input_parts[-1])  # If this fails, it's a string label
        except ValueError:
            input_parts = input_parts[:-1]  # Drop the label

        input_list = [float(x.strip()) for x in input_parts]

        if len(input_list) != len(FEATURE_COLUMNS):
            st.error(f"🚫 Expected {len(FEATURE_COLUMNS)} features but got {len(input_list)}.")
        else:
            input_df = pd.DataFrame([input_list], columns=FEATURE_COLUMNS)

            # Preprocessing pipeline
            scaled = scaler.transform(input_df)
            var_filtered = var_thresh.transform(scaled)
            mi_selected = pd.DataFrame(var_filtered, columns=np.array(FEATURE_COLUMNS)[var_thresh.get_support()])
            final_features = mi_selected[mi_features].iloc[:, roa_features]

            prediction = model.predict(final_features)
            predicted_label = label_encoder.inverse_transform(prediction)[0]

            st.success(f"✅ Predicted Attack Type: **{predicted_label}**")

    except Exception as e:
        st.error(f"🚫 Error: {e}")