import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import xgboost as xgb  # Adding XGBoost
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score, accuracy_score, f1_score, precision_score, recall_score, classification_report
import logging
import os
from datetime import datetime
import argparse
import joblib
import json
from sklearn.impute import SimpleImputer, KNNImputer
import zarr
import wandb
import time
import psutil  # For hardware details
from collections import Counter
from sklearn.utils.class_weight import compute_class_weight
from sklearn.model_selection import train_test_split
from scipy.ndimage import gaussian_filter
from sklearn.preprocessing import MinMaxScaler, StandardScaler
import gc
import sys

## all plotting routines 
def plot_data(data, title, filename, vmin=None, vmax=None):
    """
    Plot data as an image and save to a file.
    
    Parameters:
        data (np.ndarray): Data to be plotted.
        title (str): Title of the plot.
        filename (str): File path to save the plot.
        vmin (float, optional): Minimum data value for colormap scaling. Defaults to None.
        vmax (float, optional): Maximum data value for colormap scaling. Defaults to None.
    """
    plt.figure(figsize=(10, 6))
    # Pass vmin and vmax to imshow; if they are None, matplotlib will auto-scale.
    im = plt.imshow(np.flipud(data), cmap='viridis', origin='upper', vmin=vmin, vmax=vmax)
    plt.colorbar(im, label='Value')
    plt.title(title)
    plt.xlabel('Longitude')
    plt.ylabel('Latitude')
    plt.savefig(filename)
    plt.close()

def plot_histogram(data1, data2, title, filename, xlabel="Value", ylabel="Frequency"):
    plt.figure(figsize=(10, 6))
    plt.hist(data1, bins=50, alpha=0.7, label='Training Data')
    plt.hist(data2, bins=50, alpha=0.7, label='Testing Data')
    plt.legend()
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.savefig(filename)
    plt.close()
 
def plot_residuals(y_true, residuals, title, filename, xlabel="True Values", ylabel="Residuals"):
    plt.figure(figsize=(10, 6))
    plt.scatter(y_true, residuals, alpha=0.5)
    plt.axhline(0, color='red', linestyle='--')
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.savefig(filename)
    plt.close()          
    
# Start tracking the training time
start_time = time.time()
wandb_state =False
# Start tracking the training time
start_time = time.time()

# Parse command-line arguments
parser = argparse.ArgumentParser(description="Train a regression model to predict Leaf Area Index (LAI).")
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', 'linear_regression'], help="Model type to use")
parser.add_argument("--imputation", type=str, choices=['mean', 'median', 'most_frequent', 'knn'], 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("--save_data_path", type=str, required=True, help="Path to save data and model outputs")
parser.add_argument("--autoregressive", action='store_true', help="Include autoregressive feature from ground truth data (year +1)")
parser.add_argument("--autoregressive_smoothing", action='store_true', help="Smooth the autoregression values for reducing the importance of the LAI LKS", default=False)
parser.add_argument("--feature_normalise", action='store_true', help="Normalise the features to the max value", default=False)
parser.add_argument("--ground_truth_normalise", action='store_true', help="Normalise the ground truth to the max value", default=False)
parser.add_argument("--ground_truth_log_transform", action='store_true', help="Apply log transform to ground truth", default=False)   
parser.add_argument("--model_timestamp", type=str, required=True, help="Model timestamp from training phase")
parser.add_argument(
    "--task_type", 
    type=str, 
    required=True, 
    choices=['classification', 'regression'], 
    help="Task type: 'classification' or 'regression'"
)

# TODO in case normalise , it should expect min and max values! 

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.save_data_path
autoregressive = args.autoregressive
autoregressive_smoothing = args.autoregressive_smoothing
timestamp = args.model_timestamp
feature_normalise = args.feature_normalise
ground_truth_normalise = args.ground_truth_normalise
ground_truth_log_transform = args.ground_truth_log_transform
 
task_type = args.task_type

lower_bound = 0
upper_bound = 7
sigma = 10
mode='nearest'
    
if wandb_state: 
    # Initialize W&B
    wandb.init(
        project="leaf-area-index-regression",
        config={
            "start_year": start_year,
            "end_year": end_year,
            "region": region,
            "model_type": model_type,
            "imputation": imputation_method,
            "train_years": train_years,
            "test_years": test_years,
            "n_estimators": 2000 if model_type == "xgboost" else None,
        },
        mode="offline"  # Run in offline mode
)

# Define bin edges for 14 bins
bin_edges = np.arange(0, 7.25, 0.5)

# Configure logging
current_path = os.getcwd()
logs_path = os.path.join(current_path, "Logs")
os.makedirs(logs_path, exist_ok=True)
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} with the task type: {task_type}. The Autoregressive is: {autoregressive} and smoothing is: {autoregressive_smoothing}.")
logging.info(f"Timestamp is : {timestamp}")

# Set paths for saving/loading processed data and metadata
metadata_path = os.path.join(save_data_path, f'metadata_training_{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}")

# Extract feature order dynamically
feature_order = metadata.get('feature_order', {})
logging.info(f"Feature order from metadata: {feature_order}")

# Extract additional configurations
include_luh2 = metadata.get('include_luh2', True) # Defaults is True
include_hilda = metadata.get('include_hilda', True) # Defaults is True
include_esa_cci = metadata.get('include_esa_cci', False) # Defaults is True
lat_dim = metadata.get('lat_dim')
lon_dim = metadata.get('lon_dim')

# if metadata.get("log_transform", False):  # Check if log_transform was applied
#     log_transform = True
# else:
#     log_transform = False
   
# Use the metadata's data path
data_path = metadata.get('data_path', save_data_path)
logging.info(f"Using data path from metadata: {data_path}")
logging.info(f"Expected dimensions (lat x lon): {lat_dim} x {lon_dim}")

# Dynamically adjust feature importance mapping
feature_names = [feature_order[str(i)] for i in range(len(feature_order))]

# Logging feature importance with feature names
def log_feature_importance(importance, feature_names, save_path, timestamp):
    feature_importance_named = [
        {"feature": feature_names[int(k[1:])], "importance": v}
        for k, v in sorted(importance.items(), key=lambda x: x[1], reverse=True)
    ]
    logging.info(f"Feature importance: {feature_importance_named}")

    # Save feature importance to a file
    feature_importance_path = os.path.join(save_path, f'feature_importance_XGBoost_{timestamp}.json')
    with open(feature_importance_path, 'w') as f:
        json.dump(feature_importance_named, f, indent=4)
    logging.info(f"Feature importance with names saved to {feature_importance_path}")

def load_data_for_years(years, region, data_path):
    features_list = []
    LAI_list = []

    for year in years:
        features_year_path = os.path.join(data_path, f'features_training_{year}_{region}.zarr')
        LAI_year_path = os.path.join(data_path, f'LAI_{year}_{region}.zarr')

        # Check if feature file exists
        if not os.path.exists(features_year_path):
            logging.error(f"Feature file not found: {features_year_path}")
            sys.exit(1)

        # Check if LAI file exists
        if not os.path.exists(LAI_year_path):
            logging.error(f"LAI file not found: {LAI_year_path}")
            sys.exit(1)

        try:
            features = zarr.load(features_year_path)
            LAI = zarr.load(LAI_year_path)
        except Exception as e:
            logging.error(f"Error loading data for year {year}: {e}")
            sys.exit(1)

        # Check if loaded data is not empty and has correct dimensions
        if features.size == 0:
            logging.error(f"Features data is empty for year {year}: {features_year_path}")
            sys.exit(1)
        if LAI.size == 0:
            logging.error(f"LAI data is empty for year {year}: {LAI_year_path}")
            sys.exit(1)

        if features.ndim == 0:
            logging.error(f"Features data is zero-dimensional for year {year}: {features_year_path}")
            sys.exit(1)
        if LAI.ndim == 0:
            logging.error(f"LAI data is zero-dimensional for year {year}: {LAI_year_path}")
            sys.exit(1)

        # DEBUG: Log shapes for each year
        logging.info(f"Year {year} - features shape: {features.shape}, LAI shape: {LAI.shape}")
        features_list.append(features)
        LAI_list.append(LAI)

    # After loading all years, check if lists are not empty
    if not features_list:
        logging.error("No features were loaded. Exiting.")
        sys.exit(1)
    if not LAI_list:
        logging.error("No LAI data was loaded. Exiting.")
        sys.exit(1)

    try:
        features = np.concatenate(features_list, axis=0)
        LAI = np.concatenate(LAI_list, axis=0)
        # DEBUG: Log concatenated shapes
        logging.info(f"After concatenation: features shape: {features.shape}, LAI shape: {LAI.shape}")
    except ValueError as ve:
        logging.error(f"Error concatenating arrays: {ve}")
        sys.exit(1)

    return features, LAI

# Define normalization range
def normalize_data(data, lower_bound, upper_bound):
    return (data - lower_bound) / (upper_bound - lower_bound)

def denormalize_data(data, lower_bound, upper_bound):
    return data * (upper_bound - lower_bound) + lower_bound

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

# Apply normalization to ground truth
if ground_truth_normalise:
    logging.info(f"Applying manual normalization to ground truth using range [{lower_bound}, {upper_bound}].")
    y_train = normalize_data(y_train, lower_bound, upper_bound)
    y_test = normalize_data(y_test, lower_bound, upper_bound)
    logging.info("Ground truth normalization applied successfully.")

if ground_truth_log_transform:
    logging.info("Applying log transformation to ground truth data.")
    y_train = np.log1p(y_train)
    y_test = np.log1p(y_test)
    logging.info("Log transformation applied successfully.")
        
if not autoregressive:
    logging.info("Smoothing autoregressive feature using Gaussian filter.")
    logging.info("Autoregressive feature is disabled. Overwriting the first feature with zeros.")
    X_train[:, 0] = 0
    X_test[:, 0] = 0

#  Smoothing of the autoregressive LAI   
if autoregressive and autoregressive_smoothing:
    
    logging.info("Autoregressive feature and smoothing of Autoregressive feature are active.") 
    num_samples = X_train.shape[0]
    num_time_steps = num_samples // (lat_dim * lon_dim)
    if num_time_steps * lat_dim * lon_dim != num_samples:
        raise ValueError(f"Mismatch between spatial dimensions and samples. "
                        f"Expected {lat_dim * lon_dim * num_time_steps} samples, but got {num_samples}.")

    logging.info(f"Applying smoothing to autoregressive feature over {num_time_steps} time steps.")
    # Reshape for smoothing
    X_train_reshaped = X_train[:, 0].reshape(num_time_steps, lat_dim, lon_dim)
    smoothed_train = gaussian_filter(X_train_reshaped, sigma=sigma, mode=mode)
    # Use explicit C order for flattening
    X_train[:, 0] = smoothed_train.flatten(order='C')
    
    # TODO : adapt the code here to test for more than one year data. 
    X_test_reshaped = X_test[:, 0].reshape(1, lat_dim, lon_dim)  # Assuming X_test is for one time step
    smoothed_test = gaussian_filter(X_test_reshaped, sigma=sigma, mode=mode)
    X_test[:, 0] = smoothed_test.flatten(order='C')

logging.info(f"Autoregressive smoothing completed. X_train and X_test shapes: {X_train.shape}, {X_test.shape}")

# Impute missing values
if imputation_method in ['mean', 'median', 'most_frequent']:
    imputer = SimpleImputer(strategy=imputation_method)
else:
    imputer = KNNImputer(n_neighbors=5)

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

prior_to_bin = time.time() 

if task_type == "classification":
    # Bin the ground truth values into classes
    y_train_binned = np.digitize(y_train, bins=bin_edges, right=False) - 1  # Subtract 1 to make bins 0-indexed
    y_test_binned = np.digitize(y_test, bins=bin_edges, right=False) - 1
    logging.info("Binned ground truth into classes for classification task.")

    # Compute class weights for classification
    class_weights = compute_class_weight(
        class_weight='balanced',
        classes=np.unique(y_train_binned),
        y=y_train_binned
    )
    class_weight_dict = dict(zip(np.unique(y_train_binned), class_weights))
    sample_weights = np.array([class_weight_dict[label] for label in y_train_binned])
if task_type == "regression":
    # For regression, use original LAI values
    logging.info("Using original LAI values for regression task.")
    y_train_binned = y_train  # No binning
    y_test_binned = y_test
    sample_weights = None  # No sample weights for regression

# Capture training duration
bining_duration = time.time() - prior_to_bin
logging.info(f"Bining for classification/ Preprocessing for regression is completed in {bining_duration:.2f} seconds.")


# Function to train and evaluate models
def train_evaluate_model(model, X_train, X_test, y_train, y_test, ground_truth_normalise, lower_bound, upper_bound, ground_truth_log_transform, model_name, test_years, region, save_data_path, timestamp, sample_weights=None):
    # Split training data into train and validation subsets
    X_train_sub, X_val, y_train_sub, y_val = train_test_split(X_train, y_train, test_size=0.2, random_state=42)
    
    if task_type == "classification":
        # Recalculate class weights for validation set
        val_weights = np.array([
            class_weight_dict[label] for label in y_val
        ]) if 'class_weight_dict' in locals() else compute_class_weight(
            class_weight='balanced',
            classes=np.unique(y_val),
            y=y_val
        )
    
        # Fit the classifier
        if sample_weights is not None:
            model.fit(
                X_train_sub, y_train_sub,
                sample_weight=sample_weights,
                eval_set=[(X_train_sub, y_train_sub), (X_val, y_val)],
                verbose=10
            )
        else:
            model.fit(
                X_train_sub, y_train_sub,
                eval_set=[(X_train_sub, y_train_sub), (X_val, y_val)],
                verbose=10
            )
     # Fit the regression 
    elif task_type == "regression":
        model.fit(
            X_train_sub, y_train_sub,
            eval_set=[(X_train_sub, y_train_sub), (X_val, y_val)], 
            verbose=10
        )
        
    # Extract evaluation results
    evals_result = model.evals_result()

    # Log evaluation results
    # if evals_result:
    #     logging.info(f"Evaluation results: {evals_result}")
    
    # Extract feature importance for XGBoost
    if model_name == "XGBoost":
      # Retrieve feature importance based on gain
        feature_importance = model.get_booster().get_score(importance_type='gain')
        log_feature_importance(feature_importance, feature_names, save_data_path, timestamp)
        
    y_pred = model.predict(X_test)
    
    # DEBUG: Log flattened sizes for predictions
    logging.info(f"Prediction shape: {y_pred.shape}")
    logging.info(f"Ground Truth shape: {y_test.shape}")
    logging.info(f"Expected shape from metadata: ({lat_dim}, {lon_dim})")
    logging.info(f"Flattened y_pred length: {y_pred.flatten(order='C').size} vs expected: {lat_dim * lon_dim}")

    # Load the water mask for the test years and apply it to the predictions
    water_masks = []
    for year in test_years:
        water_mask_path = os.path.join(save_data_path, f'water_mask_{region}_{year}.zarr')
        if os.path.exists(water_mask_path):
            water_mask = zarr.load(water_mask_path)
            logging.info(f"Water mask shape for year {year}: {water_mask.shape}")
            water_masks.append(water_mask.flatten(order='C'))  # Use explicit C order for flattening
        else:
            logging.warning(f"Water mask for year {year} not found.")

    # Inverse log transformation
    if ground_truth_log_transform:  # Check if log_transform was applied
        logging.info("Applying inverse log transformation to predictions and ground truth.")
        y_pred = np.expm1(y_pred)
        y_test = np.expm1(y_test)
        logging.info("Inverse log transformation applied to prediction and test dataset successfully.")

    if ground_truth_normalise:
        logging.info("Denormalizing predictions back to the original range.")
        y_pred = denormalize_data(y_pred, lower_bound, upper_bound)
        y_test = denormalize_data(y_test, lower_bound, upper_bound)
        logging.info("Denormalization of predictions completed.")

    ########
    ########
    ## POST PROCESSING ANALYSIS 
    
    # Plot target distribution 
    target_dist_path = os.path.join(save_data_path, f'target_distribution_{region}_{year}_{timestamp}.png')
    plot_histogram(
        y_train,
        y_test,
        "Target Distribution",
        target_dist_path
    )
    logging.info(f"Target distribution plot saved at: {target_dist_path}")

    # Plot residuals
    residuals = y_test - y_pred
    residuals_plot_path = os.path.join(save_data_path, f'residuals_plot_{region}_{year}_{timestamp}.png')
    plot_residuals(
        y_test,
        residuals,
        "Residuals Plot",
        residuals_plot_path
    )
    logging.info(f"Residuals plot saved at: {residuals_plot_path}")

    ##########
    ##########
    # POST_PROCESSING & Regularization 

    # Bias Correction
    # simple bias correction 
    bias_correction = np.mean(y_test - y_pred)
    y_pred_corrected = y_pred + abs(bias_correction)
    logging.info(f"Bias Correction: {abs(bias_correction)}")
    
    # Optionally clip corrected predictions
    y_pred_corrected = np.clip(y_pred_corrected, lower_bound, upper_bound)
    logging.info("Finished bias correction. Predictions clipped to range.")
    
   # Print summary of corrected and scaled predictions
    print("The Ground Truth:")
    print(f"Min: {np.min(y_test)}, Max: {np.max(y_test)}, Mean: {np.mean(y_test)}")
    print("Summary of Raw Predictions:")
    print(f"Min: {np.min(y_pred)}, Max: {np.max(y_pred)}, Mean: {np.mean(y_pred)}")
    print("Summary of Corrected Predictions:")
    print(f"Min: {np.min(y_pred_corrected)}, Max: {np.max(y_pred_corrected)}, Mean: {np.mean(y_pred_corrected)}")
    
    logging.info("The Ground Truth:")
    logging.info(f"Min: {np.min(y_test)}, Max: {np.max(y_test)}, Mean: {np.mean(y_test)}")
    logging.info("Summary of Raw Predictions:")
    logging.info(f"Min: {np.min(y_pred)}, Max: {np.max(y_pred)}, Mean: {np.mean(y_pred)}")
    logging.info("Summary of Corrected Predictions:")
    logging.info(f"Min: {np.min(y_pred_corrected)}, Max: {np.max(y_pred_corrected)}, Mean: {np.mean(y_pred_corrected)}")
    
    # Save adjusted predictions
    # bias_corrected_path = os.path.join(save_data_path, f'bias_corrected_predictions_{timestamp}.zarr')
    # zarr.save(bias_corrected_path, y_pred_corrected)
    # logging.info(f"Bias-corrected predictions saved at: {bias_corrected_path}")    
    
    # Concatenate water masks and ensure alignment with y_pred
    if water_masks:
        combined_water_mask = np.concatenate(water_masks, axis=0)
        logging.info(f"Combined water mask shape: {combined_water_mask.shape}")
        logging.info(f"Predictions shape: {y_pred.shape}")
        if combined_water_mask.shape[0] != y_pred.shape[0]:
            raise ValueError("Mismatch between water mask and predictions dimensions. Check preprocessing pipeline.")
        # Apply water mask
        y_pred = y_pred * combined_water_mask
        y_pred_corrected = y_pred_corrected * combined_water_mask
        logging.info("Applied water mask to predictions.")
    else:
        logging.warning("No valid water masks found for the test years. Predictions were not masked.")

    # Plot histograms for corrected and scaled predictions
    corrected_histogram_path = os.path.join(save_data_path, f'histogram_corrected_predictions_{region}_{year}_{timestamp}.png')
    
    plot_histogram(
        y_pred_corrected,
        y_test,
        "Histogram - Corrected Predictions vs True Values",
        corrected_histogram_path
    )
    logging.info(f"Corrected predictions histogram saved at: {corrected_histogram_path}")

    # Plot residuals for corrected predictions
    corrected_residuals = y_test - y_pred_corrected
    corrected_residuals_plot_path = os.path.join(save_data_path, f'residuals_corrected_predictions_{region}_{year}_{timestamp}.png')
    plot_residuals(
        y_test,
        corrected_residuals,
        "Residuals Plot - Corrected Predictions",
        corrected_residuals_plot_path
    )
    logging.info(f"Residuals plot for corrected predictions saved at: {corrected_residuals_plot_path}")
    
    # Plot predictions for regression
    if task_type == "regression":
        
        # including padding , in case the size of the prediction doesn match the target 
        def reshape_with_padding(arr, target_shape):
            """
            Reshape a flattened array to the target shape with padding or trimming.

            Parameters:
                arr (np.ndarray): Flattened array to reshape.
                target_shape (tuple): Desired shape (lat_dim, lon_dim).

            Returns:
                np.ndarray: Reshaped array with padding or trimming.
            """
            target_size = target_shape[0] * target_shape[1]
            current_size = arr.size
            if current_size != target_size:
                logging.warning(f"reshape_with_padding: Array size {current_size} != target size {target_size}.")
            if current_size < target_size:
                # Pad with zeros if the array is smaller
                padding_size = target_size - current_size
                padded_arr = np.pad(arr, (0, padding_size), mode='constant', constant_values=0)
                reshaped_arr = padded_arr.reshape(target_shape, order='C')
            elif current_size > target_size:
                # Trim the array if it's larger
                reshaped_arr = arr[:target_size].reshape(target_shape, order='C')
            else:
                # Reshape directly if sizes match
                reshaped_arr = arr.reshape(target_shape, order='C')

            return reshaped_arr
        
        predicted_data = reshape_with_padding(y_pred, (lat_dim, lon_dim)) 
        true_data = reshape_with_padding(y_test, (lat_dim, lon_dim))       
        corrected_data = reshape_with_padding(y_pred_corrected, (lat_dim, lon_dim))

        # Paths for saving plots
        pred_plot_path = os.path.join(save_data_path, f'prediction_plot_{model_name}_{region}_{year}_{timestamp}.png')
        true_plot_path = os.path.join(save_data_path, f'true_plot_{model_name}__{region}_{year}_{timestamp}.png')
        corrected_plot_path = os.path.join(save_data_path, f'prediction_plot_corrected_{model_name}_{region}_{year}_{timestamp}.png')

        # Plot predicted values
        plot_data(
            predicted_data,
            f"Predicted LAI - {region} ({start_year}-{end_year})",
            pred_plot_path,
            vmin=lower_bound,   # force the minimum to 0
            vmax=upper_bound    # force the maximum to 7
        )
        logging.info(f"Prediction plot saved at: {pred_plot_path}")

        # Plot true values
        plot_data(
            true_data,
            f"True LAI - {region} ({start_year}-{end_year})",
            true_plot_path,
            vmin=lower_bound,
            vmax=upper_bound
        )
        logging.info(f"True value plot saved at: {true_plot_path}")
        # Plot bias-corrected predictions
        plot_data(
            corrected_data,
            f"Bias-Corrected Predicted LAI - {region} ({start_year}-{end_year})",
            corrected_plot_path,
            vmin=lower_bound,
            vmax=upper_bound
        )
        logging.info(f"Bias-corrected prediction plot saved at: {corrected_plot_path}")
       
    # Plot predictions for classification
    if task_type == "classification":
        # Reshape predictions and true values for plotting
        predicted_classes = y_pred.reshape(lat_dim, lon_dim)
        true_classes = y_test.reshape(lat_dim, lon_dim)

        # Define colormap for classification
        cmap = plt.get_cmap('tab20')  # Adjust colormap as needed for your classes

        # Paths for saving plots
        pred_class_plot_path = os.path.join(save_data_path, f'predicted_classes_plot_{model_name}_{timestamp}.png')
        true_class_plot_path = os.path.join(save_data_path, f'true_classes_plot_{model_name}_{timestamp}.png')

        # Plot predicted class labels
        plot_data(
            predicted_classes,
            f"Predicted Classes - {region} ({start_year}-{end_year})",
            pred_class_plot_path
        )
        logging.info(f"Predicted classes plot saved at: {pred_class_plot_path}")

        # Plot true class labels
        plot_data(
            true_classes,
            f"True Classes - {region} ({start_year}-{end_year})",
            true_class_plot_path
        )
        logging.info(f"True classes plot saved at: {true_class_plot_path}")

    
    
    # Compute regression/classification metrics
    if model_name == 'XGBoost':
        if task_type == "classification":
            # Classification metrics
            accuracy = accuracy_score(y_test, y_pred)
            f1 = f1_score(y_test, y_pred, average='weighted')
            precision = precision_score(y_test, y_pred, average='weighted')
            recall = recall_score(y_test, y_pred, average='weighted')
            class_report = classification_report(y_test, y_pred, zero_division=0)
            
            # Log metrics to console
            logging.info(f'Accuracy of {model_name}: {accuracy:.4f}')
            logging.info(f'F1 Score of {model_name}: {f1:.4f}')
            logging.info(f'Precision of {model_name}: {precision:.4f}')
            logging.info(f'Recall of {model_name}: {recall:.4f}')
            logging.info(f'Classification Report of {model_name}:\n{class_report}')
            
            # Log metrics to W&B
            if wandb_state: 
                wandb.log({
                    'accuracy': accuracy,
                    'f1_score': f1,
                    'precision': precision,
                    'recall': recall
                })
        
        elif task_type == "regression":
            # Regression metrics
            # Regression metrics for raw predictions
            mse = mean_squared_error(y_test, y_pred)
            rmse = np.sqrt(mse)
            mae = mean_absolute_error(y_test, y_pred)
            r2 = r2_score(y_test, y_pred)

            # Log metrics for raw predictions
            logging.info(f'Metrics for {model_name} (Raw Predictions):')
            logging.info(f'MSE: {mse:.4f}')
            logging.info(f'RMSE: {rmse:.4f}')
            logging.info(f'MAE: {mae:.4f}')
            logging.info(f'R²: {r2:.4f}')

            # Regression metrics for bias-corrected predictions
            mse_corrected = mean_squared_error(y_test, y_pred_corrected)
            rmse_corrected = np.sqrt(mse_corrected)
            mae_corrected = mean_absolute_error(y_test, y_pred_corrected)
            r2_corrected = r2_score(y_test, y_pred_corrected)

            # Log metrics for bias-corrected predictions
            logging.info(f'Metrics for {model_name} (Bias-Corrected Predictions):')
            logging.info(f'MSE: {mse_corrected:.4f}')
            logging.info(f'RMSE: {rmse_corrected:.4f}')
            logging.info(f'MAE: {mae_corrected:.4f}')
            logging.info(f'R²: {r2_corrected:.4f}')

            
            # Log metrics to W&B
            if wandb_state: 
                wandb.log({
                    'mse': mse,
                    'rmse': rmse,
                    'mae': mae,
                    'r2_score': r2
            })

        gc.collect()

    # Plot evaluation metrics for regression
    if task_type == "regression" and evals_result:
        plot_filename = os.path.join(save_data_path, f'rmse_plot_{model_name}__{region}_{timestamp}.png')
        plt.figure(figsize=(10, 6))
        plt.plot(evals_result['validation_0']['rmse'], label='Train RMSE')
        plt.plot(evals_result['validation_1']['rmse'], label='Validation RMSE')
        plt.xlabel('Boosting Rounds')
        plt.ylabel('RMSE')
        plt.legend()
        plt.title(f'{model_name} Training RMSE')
        plt.savefig(plot_filename)
        plt.close()
        logging.info(f"RMSE plot saved as {plot_filename}")


    # Plot evaluation metrics for classification
    if task_type == "classification" and evals_result:
        plot_filename = os.path.join(save_data_path, f'log_loss_plot_{model_name}_{timestamp}.png')
        plt.figure(figsize=(10, 6))
        plt.plot(evals_result['validation_0']['mlogloss'], label='Train Log Loss')
        plt.plot(evals_result['validation_1']['mlogloss'], label='Validation Log Loss')
        plt.xlabel('Boosting Rounds')
        plt.ylabel('Log Loss')
        plt.legend()
        plt.title(f'{model_name} Training and Validation Log Loss')
        plt.savefig(plot_filename)
        plt.close()
        logging.info(f"Log loss plot saved as {plot_filename}")
 
       
    # --- POSTPROCESSING INFO ---
    postprocessing_info = {
        
        "autoregression": {
            "autoregression_applied": bool(autoregressive),
            "autoregressive_smoothing_applied": bool(autoregressive_smoothing),
            "mode": mode,
            "sigma": sigma,            
            "description": "Autoregression and smoothing method indformation ."
        },
        "bias_correction": {
            "method": "mean_offset",
            "bias_value": abs(bias_correction),
            "description": "Bias correction by adding mean offset between y_test and y_pred."
        },
        "clipping_range": {
            "lower_bound": lower_bound,
            "upper_bound": upper_bound
        },
        "water_mask": {
            "applied": bool(water_masks),
            "years_processed": test_years if water_masks else None
        },
        "metrics_summary": {
            "raw_predictions": {
                "min": float(np.min(y_pred)),
                "max": float(np.max(y_pred)),
                "mean": float(np.mean(y_pred))
            },
            "bias_corrected_predictions": {
                "min": float(np.min(y_pred_corrected)),
                "max": float(np.max(y_pred_corrected)),
                "mean": float(np.mean(y_pred_corrected))
            }
        },
        "plots_generated": {
            "target_distribution": target_dist_path,
            "residuals_raw": residuals_plot_path,
            "residuals_corrected": corrected_residuals_plot_path,
            "histogram_corrected_predictions": corrected_histogram_path
        }
    }
    # Save the model
    model_file_path = os.path.join(save_data_path, f'{model_name}_LAI_{start_year}_{end_year}_{region}_{timestamp}.joblib')
    joblib.dump(model, model_file_path)
    logging.info(f'{model_name} model saved successfully at {model_file_path}.')

    return postprocessing_info

# Initialize and train the model based on the chosen architecture
if model_type == "xgboost":
    if task_type == "classification":
        model = xgb.XGBClassifier(
            n_estimators=500,
            max_depth=8,
            learning_rate=0.02,
            subsample=0.8,
            reg_alpha=0.5,
            reg_lambda=1.0,
            colsample_bytree=0.8,
            tree_method='hist',
            random_state=42,
            device='cuda',
            verbosity=2,
            eval_metric='mlogloss',
            early_stopping_rounds=50  # Early stopping here
        )
    elif task_type == "regression":
        model = xgb.XGBRegressor(
            n_estimators=700,
            max_depth=12,
            learning_rate=0.01,
            subsample=0.8,
            reg_alpha=0.5,
            reg_lambda=1.0,
            colsample_bytree=0.8,
            tree_method='hist',
            random_state=42,
            device='cuda',
            verbosity=2,
            eval_metric='rmse',
            early_stopping_rounds=50
        )
    postprocessing_info = train_evaluate_model(model, X_train, X_test, y_train_binned, y_test_binned, ground_truth_normalise, lower_bound, upper_bound, ground_truth_log_transform, "XGBoost", test_years, region, save_data_path, timestamp, sample_weights)

elif model_type == "linear_regression":
    model = LinearRegression()
    train_evaluate_model(model, X_train, X_test, y_train, y_test, "LinearRegression")
else:
    raise ValueError("Invalid model type specified. Choose from 'xgboost', 'linear_regression'.")

# Capture training duration
training_duration = time.time() - start_time
logging.info(f"Training completed in {training_duration:.2f} seconds.")

# Capture hardware details
hardware_info = {
    'num_cores': os.cpu_count(),
    'memory_gb': psutil.virtual_memory().total / (1024 ** 3),
    'gpu_used': model_type == 'xgboost'
}

def convert_numpy_types(obj):
    """
    Recursively convert NumPy data types to standard Python types.
    """
    if isinstance(obj, (np.integer, int)):
        return int(obj)
    elif isinstance(obj, (np.floating, float)):
        return float(obj)
    elif isinstance(obj, np.ndarray):
        return obj.tolist()  # Convert numpy arrays to lists
    elif isinstance(obj, dict):
        return {str(k): convert_numpy_types(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [convert_numpy_types(v) for v in obj]
    elif isinstance(obj, tuple):
        return tuple(convert_numpy_types(v) for v in obj)
    else:
        return obj


# Create metadata dictionary for training process
training_metadata = {
    'start_year': start_year,
    'end_year': end_year,
    'region': region,
    'model_type': model_type,
    'train_years': [convert_numpy_types(year) for year in train_years],
    'test_years': [convert_numpy_types(year) for year in test_years],
    'imputation': imputation_method,
    'n_estimators': model.get_params().get('n_estimators', None),
    'max_iter': model.get_params().get('max_iter', None),
    'solver': model.get_params().get('solver', None),
    'timestamp': timestamp,
    'X_train_shape': convert_numpy_types(X_train.shape),
    'X_test_shape': convert_numpy_types(X_test.shape),
    'random_seed': 42,
    'training_duration': training_duration,
    'model_hyperparameters': {str(k): convert_numpy_types(v) for k, v in model.get_params().items()},
    'hardware_info': {str(k): convert_numpy_types(v) for k, v in hardware_info.items()},
    'ground_truth_preprocessing': {
        'log_transform': ground_truth_log_transform,
        'normalization': ground_truth_normalise,
        'normalization_range': {
            'lower_bound': lower_bound,
            'upper_bound': upper_bound
        } if ground_truth_normalise else None
    },
    'postprocessing': postprocessing_info
}


# Save training metadata to a JSON file
training_metadata_filename = f'{model_type}_LAI_training_metadata_{start_year}_{end_year}_{region}_{timestamp}.json'
training_metadata_path = os.path.join(save_data_path, training_metadata_filename)
gc.collect()

training_metadata_cleaned = convert_numpy_types(training_metadata)

with open(training_metadata_path, 'w') as f:
    json.dump(training_metadata_cleaned, f, indent=4)
    gc.collect()


logging.info(f'Training metadata saved successfully to {training_metadata_path}.')

logging.info("Training and evaluation complete.")
if wandb_state: 
    wandb.finish()