import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler
import joblib
import gradio as gr

# Load model IDS
model = joblib.load("model_cb2.pkl")

# Label mapping
LABEL_CLASS = {
    0: "Benign",
    1: "Bot",
    2: "DDOS attack-HOIC",
    3: "DDOS attack-LOIC-UDP",
    4: "DDoS attacks-LOIC-HTTP",        
    5: "DoS attacks-GoldenEye",
    6: "DoS attacks-Hulk",
    7: "DoS attacks-Slowloris",
    8: "Infilteration",
    9: "SSH-Bruteforce"
}

def preprocess_data(df_all):
    # Hapus kolom tidak penting
    df_all.drop(columns=['Unnamed: 0', 'Timestamp'], inplace=True, errors='ignore')

    # Ganti inf dengan NaN
    df_all['Flow Byts/s'].replace([np.inf, -np.inf], np.nan, inplace=True)
    df_all['Flow Pkts/s'].replace([np.inf, -np.inf], np.nan, inplace=True)

    # Isi NaN
    df_all['Flow Byts/s'].fillna(df_all['Flow Byts/s'].median(), inplace=True)
    df_all['Flow Pkts/s'].fillna(df_all['Fwd Pkts/s'] + df_all['Bwd Pkts/s'], inplace=True)

    # Normalisasi
    df_all = df_all[['Fwd Seg Size Min', 'Bwd IAT Tot', 'Bwd IAT Max', 'Bwd IAT Std',
       'Bwd IAT Mean', 'PSH Flag Cnt', 'Bwd Pkt Len Min', 'Flow IAT Std',
       'Flow IAT Max', 'Fwd IAT Std', 'Fwd IAT Max', 'Idle Max', 'Idle Mean',
       'Bwd IAT Min', 'Bwd Pkt Len Max', 'Idle Min', 'Bwd Pkts/s',
       'Flow Pkts/s', 'Bwd Seg Size Avg', 'Bwd Pkt Len Mean', 'Fwd PSH Flags',
       'SYN Flag Cnt', 'Pkt Len Max', 'Fwd Pkt Len Min', 'Pkt Len Min',
       'ACK Flag Cnt', 'Init Bwd Win Byts', 'Fwd Header Len',
       'TotLen Fwd Pkts', 'Subflow Fwd Byts', 'Subflow Fwd Pkts',
       'Tot Fwd Pkts', 'Fwd Act Data Pkts', 'Dst Port', 'ECE Flag Cnt',
       'RST Flag Cnt']]
    # scaler = StandardScaler()
    # df_all[df_all.columns] = scaler.fit_transform(df_all[df_all.columns])
    
    return df_all

def predict_ids(file):
    try:
        df_all = pd.read_csv(file.name)
        df_processed = preprocess_data(df_all)
        pred_indices = model.predict(df_processed)

        # Buat DataFrame hasil
        hasil = pd.DataFrame({
            "Index": range(len(pred_indices)),
            "Predicted Label": [LABEL_CLASS[i] for i in pred_indices.flatten()]
        })

        return hasil
    except Exception as e:
        return f"Terjadi error saat memproses file: {str(e)}"

# Gradio UI
demo = gr.Interface(
    fn=predict_ids,
    inputs=gr.File(label="Upload CSV File", file_types=[".csv"]),
    outputs=gr.Dataframe(label="Hasil Prediksi (Index dan Label)"),
    title="Predicting IDS",
    description="Upload file CSV untuk memprediksi jenis aktivitas jaringan (Benign, BruteForce, dll)"
)

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