import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.preprocessing import LabelEncoder, StandardScaler
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.metrics import (classification_report, accuracy_score, confusion_matrix)

# Load Training Set file
file_path = 'kdd_train.csv'  # Finds training data
df = pd.read_csv(file_path)

# Handle missing values
df.replace(" ", np.nan, inplace=True)
df.dropna(inplace=True)

# Encode categorical features
label_encoders = {}
for col in df.select_dtypes(include=['object']).columns:
    le = LabelEncoder()
    df[col] = le.fit_transform(df[col])
    label_encoders[col] = le

#Features & target
x = df.drop("label", axis=1)
y = df["label"]

# Binary labels
df["label"] = df["label"].apply(lambda x: 0 if x == 0 else 1)

# Train-test split
X_train, X_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=42, stratify=y)

# Feature scaling
scaler = StandardScaler()
x_train_scaled = scaler.fit_transform(X_train)
x_test_scaled = scaler.transform(X_test)


# For RandForest
rf = RandomForestClassifier(random_state=42)

rf_params = {"n_estimators": [100, 200], "max_depth": [None, 10, 20], "min_samples_split": [2, 5]}

rf_grid = GridSearchCV(rf, rf_params, cv=3, n_jobs=-1, scoring="f1_macro")
rf_grid.fit(x_train, y_train)

rf_best = rf_grid.best_estimator_


# For GradBoost
gb = GradientBoostingClassifier(random_state=42)

gb_params = {
    "n_estimators": [100, 200], "learning_rate": [0.05, 0.1], "max_depth": [3, 5]
}

gb_grid = GridSearchCV(gb, gb_params, cv=3, scoring="f1_macro")
gb_grid.fit(x_train_scaled, y_train)
gb_best = gb_grid.best_estimator_

# Evaluating Model
def evaluate_model(model, X_test, y_test, model_name):
    y_pred = model.predict(X_test)

    print(f"\n{model_name}")
    print("Accuracy:", accuracy_score(y_test, y_pred))
    print("\nClassification Report:\n", classification_report(y_test, y_pred))

    # Create Confusion Matrix
    cm = confusion_matrix(y_test, y_pred)
    tn, fp, fn, tp = cm.ravel()

    print("\nConfusion Matrix:")
    print(f"True Positives: {tp}")
    print(f"False Positives: {fp}")
    print(f"True Negatives: {tn}")
    print(f"False Negatives: {fn}")

    labels = np.array([[f"TN\n{tn}", f"FP\n{fp}", f"FN\n{fn}", f"TP\n{tp}"]])

    sns.heatmap(cm, cmap="Blues", fmt="d")
    plt.title(f"{model_name} Confusion Matrix")
    plt.xlabel("Predicted")
    plt.ylabel("Actual")
    plt.show()

evaluate_model(rf_best, X_test, y_test, "Random Forest")
evaluate_model(gb_best, X_test_scaled, y_test, "Gradient Boosting")

rf_importance = pd.Series(rf_best.feature_importances_, index=X_train.columns).sort_values(ascending=False)

plt.figure(figsize = (10, 6))
rf_importance.head(10).plot(kind="bar")
plt.title("10 Most Important Features (Random Forest)")
plt.ylabel("Importance Score")
plt.show()