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, HistGradientBoostingClassifier
from sklearn.metrics import (classification_report, accuracy_score, confusion_matrix)
from sklearn.inspection import permutation_importance

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

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

for col in df.select_dtypes(include=['object']).columns:
    df[col] = LabelEncoder().fit_transform(df[col])

target_col = df.columns[-1]

counts = df[target_col].value_counts()
small_classes = counts[counts < 10].index # Only counts classes with 10+ samples
df = df[~df[target_col].isin(small_classes)]

normal_label = df[target_col].value_counts().idxmax()
df["binary_label"] = (df[target_col] != normal_label).astype(int)
x = df.drop([target_col, "binary_label"], axis=1)
y = df ["binary_label"]

# 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)

# ////////////////////////////////////////////////////////////////////////////////////////////////////////////////

# ///////////////////////////////////////////////// RandForest ///////////////////////////////////////////////////

rf = RandomForestClassifier(random_state=42)

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

rf_grid = GridSearchCV(rf, rf_params, cv=2, n_jobs=-1, scoring="f1_macro")
print("\nBeginning grid search (RandForest):") # Training takes a long time - this verifies the program is actually running
rf_grid.fit(X_train, y_train)
print("Done") # Verifies training is finished
rf_best = rf_grid.best_estimator_

# ////////////////////////////////////////////////////////////////////////////////////////////////////////////////

# ///////////////////////////////////////////////// GradBoost ////////////////////////////////////////////////////

gb = HistGradientBoostingClassifier(random_state=42)

gb_params = {"max_iter": [50, 100], "max_depth": [None, 10], "learning_rate": [0.05, 0.1]} # Parameters

gb_grid = GridSearchCV(gb, gb_params, cv=2, scoring="f1_macro")
print("\nBeginning grid search (GradBoost):") # Training takes a long time - this verifies the program is actually running
gb_grid.fit(X_train, y_train)
print("Done") # Verifies training is finished
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))

    cm = confusion_matrix(y_test, y_pred) # Creates confusion matrix

    sns.heatmap(cm, annot=True, fmt="d", cmap="Blues", xticklabels=["Negative", "Positive"], yticklabels=["Negative", "Positive"]) # Formatting for conf. matrix
    plt.title(f"{model_name} Confusion Matrix") # Title
    plt.xlabel("Predicted") # X Label
    plt.ylabel("Actual") # Y Label
    plt.show()

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

rf_importance = pd.Series(rf_best.feature_importances_, index=X_train.columns).sort_values(ascending=False) # Importance chart - RandForest
plt.figure(figsize = (10, 6))
rf_importance.head(10).plot(kind="bar") # Bar chart
plt.title("10 Most Important Features (Random Forest)")
plt.ylabel("Importance Score")
plt.xticks(rotation=45, ha="right") # Chart formatting
plt.tight_layout()
plt.show()

permutation_importance = permutation_importance(gb_best, X_test, y_test, n_repeats=5, random_state=42, scoring="f1") # Importance chart - GradBoost (Hist)
gb_importance = pd.Series(permutation_importance.importances_mean, index=X_train.columns).sort_values(ascending=False)
plt.figure(figsize = (10, 6))
gb_importance.head(10).plot(kind="bar") # Bar chart
plt.title("10 Most Important Features (Gradient Boosting)")
plt.ylabel("Importance Score")
plt.xticks(rotation=45, ha="right") # Chart formatting
plt.tight_layout()
plt.show()

Beginning grid search (RandForest):
Done

Beginning grid search (GradBoost):
Done

Random Forest
Accuracy: 0.9990471653168176

Classification Report:
               precision    recall  f1-score   support

           0       1.00      1.00      1.00     13469
           1       1.00      1.00      1.00     11719

    accuracy                           1.00     25188
   macro avg       1.00      1.00      1.00     25188
weighted avg       1.00      1.00      1.00     25188

Gradient Boosting
Accuracy: 0.9990471653168176

Classification Report:
               precision    recall  f1-score   support

           0       1.00      1.00      1.00     13469
           1       1.00      1.00      1.00     11719

    accuracy                           1.00     25188
   macro avg       1.00      1.00      1.00     25188
weighted avg       1.00      1.00      1.00     25188