preprocessing_cicids.py

import os
import pandas as pd
import numpy as np
import json
from sklearn.utils import resample

RAW_DIR = "data/cicids"            # Folder with raw CSVs
PROCESSED_DIR = "data/processed"   # Where output will go
os.makedirs(PROCESSED_DIR, exist_ok=True)

# ===== CONFIG =====
PER_CLASS_MAX = 50_000   # Max rows per class after downsampling (tune as needed)

def preprocess():
    print("📂 Loading CICIDS raw data...")
    csv_files = [f for f in os.listdir(RAW_DIR) if f.endswith(".csv")]
    if not csv_files:
        print(f"❌ No CSV files found in '{RAW_DIR}'. Exiting.")
        return

    # --- 1. Global Label Mapping ---
    print("🗺️ Scanning all files to create a consistent label map...")
    all_labels = set()
    for file in csv_files:
        file_path = os.path.join(RAW_DIR, file)
        try:
            df_labels = pd.read_csv(file_path, usecols=[" Label"], engine="python", encoding="ISO-8859-1")
            df_labels.columns = [col.strip() for col in df_labels.columns]
            all_labels.update(df_labels["Label"].astype(str).str.strip().unique())
        except Exception as e:
            print(f"⚠️ Could not read labels from {file}: {e}")
            continue

    # Stable sorted mapping
    global_label_map = {label: idx for idx, label in enumerate(sorted(all_labels))}
    print(f"✅ Found {len(global_label_map)} unique labels across all files.")

    # --- 2. Load, Clean, Encode, Downsample ---
    all_dfs = []

    for file in csv_files:
        file_path = os.path.join(RAW_DIR, file)
        print(f"   -> Processing {file}...")

        try:
            # Try utf-8, fallback
            try:
                df = pd.read_csv(file_path, engine="python")
            except UnicodeDecodeError:
                df = pd.read_csv(file_path, engine="python", encoding="ISO-8859-1")

            # Strip cols
            df.columns = [col.strip() for col in df.columns]

            if "Label" not in df.columns:
                print(f"⚠️ Skipping {file}: no 'Label' column.")
                continue

            df["Label"] = df["Label"].astype(str).str.strip()
            df["LabelMapped"] = df["Label"].map(global_label_map)

            # Fill missing values
            for col in df.columns:
                if df[col].dtype == "object":
                    df[col] = df[col].fillna("missing").astype(str)
                else:
                    df[col] = df[col].fillna(0)

            # --- Downsampling ---
            dfs_down = []
            for label, group in df.groupby("LabelMapped"):
                if len(group) > PER_CLASS_MAX:
                    group = resample(group, replace=False, n_samples=PER_CLASS_MAX, random_state=42)
                dfs_down.append(group)
            df_balanced = pd.concat(dfs_down)

            # Save parquet per file
            out_file = os.path.join(PROCESSED_DIR, f"processed_{os.path.splitext(file)[0]}.parquet")
            df_balanced.to_parquet(out_file, index=False)
            print(f"✅ {file} → {out_file} (rows: {len(df_balanced)}, classes: {df_balanced['LabelMapped'].nunique()})")

            all_dfs.append(df_balanced)

        except Exception as e:
            print(f"⚠️ Skipping {file}: {e}")
            continue

    # --- 3. Combined Dataset ---
    if all_dfs:
        full_df = pd.concat(all_dfs, ignore_index=True)

        # Re-check object cols
        for col in full_df.select_dtypes(include=["object"]).columns:
            full_df[col] = full_df[col].fillna("missing").astype(str)

        combined_file = os.path.join(PROCESSED_DIR, "cicids_full.parquet")
        full_df.to_parquet(combined_file, index=False)
        print(f"✅ Combined dataset saved → {combined_file} ({len(full_df)} rows, {full_df['LabelMapped'].nunique()} classes)")

        # Save global label mapping
        with open(os.path.join(PROCESSED_DIR, "label_map.json"), "w") as f:
            json.dump(global_label_map, f, indent=2)
        print("📝 Label mapping saved → label_map.json")

    print("📌 Preprocessing complete.")


if __name__ == "__main__":
    preprocess()