import os
import shap
import numpy as np
import pandas as pd
import h2o
from h2o.automl import H2OAutoML
from sklearn.preprocessing import StandardScaler, RobustScaler, PolynomialFeatures, PowerTransformer
from sklearn.impute import SimpleImputer
from sklearn.pipeline import Pipeline
from sklearn.model_selection import KFold
from sklearn.feature_selection import SelectKBest, f_regression
from sklearn.metrics import mean_squared_error
from sklearn.cluster import KMeans

from combine_datasets import combine_datasets

# Initialize H2O
h2o.init()


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

    # Convert cluster centers to H2OFrame
    cluster_centers_h2o = h2o.H2OFrame(cluster_centers)

    # Convert X_test to H2OFrame
    X_test_h2o = h2o.H2OFrame(X_test)

    # Compute SHAP values on cluster centers
    explainer = shap.KernelExplainer(lambda x: model.predict(h2o.H2OFrame(x)).as_data_frame().values.flatten(),
                                     cluster_centers_h2o.as_data_frame().values)
    shap_values = explainer.shap_values(X_test_h2o.as_data_frame().values)

    return shap_values


# Function to load and combine datasets
def load_and_combine_datasets(use_ndvi=True):
    data = combine_datasets(
        "./data/potatoes_dataset.csv",
        "./data/updated_soil_data_with_awc.csv",
        "./data/final_augmented_climate_data.csv",
        "./data/NDVI.csv",
        ndvi=use_ndvi
    )
    return data


# Main experiment function with H2O AutoML
def run_h2o_experiment(use_ndvi=True, automl_time=10, output_base_directory="./results_h2o"):
    data = load_and_combine_datasets(use_ndvi=use_ndvi)
    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'))
        ])
    ]

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

    metrics_df = pd.DataFrame(columns=['Model', 'Pipeline', 'MSE', 'SHAP_Stability'])

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

    for pipeline_idx, pipeline in enumerate(pipelines):
        X_processed, y = preprocess_data(data, pipeline)

        X_h2o = h2o.H2OFrame(X_processed)
        y_h2o = h2o.H2OFrame(y.to_frame())
        X_h2o['target'] = y_h2o

        shap_stabilities = []
        all_shap_values = []

        for fold, (train_idx, test_idx) in enumerate(kf.split(X_processed)):
            train = X_h2o[train_idx.tolist(), :]
            test = X_h2o[test_idx.tolist(), :]

            aml = H2OAutoML(max_runtime_secs=int(automl_time * 3600 / kf.get_n_splits()), max_models=5, seed=42,
                            sort_metric='MSE')
            aml.train(x=list(X_h2o.columns[:-1]), y='target', training_frame=train)

            best_model = aml.leader

            leaderboard = aml.leaderboard.as_data_frame()
            leaderboard_csv_path = os.path.join(output_directory, f"leaderboard_pipeline_{pipeline_idx + 1}_fold_{fold + 1}.csv")
            leaderboard.to_csv(leaderboard_csv_path, index=False)

            model_output_dir = os.path.join(output_directory, f"pipeline_{pipeline_idx + 1}")
            os.makedirs(model_output_dir, exist_ok=True)
            model_file_path = os.path.join(model_output_dir, f"h2o_best_model_{pipeline_idx + 1}_fold_{fold + 1}.zip")
            h2o.save_model(best_model, path=model_file_path)

            predictions = best_model.predict(test).as_data_frame()['predict']

            mse = mean_squared_error(test['target'].as_data_frame().values, predictions.values)

            shap_values = compute_shap_values_with_kmeans(best_model, train.as_data_frame().values,
                                                          test.as_data_frame().values)
            shap_stability = np.std(shap_values, axis=0).mean()
            shap_stabilities.append(shap_stability)
            all_shap_values.append(shap_values)


            print(f"Completed H2O AutoML and SHAP computation for Pipeline {pipeline_idx + 1}, Fold {fold + 1}")

        mean_shap_stability = np.mean(shap_stabilities)
        for fold_idx, shap_values in enumerate(all_shap_values):
            shap_file_name = f"shap_values_pipeline_{pipeline_idx + 1}_fold_{fold_idx + 1}.npy"
            shap_file_path = os.path.join(model_output_dir, shap_file_name)
            np.save(shap_file_path, shap_values)

        metrics_df = pd.concat([metrics_df, pd.DataFrame([{
            'Model': best_model.model_id,
            'Pipeline': f"Pipeline_{pipeline_idx + 1}",
            'MSE': mse,
            'SHAP_Stability': mean_shap_stability
        }])], ignore_index=True)

    metrics_csv_path = os.path.join(output_directory, "metrics_summary.csv")
    metrics_df.to_csv(metrics_csv_path, index=False)

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

    h2o.shutdown(prompt=False)


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description='Run H2O AutoML for different models.')
    parser.add_argument('--time', type=float, required=True, help='Time in hours for the AutoML process to run.')
    args = parser.parse_args()

    # Run the H2O AutoML experiment
    run_h2o_experiment(use_ndvi=True, automl_time=args.time, output_base_directory="./results_h2o")

