import numpy as np
from os import path
from sklearn.preprocessing import MinMaxScaler
from keras import backend as K
from tensorflow import keras
import warnings
from dataPre import loadCsv, dataset_pre
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.ensemble import AdaBoostClassifier, RandomForestClassifier
from sklearn.neighbors import KNeighborsClassifier
from sklearn.tree import DecisionTreeClassifier
from sklearn.neural_network import MLPClassifier
from sklearn.linear_model import SGDClassifier
from sklearn.naive_bayes import GaussianNB
from sklearn.metrics import precision_score, recall_score, f1_score, matthews_corrcoef, accuracy_score
from deap import base, creator, tools, algorithms
import random

warnings.filterwarnings("ignore", category=Warning)

TL = 4

train_paths = [
    "data/UNSW_NB15_Train201.csv",
    "data/UNSW_NB15_Train202.csv",
    "data/UNSW_NB15_Train203.csv",
    "data/UNSW_NB15_Train204.csv",
    "data/UNSW_NB15_Train205.csv"
]

test_path = 'data/UNSW_NB15_TestBin.csv'

# Load data
train_data = [loadCsv(path) for path in train_paths]
test_data = loadCsv(test_path)

# Preprocess data
scaler = MinMaxScaler()
train_data_scaled = [scaler.fit_transform(data[:, :196]) for data in train_data]
test_data_scaled = scaler.fit_transform(test_data[:, :196])

x_train = [dataset_pre(data, TL) for data in train_data_scaled]
x_train = [np.reshape(data, (-1, 196)) for data in x_train]  # Reshape without TL for non-federated learning
x_test = dataset_pre(test_data_scaled, TL)
x_test = np.reshape(x_test, (-1, 196))

# Labels
y_train = [data[:, 196] for data in train_data]
y_train = np.concatenate(y_train, axis=0)  # Concatenate all y_train for non-federated learning
y_test = test_data[:, 196]

# Combine all training data
X_train = np.concatenate(x_train, axis=0)

# Define the models to be evaluated
models = {
    "AdaBoostM1": AdaBoostClassifier(n_estimators=50, random_state=42),
    "KNN": KNeighborsClassifier(n_neighbors=5),
    "J48": DecisionTreeClassifier(random_state=42),  # J48 is equivalent to DecisionTree in sklearn
    "MLP": MLPClassifier(hidden_layer_sizes=(100,), max_iter=300, random_state=42),
    "SGD": SGDClassifier(max_iter=1000, tol=1e-3, random_state=42),
    "NaiveBayes": GaussianNB(),
    "LMT": RandomForestClassifier(random_state=42)  # Logistic Model Tree not directly available in sklearn, using RF
}

# Function to evaluate the model
def evaluate_model(model, X_test, y_test):
    y_pred = model.predict(X_test)
    precision = precision_score(y_test, y_pred, zero_division=1)
    recall = recall_score(y_test, y_pred, zero_division=1)
    f1 = f1_score(y_test, y_pred, zero_division=1)
    mcc = matthews_corrcoef(y_test, y_pred)
    accuracy = accuracy_score(y_test, y_pred)

    return precision, recall, f1, mcc, accuracy

# Train and evaluate each model
results = {}

for name, model in models.items():
    model.fit(X_train, y_train)
    results[name] = evaluate_model(model, x_test, y_test)

# Print the evaluation results
for name, metrics in results.items():
    print(f"Model: {name}")
    print(f"Precision: {metrics[0]:.4f}")
    print(f"Recall: {metrics[1]:.4f}")
    print(f"F1-Score: {metrics[2]:.4f}")
    print(f"MCC: {metrics[3]:.4f}")
    print(f"Accuracy: {metrics[4]:.4f}")
    print("-" * 30)

# Further optimization with Genetic Algorithm for the best model
best_model_name = max(results, key=lambda name: results[name][-1])
best_model = models[best_model_name]

# Define the evaluation function for the Genetic Algorithm
def eval_forest(individual):
    n_trees = individual[0]
    max_depth = individual[1]
    
    # Create and train the RandomForest model
    model = RandomForestClassifier(n_estimators=n_trees, max_depth=max_depth, random_state=42)
    model.fit(X_train, y_train)
    
    # Predict on the test set
    y_pred = model.predict(x_test)
    
    # Calculate accuracy
    accuracy = accuracy_score(y_test, y_pred)
    return (accuracy,)

# Define the parameters for the Genetic Algorithm
creator.create("FitnessMax", base.Fitness, weights=(1.0,))
creator.create("Individual", list, fitness=creator.FitnessMax)

toolbox = base.Toolbox()
toolbox.register("attr_int", random.randint, 10, 200)  # Number of trees
toolbox.register("attr_depth", random.randint, 1, 20)  # Maximum depth of trees
toolbox.register("individual", tools.initCycle, creator.Individual, (toolbox.attr_int, toolbox.attr_depth), n=1)
toolbox.register("population", tools.initRepeat, list, toolbox.individual)
toolbox.register("mate", tools.cxBlend, alpha=0.5)
toolbox.register("mutate", tools.mutGaussian, mu=0, sigma=1, indpb=0.2)
toolbox.register("select", tools.selTournament, tournsize=3)
toolbox.register("evaluate", eval_forest)

# Set up the Genetic Algorithm
population = toolbox.population(n=20)
ngen = 10
cxpb = 0.5
mutpb = 0.2

# Run the Genetic Algorithm
algorithms.eaSimple(population, toolbox, cxpb=cxpb, mutpb=mutpb, ngen=ngen, verbose=True)

# Select the best individual
best_individual = tools.selBest(population, k=1)[0]
print('Best individual:', best_individual)

# Train the final RandomForest model with the best parameters
final_model = RandomForestClassifier(n_estimators=best_individual[0], max_depth=best_individual[1], random_state=42)
final_model.fit(X_train, y_train)

# Evaluate the final model
final_metrics = evaluate_model(final_model, x_test, y_test)
print(f"Final RandomForest Model Evaluation")
print(f"Precision: {final_metrics[0]:.4f}")
print(f"Recall: {final_metrics[1]:.4f}")
print(f"F1-Score: {final_metrics[2]:.4f}")
print(f"MCC: {final_metrics[3]:.4f}")
print(f"Accuracy: {final_metrics[4]:.4f}")
