import os
import time
import numpy as np
import pandas as pd
import warnings
import shap
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
from sklearn.exceptions import ConvergenceWarning
from joblib import dump

from combine_datasets import combine_datasets
from algorithms import lasso, random_forest, gradient_boosting, decision_tree_regressor, ridge_regressor, stacking_lasso

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


# Preprocess the data and separate features and target
def preprocess_data(input_data, pipeline):
    X = input_data.drop(columns=['yield_t/ha'], errors='ignore')
    y = input_data['yield_t/ha']

    # Fit-transform the pipeline steps without `y` first
    for name, step in pipeline.steps:
        if name != 'feature_selection':
            X = step.fit_transform(X)

    # Apply `SelectKBest` separately, which requires `y`
    if 'feature_selection' in [name for name, _ in pipeline.steps]:
        X = pipeline.named_steps['feature_selection'].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):
    # Fit KMeans to the training data with explicit n_init parameter
    kmeans = KMeans(n_clusters=n_clusters, n_init=10, random_state=42).fit(X_train)
    cluster_centers = kmeans.cluster_centers_

    # Compute SHAP values on cluster centers
    explainer = shap.KernelExplainer(model.predict, cluster_centers)
    shap_values = explainer.shap_values(X_test)

    return shap_values


# Iterative grid search with time monitoring
def iterative_grid_search(model, param_grid, X, y, cv, time_limit):
    best_model = None
    best_score = -np.inf
    start_time = time.time()

    # Manually iterate over parameter combinations
    for params in ParameterGrid(param_grid):
        elapsed_time = time.time() - start_time
        if elapsed_time > time_limit:
            print("Time limit exceeded. Stopping search.")
            break

        model.set_params(**params)
        scores = []

        for train_idx, test_idx in cv.split(X, y):
            X_train, X_test = X[train_idx], X[test_idx]
            y_train, y_test = y[train_idx], y[test_idx]
            model.fit(X_train, y_train)
            predictions = model.predict(X_test)
            score = -mean_squared_error(y_test, predictions)
            scores.append(score)

        mean_score = np.mean(scores)
        if mean_score > best_score:
            best_score = mean_score
            best_model = model

    return best_model


# Main experiment function
def run_experiment_with_dynamic_grid(total_hours=10.0, output_base_directory="./results"):
    time_per_algorithm = total_hours * 3600 / 5

    # Load and combine datasets
    data = combine_datasets(
        "./data/potatoes_dataset.csv",
        "./data/updated_soil_data_with_awc.csv",
        "./data/final_augmented_climate_data.csv",
        "./data/NDVI.csv",
        ndvi=True
    )

    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'))
        ])
    ]

    models = {
        'lasso': lasso,
        'ridge': ridge_regressor,
        'random_forest': random_forest,
        'gradient_boosting': gradient_boosting,
        'decision_tree': decision_tree_regressor,
        'stacking_ensemble': stacking_lasso
    }

    output_directory = os.path.join(output_base_directory, f"grid_search_results_{total_hours}h")
    os.makedirs(output_directory, exist_ok=True)

    kf = KFold(n_splits=5, shuffle=True, random_state=42)

    for model_name, model_func in models.items():
        best_pipeline_result = None

        for pipeline_idx, pipeline in enumerate(pipelines):
            X, y = preprocess_data(data, pipeline)
            model, params = model_func()

            best_model = iterative_grid_search(
                model, params, X, y, cv=kf, time_limit=time_per_algorithm
            )

            shap_stabilities = []
            for fold, (train_idx, test_idx) in enumerate(kf.split(X, y)):
                X_train_fold, X_test_fold = X[train_idx], X[test_idx]
                y_train_fold = y[train_idx]

                best_model.fit(X_train_fold, y_train_fold)
                shap_values = compute_shap_values_with_kmeans(best_model, X_train_fold, X_test_fold)
                fold_shap_stability = np.std(shap_values, axis=0).mean()
                shap_stabilities.append(fold_shap_stability)

            shap_stability = np.mean(shap_stabilities)
            predictions = best_model.predict(X)
            mse = mean_squared_error(y, predictions)

            if best_pipeline_result is None or (
                    mse < best_pipeline_result['mse'] and shap_stability < best_pipeline_result['shap_stability']):
                best_pipeline_result = {
                    'model': best_model,
                    'pipeline_idx': pipeline_idx,
                    'mse': mse,
                    'shap_stability': shap_stability
                }

        if best_pipeline_result:
            model_output_dir = os.path.join(output_directory, model_name)
            os.makedirs(model_output_dir, exist_ok=True)

            model_file_path = os.path.join(model_output_dir, f"{model_name}_best_model.joblib")
            dump(best_pipeline_result['model'], model_file_path)

            metrics_df = pd.DataFrame({
                'Model': [model_name],
                'Pipeline': [f"Pipeline_{best_pipeline_result['pipeline_idx'] + 1}"],
                'MSE': [best_pipeline_result['mse']],
                'SHAP_Stability': [best_pipeline_result['shap_stability']]
            })
            metrics_csv_path = os.path.join(output_directory, "metrics_summary_GS.csv")
            metrics_df.to_csv(metrics_csv_path, index=False, mode='a', header=not os.path.exists(metrics_csv_path))

            print(f"Saved best results for {model_name} from Pipeline {best_pipeline_result['pipeline_idx'] + 1}")

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

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description='Run Grid Search for different models.')
    parser.add_argument('--time', type=float, required=True, help='Time for the model to run in hours')
    args = parser.parse_args()
    run_experiment_with_dynamic_grid(total_hours=args.time, output_base_directory="./results_grid_search")
