import os
import time
import numpy as np
import pandas as pd
from sklearn.model_selection import KFold, ParameterGrid
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler, PolynomialFeatures, RobustScaler, PowerTransformer
from sklearn.impute import SimpleImputer
from sklearn.feature_selection import SelectKBest, f_regression
from sklearn.metrics import mean_squared_error
from sklearn.cluster import KMeans
import shap
from deap import base, creator, tools, algorithms
from algorithms import lasso, random_forest, gradient_boosting, decision_tree_regressor, ridge_regressor, stacking_lasso

import argparse

creator.create("FitnessMin", base.Fitness, weights=(-1.0, -1.0))  # Minimize both objectives
creator.create("Individual", list, fitness=creator.FitnessMin)

def load_dataset():

    # Read the CSV path from the temporary file
    with open("last_csv_path.txt", "r") as f:
        csv_path = f.read().strip()

    if not os.path.exists(csv_path):
        raise FileNotFoundError(f"CSV file not found: {csv_path}")

    # Load the dataset
    data = pd.read_csv(csv_path)
    return data

# Preprocess the data and separate features and target
def preprocess_data(input_data, pipeline, k_value, target_column):
    pipeline.named_steps['feature_selection'].set_params(k=k_value)
    X = input_data.drop(columns=[target_column], errors='ignore')
    print(X.columns)    
    y = input_data[target_column]
    X = pipeline.fit_transform(X, y)
    return X, y

# Function to compute SHAP values using KMeans clustering
def compute_shap_values_with_kmeans(model, X_train, X_test, n_clusters=10):
    kmeans = KMeans(n_clusters=n_clusters, n_init=10, random_state=42).fit(X_train)
    cluster_centers = kmeans.cluster_centers_

    explainer = shap.KernelExplainer(model.predict, cluster_centers)
    shap_values = explainer.shap_values(X_test)

    return shap_values

# Safe model training and evaluation
def safe_model_fit_evaluate(model, X_train, y_train, X_test, y_test):
    try:
        model.fit(X_train, y_train)
        predictions = model.predict(X_test)
        mse = mean_squared_error(y_test, predictions)
        return mse, predictions
    except (ValueError, OverflowError) as e:
        print(f"Model evaluation failed due to: {e}")
        return float('inf'), None

# Updated fitness function to evaluate MSE and SHAP stability using KMeans for SHAP
def evaluate_individual(individual, data, kf):
    model_name, model_params, pipeline_idx, k_value = individual[0], individual[1], individual[2], individual[3]

    # Get the model and pipeline
    model_func = model_space[model_name]
    model, params = model_func()
    model.set_params(**model_params)

    pipeline = pipelines[pipeline_idx]
    X_processed, y_processed = preprocess_data(data, pipeline, k_value)

    mse_scores = []
    shap_stabilities = []
    all_shap_values = []

    for fold, (train_idx, test_idx) in enumerate(kf.split(X_processed, y_processed)):
        X_train_fold, X_test_fold = X_processed[train_idx], X_processed[test_idx]
        y_train_fold, y_test_fold = y_processed[train_idx], y_processed[test_idx]

        # Safe model fitting and evaluation
        mse, predictions = safe_model_fit_evaluate(model, X_train_fold, y_train_fold, X_test_fold, y_test_fold)
        mse_scores.append(mse)

        if predictions is not None:
            # Compute SHAP values using KMeans clustering
            shap_values = compute_shap_values_with_kmeans(model, X_train_fold, X_test_fold)
            shap_stability = np.std(shap_values, axis=0).mean()
            shap_stabilities.append(shap_stability)
            all_shap_values.append(shap_values)

    # Return average MSE and SHAP stability
    return np.mean(mse_scores), np.mean(shap_stabilities), all_shap_values

# Run NSGA-II experiment with a time limit
def run_nsga_experiment(output_base_directory="./results_nsga",
                        population_size=30,
                        n_generations=50,
                        time_limit=36000):
    # Load the dataset
    data = load_dataset()

    # Pipelines for preprocessing
    global pipelines
    pipelines = [
        Pipeline([
            ('normalize', StandardScaler()),
            ('imputer', SimpleImputer(strategy='mean')),
            ('poly_features', PolynomialFeatures(degree=2)),
            ('feature_selection', SelectKBest(score_func=f_regression, k=25))
        ]),
        Pipeline([
            ('robust', RobustScaler()),
            ('feature_selection', SelectKBest(score_func=f_regression, k=25)),
            ('power_transformation', PowerTransformer(method='yeo-johnson'))
        ])
    ]

    # Model space with their corresponding hyperparameter grids
    global model_space
    model_space = {
        'lasso': lasso,
        'ridge': ridge_regressor,
        'random_forest': random_forest,
        'gradient_boosting': gradient_boosting,
        'decision_tree': decision_tree_regressor,
        'stacking_ensemble': stacking_lasso
    }

    # Create output directory with timestamp
    timestamp = time.strftime("%Y%m%d_%H%M%S")
    output_directory = os.path.join(output_base_directory, f"nsga_{time_limit // 3600}h_{timestamp}")
    os.makedirs(output_directory, exist_ok=True)

    # Define cross-validation strategy
    kf = KFold(n_splits=5, shuffle=True, random_state=42)

    # Convert parameter grid to a list of tuples for DEAP
    param_space = [(model_name, param_values, i, k_value) for model_name, model_func in model_space.items()
                   for param_values in list(ParameterGrid(model_func()[1])) for i in range(len(pipelines))
                   for k_value in range(10, 31, 5)]  # Range for k in SelectKBest

    # Initialize DEAP toolbox
    toolbox = base.Toolbox()
    toolbox.register("attr_float", lambda: np.random.choice(len(param_space), 4))  # Generate a list with 4 elements
    toolbox.register("individual", tools.initIterate, creator.Individual, toolbox.attr_float)
    toolbox.register("population", tools.initRepeat, list, toolbox.individual)
    toolbox.register("mate", tools.cxTwoPoint)
    toolbox.register("mutate", tools.mutPolynomialBounded, low=0, up=len(param_space) - 1, eta=0.5, indpb=0.2)
    toolbox.register("select", tools.selNSGA2)
    toolbox.register("evaluate", lambda ind: evaluate_individual(param_space[int(ind[0])], data, kf))

    # Run NSGA-II algorithm
    start_time = time.time()
    pop = toolbox.population(n=population_size)

    # Evaluate the initial population
    fits = toolbox.map(toolbox.evaluate, pop)
    for fit, ind in zip(fits, pop):
        ind.fitness.values = fit[:2]
        ind.shap_values = fit[2]

    # Use the number of parents and offspring in the evolution process
    for gen in range(n_generations):
        if time.time() - start_time > time_limit:
            print("Time limit exceeded, stopping evolution.")
            break

        offspring = algorithms.varAnd(pop, toolbox, cxpb=0.7, mutpb=0.2)
        fits = toolbox.map(toolbox.evaluate, offspring)

        for fit, ind in zip(fits, offspring):
            ind.fitness.values = fit[:2]  # Ensure only the first two values are assigned to fitness
            ind.shap_values = fit[2]

        # Select the next generation of parents from the combined pool of parents and offspring
        pop = toolbox.select(pop + offspring, k=population_size)

    for ind in pop:
        for fold_idx, shap_values in enumerate(ind.shap_values):
            shap_output_path = os.path.join(output_directory, f"shap_values_{int(ind[0])}_fold_{fold_idx + 1}.npy")
            np.save(shap_output_path, shap_values)

    # Extract and save the Pareto front
    pareto_front = tools.sortNondominated(pop, len(pop), first_front_only=True)[0]
    pareto_results_df = pd.DataFrame([param_space[int(ind[0])] + ind.fitness.values for ind in pareto_front],
                                     columns=['Model', 'Params', 'Pipeline', 'K', 'MSE', 'SHAP_Stability'])
    pareto_results_df.to_csv(os.path.join(output_directory, "nsga_pareto_front.csv"), index=False)

    # Save the final population
    results_df = pd.DataFrame([param_space[int(ind[0])] + ind.fitness.values for ind in pop],
                              columns=['Model', 'Params', 'Pipeline', 'K', 'MSE', 'SHAP_Stability'])
    results_df.to_csv(os.path.join(output_directory, "nsga_results.csv"), index=False)

    print(f"All results saved to {output_directory}")


if __name__ == "__main__":
    run_nsga_experiment(
        output_base_directory="./results_nsga",
        population_size=80,  # Larger population size for a comprehensive search
        n_generations=100,  # Increased number of generations
        num_parents=30,  # Increased number of parents
        num_offspring=50,  # Increased number of offspring
        time_limit=108000  # 20 hours (20 * 3600 seconds)
    )

