import pandas as pd
import numpy as np

from sklearn.model_selection import train_test_split, GridSearchCV, cross_val_score
from sklearn.preprocessing import StandardScaler
from sklearn.feature_selection import SelectKBest, mutual_info_classif

from sklearn.linear_model import LogisticRegression
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.svm import SVC
from sklearn.neural_network import MLPClassifier

from sklearn.metrics import (
    accuracy_score, precision_score, recall_score, f1_score,
    confusion_matrix, classification_report, roc_auc_score, roc_curve
)

from imblearn.over_sampling import SMOTE

import matplotlib.pyplot as plt
import seaborn as sns

"""**LOAD DATA**"""

df = pd.read_csv("/content/feature_vectors_syscallsbinders_frequency_5_Cat.csv")
feature_info = pd.read_csv("/content/dataset-features-categories.csv")

print("Dataset shape:", df.shape)
print(df.head())

"""**DATA CLEANING**"""

df = df.drop_duplicates()
df = df.fillna(0)

"""**DEFINE FEATURES & LABEL**"""

X = df.drop("Class", axis=1)
y = df["Class"]

print("\nClass distribution:")
print(y.value_counts())

"""**TRAIN-TEST SPLIT**"""

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.3, random_state=42, stratify=y
)

"""**FEATURE SCALING**"""

scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

"""**FEATURE SELECTION**"""

selector = SelectKBest(score_func=mutual_info_classif, k=100)

selector.fit(X_train_scaled, y_train)

X_train_sel = selector.transform(X_train_scaled)
X_test_sel = selector.transform(X_test_scaled)

selected_features = X.columns[selector.get_support()]
print("\nSelected features:", len(selected_features))

"""**HANDLE CLASS IMBALANCE (SMOTE)**"""

sm = SMOTE(random_state=42)
X_train_res, y_train_res = sm.fit_resample(X_train_sel, y_train)

"""**MODEL DEFINITIONS**"""

models = {
    "Logistic Regression": LogisticRegression(max_iter=1000),
    "Decision Tree": DecisionTreeClassifier(),
    "Random Forest": RandomForestClassifier(),
    "SVM": SVC(kernel='linear', probability=True),
    "Neural Network": MLPClassifier(hidden_layer_sizes=(64, 32), max_iter=200, early_stopping=True)
}

results = {}

"""**TRAIN & EVALUATE MODELS**"""

for name, model in models.items():
    print(f"\n🔹 Training: {name}")

    model.fit(X_train_res, y_train_res)

    y_pred = model.predict(X_test_sel)
    # For multiclass ROC AUC, predict_proba returns probabilities for each class
    y_prob = model.predict_proba(X_test_sel)

    acc = accuracy_score(y_test, y_pred)
    # Specify average='weighted' for multiclass classification metrics
    prec = precision_score(y_test, y_pred, average='weighted')
    rec = recall_score(y_test, y_pred, average='weighted')
    f1 = f1_score(y_test, y_pred, average='weighted')
    # For multiclass ROC AUC, specify multi_class and average parameters
    auc = roc_auc_score(y_test, y_prob, multi_class='ovr', average='weighted')

    results[name] = [acc, prec, rec, f1, auc]

    print(classification_report(y_test, y_pred))

"""**RESULTS TABLE**"""

results_df = pd.DataFrame(results, index=["Accuracy", "Precision", "Recall", "F1", "AUC"]).T

print("\nModel Comparison:")
print(results_df)

"""**BEST MODEL**"""

best_model_name = results_df["F1"].idxmax()
best_model = models[best_model_name]

print("\nBest Model:", best_model_name)

"""**CONFUSION MATRIX**"""

y_pred = best_model.predict(X_test_sel)

cm = confusion_matrix(y_test, y_pred)

plt.figure(figsize=(6,4))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
plt.title(f"Confusion Matrix - {best_model_name}")
plt.xlabel("Predicted")
plt.ylabel("Actual")
plt.show()

# False positive / negative analysis
fp = cm[0][1]
fn = cm[1][0]

print("False Positives:", fp)
print("False Negatives:", fn)

"""**ROC CURVE**"""

y_prob_full = best_model.predict_proba(X_test_sel)

plt.figure(figsize=(8, 6))

classes = np.unique(y_test)

for i, class_label in enumerate(classes):
    y_test_binarized = (y_test == class_label).astype(int)
    class_probabilities = y_prob_full[:, i]

    fpr, tpr, _ = roc_curve(y_test_binarized, class_probabilities)
    auc_score = roc_auc_score(y_test_binarized, class_probabilities)

    plt.plot(fpr, tpr, label=f"Class {class_label} (AUC={auc_score:.2f})")

plt.plot([0,1],[0,1],'k--')
plt.xlabel("False Positive Rate")
plt.ylabel("True Positive Rate")
plt.title(f"ROC Curve - {best_model_name} (One-vs-Rest)")
plt.legend()
plt.show()

"""**FEATURE IMPORTANCE + CTI INTERPRETATION**"""

if best_model_name == "Random Forest":
    importances = best_model.feature_importances_

    feat_imp = pd.DataFrame({
        "feature": selected_features,
        "importance": importances
    })

    top_features = feat_imp.sort_values("importance", ascending=False).head(15)

    print("\nTop Threat Indicators (IoCs):")
    print(top_features)

    top_features.set_index("feature").plot(kind='barh')
    plt.title("Top Features (CTI Indicators)")
    plt.show()

    merged = top_features.merge(feature_info, left_on="feature", right_on="transact", how="left")
    print("\nFeature: CTI Category Mapping:")
    print(merged)

"""**CROSS-VALIDATION**"""

cv_scores = cross_val_score(best_model, X_train_res, y_train_res, cv=10)

print("\nCross-validation Accuracy:", cv_scores.mean())

"""**MODEL COMPARISON PLOT**"""

results_df.plot(kind='bar', figsize=(10,6))
plt.title("Model Performance Comparison")
plt.xticks(rotation=45)
plt.show()