import numpy as np
import pandas as pd
from collections import Counter
import xgboost as xgb
import logging
import os
from datetime import datetime
import argparse
import joblib
import json
import zarr

# Dask imports
import dask.array as da
from dask.distributed import Client
from dask_ml.impute import SimpleImputer as DaskSimpleImputer
from dask_ml.preprocessing import StandardScaler as DaskStandardScaler

# Import metrics from scikit-learn
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score, classification_report

def load_data_for_years(years, region, data_path):
    """
    Load features and lccs_class data for the specified years using Dask arrays.

    :param years: list of int, Years to load data for
    :param region: str, Region code (ES, EU, UK, AU, GLOBE)
    :param data_path: str, Base path where data is stored
    :return: tuple of dask.array, Concatenated features and lccs_class for the specified years
    """
    features_list = []
    lccs_class_list = []

    for year in years:
        features_year_path = os.path.join(data_path, f'features_{year}_{region}.zarr')
        lccs_class_year_path = os.path.join(data_path, f'lccs_class_{year}_{region}.zarr')

        features_list.append(da.from_zarr(features_year_path))
        lccs_class_list.append(da.from_zarr(lccs_class_year_path))

    features = da.concatenate(features_list, axis=0)
    lccs_class = da.concatenate(lccs_class_list, axis=0)

    return features, lccs_class

def train_evaluate_ml_model(model, X_train, X_test, y_train, y_test, model_name):
    model.fit(X_train, y_train)
    y_pred = model.predict(X_test)

    # Collect predictions and true labels
    y_pred_np = y_pred.compute()
    y_test_np = y_test.compute()

    # Compute performance metrics using scikit-learn
    accuracy = accuracy_score(y_test_np, y_pred_np)
    f1 = f1_score(y_test_np, y_pred_np, average='weighted')
    precision = precision_score(y_test_np, y_pred_np, average='weighted', zero_division=0)
    recall = recall_score(y_test_np, y_pred_np, average='weighted')

    # Log the performance metrics
    logging.info(f'Accuracy of {model_name}: {accuracy*100:.2f}%')
    logging.info(f'F1 Score of {model_name}: {f1:.2f}')
    logging.info(f'Precision of {model_name}: {precision:.2f}')
    logging.info(f'Recall of {model_name}: {recall:.2f}')

    # Log detailed classification report
    class_report = classification_report(y_test_np, y_pred_np, target_names=[f'Class {i}' for i in range(len(unique_classes))], zero_division=0)
    logging.info(f'Classification Report for {model_name}:\n{class_report}')

    # Save the model
    if model_name == "XGBoost":
        model.save_model(os.path.join(save_data_path, f'{model_name}_{start_year}_{end_year}_{region}_{timestamp}.json'))
    else:
        joblib.dump(model, os.path.join(save_data_path, f'{model_name}_{start_year}_{end_year}_{region}_{timestamp}.joblib'))
    logging.info(f'{model_name} model saved successfully.')

if __name__ == '__main__':
    # Parse command-line arguments
    parser = argparse.ArgumentParser(description="Train a model using land use data.")
    parser.add_argument("--data_path", type=str, required=True, help="Path to the directory where data is stored")
    parser.add_argument("--start_year", type=int, required=True, help="Start year for the data range")
    parser.add_argument("--end_year", type=int, required=True, help="End year for the data range")
    parser.add_argument("--region", type=str, required=True, help="Region for the data slice (ES, EU, UK, AU, GLOBE)")
    parser.add_argument("--model_type", type=str, required=True, choices=['xgboost', 'logistic_regression'], help="Model type to use")
    parser.add_argument("--imputation", type=str, choices=['mean', 'median', 'most_frequent'], default='mean', help="Imputation method for handling NaN values")
    parser.add_argument("--train_years", nargs='+', type=int, required=True, help="Years to use for training")
    parser.add_argument("--test_years", nargs='+', type=int, required=True, help="Years to use for testing")
    parser.add_argument("--grid_search", action='store_true', help="Enable grid search for hyperparameter tuning")

    args = parser.parse_args()
    start_year = args.start_year
    end_year = args.end_year
    region = args.region
    model_type = args.model_type
    imputation_method = args.imputation
    train_years = args.train_years
    test_years = args.test_years
    save_data_path = args.data_path
    grid_search_enabled = args.grid_search

    # Configure logging
    current_path = os.getcwd()
    logs_path = os.path.join(current_path, "Logs")
    os.makedirs(logs_path, exist_ok=True)
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    log_file_name = f"train_log_{start_year}_{end_year}_{region}_{timestamp}.log"
    log_file_path = os.path.join(logs_path, log_file_name)
    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', handlers=[
        logging.FileHandler(log_file_path),
        logging.StreamHandler()
    ])
    logging.info(f"Training model of type: {model_type}")

    # Set paths for saving/loading processed data and metadata
    metadata_path = os.path.join(save_data_path, f'metadata_{start_year}_{end_year}_{region}.json')

    # Load metadata
    with open(metadata_path, 'r') as f:
        metadata = json.load(f)

    # Log the loaded metadata
    logging.info(f"Loaded metadata: {metadata}")

    # Start Dask client
    client = Client()
    logging.info("Dask client started.")

    # Load data for training and testing years
    X_train, y_train = load_data_for_years(train_years, region, save_data_path)
    X_test, y_test = load_data_for_years(test_years, region, save_data_path)

    # Log the shapes of the arrays
    logging.info(f"X_train shape: {X_train.shape}")
    logging.info(f"X_test shape: {X_test.shape}")

    # Rechunk the arrays before imputation
    # Ensure that there's only one chunk along the features axis (axis=1)
    X_train = X_train.rechunk({1: -1})
    X_test = X_test.rechunk({1: -1})

    # Impute missing values
    if imputation_method in ['mean', 'median', 'most_frequent']:
        imputer = DaskSimpleImputer(strategy=imputation_method)
    else:
        raise ValueError("KNN imputation is not supported with Dask. Please choose from 'mean', 'median', 'most_frequent'.")

    X_train = imputer.fit_transform(X_train)
    X_test = imputer.transform(X_test)
    logging.info(f"Imputed missing values using {imputation_method} strategy.")

    # Adjust labels to ensure they start from 0 and go up to n_classes-1
    unique_classes = da.unique(y_train).compute()
    unique_classes = sorted(unique_classes)
    class_to_index = {float(cls): idx for idx, cls in enumerate(unique_classes)}  # Convert keys to float

    # Save class_to_index to a JSON file
    class_to_index_path = os.path.join(save_data_path, f'class_to_index_{start_year}_{end_year}_{region}_{timestamp}.json')
    with open(class_to_index_path, 'w') as f:
        json.dump(class_to_index, f)

    logging.info("Class-to-index mapping saved successfully.")

    # Map labels to indices
    def map_labels(y, class_to_index):
        return y.map_blocks(lambda x: np.vectorize(class_to_index.get)(x.astype(float)), dtype=int)

    y_train = map_labels(y_train, class_to_index)
    y_test = map_labels(y_test, class_to_index)

    # Rechunk y_train and y_test to match X_train and X_test
    y_train = y_train.rechunk({0: X_train.chunksize[0]})
    y_test = y_test.rechunk({0: X_test.chunksize[0]})

    # Verify chunk sizes
    logging.info(f"X_train chunks: {X_train.chunks}")
    logging.info(f"y_train chunks: {y_train.chunks}")

    # Initialize and train the model based on the chosen architecture
    if model_type == "xgboost":
        # Create DaskDMatrix
        dtrain = xgb.dask.DaskDMatrix(client, X_train, y_train)
        dtest = xgb.dask.DaskDMatrix(client, X_test, y_test)

        params = {
            'objective': 'multi:softprob',
            'num_class': len(unique_classes),
            'eval_metric': 'mlogloss',
            'tree_method': 'hist',
            'random_state': 42,
            # Removed n_estimators from params
        }

        # Set the number of boosting rounds (equivalent to n_estimators)
        num_boost_round = 2000

        if grid_search_enabled:
            logging.info("Starting grid search for XGBoost...")

            # Define parameter grid with 16 combinations
            param_grid = {
                'max_depth': [9, 10],
                'learning_rate': [0.035, 0.04],
                'reg_alpha': [0.4, 0.5],
                'subsample': [0.8],
                'reg_lambda': [1.0],
                'colsample_bytree': [0.8]
            }

            # Manual grid search
            from itertools import product
            param_combinations = list(product(
                param_grid['max_depth'],
                param_grid['learning_rate'],
                param_grid['reg_alpha']
            ))

            best_score = -np.inf
            best_params = None
            best_model = None

            for max_depth, learning_rate, reg_alpha in param_combinations:
                params.update({
                    'max_depth': max_depth,
                    'learning_rate': learning_rate,
                    'reg_alpha': reg_alpha,
                    'subsample': 0.8,
                    'reg_lambda': 1.0,
                    'colsample_bytree': 0.8
                })
                logging.info(f"Training with parameters: {params}")
                output = xgb.dask.train(
                    client,
                    params,
                    dtrain,
                    num_boost_round=num_boost_round,  # Correctly pass num_boost_round
                    evals=[(dtrain, 'train')],
                    verbose_eval=False
                )
                model = output['booster']
                y_pred = xgb.dask.predict(client, model, dtest)
                y_pred_np = y_pred.compute()
                y_test_np = y_test.compute()
                f1 = f1_score(y_test_np, y_pred_np, average='weighted', zero_division=0)
                logging.info(f"F1 Score with current parameters: {f1:.4f}")
                if f1 > best_score:
                    best_score = f1
                    best_params = params.copy()
                    best_model = model
            logging.info(f'Best parameters found: {best_params}')
            model = best_model

        else:
            # Use your previous hyperparameters if grid search is not enabled
            params.update({
                'max_depth': 10,
                'learning_rate': 0.04,
                'subsample': 0.8,
                'reg_alpha': 0.5,
                'reg_lambda': 1.0,
                'colsample_bytree': 0.8
            })
            output = xgb.dask.train(
                client,
                params,
                dtrain,
                num_boost_round=num_boost_round,  # Correctly pass num_boost_round
                evals=[(dtrain, 'train')],
                verbose_eval=False
            )
            model = output['booster']

        # Evaluate the model
        y_pred = xgb.dask.predict(client, model, dtest)
        y_pred = y_pred.persist()

        # Collect predictions and true labels
        y_pred_np = y_pred.compute()
        y_test_np = y_test.compute()

        # Compute performance metrics using scikit-learn
        accuracy = accuracy_score(y_test_np, y_pred_np)
        f1 = f1_score(y_test_np, y_pred_np, average='weighted', zero_division=0)
        precision = precision_score(y_test_np, y_pred_np, average='weighted', zero_division=0)
        recall = recall_score(y_test_np, y_pred_np, average='weighted', zero_division=0)

        # Log the performance metrics
        logging.info(f'Accuracy of XGBoost: {accuracy*100:.2f}%')
        logging.info(f'F1 Score of XGBoost: {f1:.2f}')
        logging.info(f'Precision of XGBoost: {precision:.2f}')
        logging.info(f'Recall of XGBoost: {recall:.2f}')

        # Log detailed classification report
        class_report = classification_report(
            y_test_np, y_pred_np,
            target_names=[f'Class {i}' for i in range(len(unique_classes))],
            zero_division=0
        )
        logging.info(f'Classification Report for XGBoost:\n{class_report}')

        # Save the model
        model_path = os.path.join(save_data_path, f'XGBoost_{start_year}_{end_year}_{region}_{timestamp}.json')
        model.save_model(model_path)
        logging.info(f'XGBoost model saved successfully at {model_path}.')

    elif model_type == "logistic_regression":
        # Initialize model with default parameters
        from dask_ml.linear_model import LogisticRegression as DaskLogisticRegression

        model = DaskLogisticRegression(
            solver='admm',
            random_state=42,
            max_iter=2000
        )

        if grid_search_enabled:
            logging.info("Starting grid search for Logistic Regression...")
            # Define parameter grid
            param_grid = {
                'C': [0.1, 1.0],
                'penalty': ['l1', 'l2'],
                'max_iter': [1000, 2000]
            }

            # Initialize GridSearchCV from Dask-ML
            from dask_ml.model_selection import GridSearchCV as DaskGridSearchCV

            grid_search = DaskGridSearchCV(
                estimator=model,
                param_grid=param_grid,
                scoring='accuracy',  # Changed to 'accuracy' since 'f1_weighted' may not be supported
                cv=3,
                n_jobs=-1,
            )

            grid_search.fit(X_train, y_train)
            # Update model with best parameters
            model = grid_search.best_estimator_
            logging.info(f'Best parameters found: {grid_search.best_params_}')

        else:
            # Use default parameters if grid search is not enabled
            model = DaskLogisticRegression(
                solver='admm',
                random_state=42,
                max_iter=2000
            )

        train_evaluate_ml_model(model, X_train, X_test, y_train, y_test, "LogisticRegression")

    else:
        raise ValueError("Invalid model type specified. Choose from 'xgboost', 'logistic_regression'.")

    logging.info("Training and evaluation complete.")
