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.multioutput import MultiOutputRegressor   
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

class XGBLoggingCallback(xgb.callback.TrainingCallback):
    def __init__(self, period=10):
        self.period = period

    def after_iteration(self, model, epoch, evals_log):
        if epoch % self.period == 0:
            msg_parts = []
            # evals_log is a dict where keys are dataset names and values are dicts of metric lists.
            for dataset, metrics in evals_log.items():
                for metric, values in metrics.items():
                    # Get the last reported value for this metric.
                    if values:
                        msg_parts.append(f"{dataset}-{metric}: {values[-1]:.4f}")
            logging.info(f"Iteration {epoch} - " + ", ".join(msg_parts))
        # Return False to indicate training should not stop.
        return False

def sanitize_params(params):
    """Return a sanitized copy of params where non-serializable items are converted to strings."""
    sanitized = {}
    for key, value in params.items():
        try:
            json.dumps(value)
            sanitized[key] = value
        except TypeError:
            # If the value is not serializable, store its string representation.
            sanitized[key] = str(value)
    return sanitized

## 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 = []
    for k, v in sorted(importance.items(), key=lambda x: x[1], reverse=True):
        try:
            idx = int(k[1:])  # assuming keys like 'f0', 'f1', etc.
        except ValueError:
            idx = None
        if idx is not None and idx < len(feature_names):
            fname = feature_names[idx]
        else:
            fname = k  # fallback to the key if no corresponding name exists
        feature_importance_named.append({"feature": fname, "importance": v})
        
    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_flat_{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)  # Should have shape (lat*lon, 12)
            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":
        # If using MultiOutputRegressor (i.e. for multi-output monthly prediction)
        if isinstance(model, MultiOutputRegressor):
            # Disable early stopping since we cannot pass a validation set with MultiOutputRegressor
            model.estimator.set_params(early_stopping_rounds=None)
            model.fit(X_train_sub, y_train_sub)
            logging.warning("DEBUG: The model is MultiOutputRegressor, early stopping disabled.")
            
            # Retrieve feature importance from the first underlying estimator
            feature_importance = model.estimators_[0].get_booster().get_score(importance_type='gain')
            log_feature_importance(feature_importance, feature_names, save_data_path, timestamp)
        else:
            model.fit(
                X_train_sub, y_train_sub,
                eval_set=[(X_train_sub, y_train_sub), (X_val, y_val)], 
                verbose=10
            )
            feature_importance = model.get_booster().get_score(importance_type='gain')
            log_feature_importance(feature_importance, feature_names, save_data_path, timestamp)
         
    # Extract evaluation results
    evals_result = model.evals_result() if hasattr(model, "evals_result") else None

    # Log evaluation results if desired
    # if evals_result:
    #     logging.info(f"Evaluation results: {evals_result}")
    
    # Extract feature importance for XGBoost if applicable
    # if model_name == "XGBoost":
    #     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)  # For multi-output, shape: (n_samples, 12)
    
    # DEBUG: Log shapes and total elements
    logging.info(f"Prediction shape: {y_pred.shape}")
    logging.info(f"Ground Truth shape: {y_test.shape}")
    logging.info(f"Expected flat shape from metadata: ({lat_dim * lon_dim}, 12)")
    logging.info(f"Total elements in y_pred: {y_pred.size} vs expected: {lat_dim * lon_dim * 12}")

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

    # Inverse log transformation if applicable
    if ground_truth_log_transform:
        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 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 completed.")

    ########
    ## POST PROCESSING ANALYSIS 
    ########

    # Plot target distribution 
    target_dist_path = os.path.join(save_data_path, f'target_distribution_{region}_{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}_{timestamp}.png')
    plot_residuals(y_test, residuals, "Residuals Plot", residuals_plot_path)
    logging.info(f"Residuals plot saved at: {residuals_plot_path}")

    # 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)}")
    y_pred_corrected = np.clip(y_pred_corrected, lower_bound, upper_bound)
    logging.info("Finished bias correction. Predictions clipped to range.")
    
    # Print summaries
    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)}")
    
    # [UPDATED] Apply water mask: tile the water mask to match the multi-output shape (lat*lon, 12)
    if water_masks:
        combined_water_mask = np.concatenate(water_masks, axis=0)  # shape: (lat*lon,)
        logging.info(f"Combined water mask shape: {combined_water_mask.shape}")
        combined_water_mask = np.tile(combined_water_mask[:, np.newaxis], (1, y_pred.shape[1]))
        logging.info(f"Tiled water mask shape: {combined_water_mask.shape}")
        if combined_water_mask.shape != y_pred.shape:
            raise ValueError("Mismatch between water mask and predictions dimensions. Check preprocessing pipeline.")
        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 predictions vs. true values.
    corrected_histogram_path = os.path.join(save_data_path, f'histogram_corrected_predictions_{region}_{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}_{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":
        # [UPDATED] For multi-output, plot the first month's predictions.
        def reshape_with_padding(arr, target_shape):
            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:
                padded_arr = np.pad(arr, (0, target_size - current_size), mode='constant', constant_values=0)
                return padded_arr.reshape(target_shape, order='C')
            elif current_size > target_size:
                return arr[:target_size].reshape(target_shape, order='C')
            else:
                return arr.reshape(target_shape, order='C')
        
        # Loop over each month (assuming y_pred, y_test, and y_pred_corrected have shape (n_samples, 12))
        # TODO : plotting should be optional. 
        for month in range(y_pred.shape[1]):
            predicted_data = reshape_with_padding(y_pred[:, month], (lat_dim, lon_dim))
            true_data = reshape_with_padding(y_test[:, month], (lat_dim, lon_dim))
            corrected_data = reshape_with_padding(y_pred_corrected[:, month], (lat_dim, lon_dim))
            
            pred_plot_path = os.path.join(save_data_path, f'prediction_plot_{model_name}_month{month+1}_{region}_{timestamp}.png')
            true_plot_path = os.path.join(save_data_path, f'true_plot_{model_name}_month{month+1}_{region}_{timestamp}.png')
            corrected_plot_path = os.path.join(save_data_path, f'prediction_plot_corrected_{model_name}_month{month+1}_{region}_{timestamp}.png')
            
            plot_data(predicted_data,
                    f"Predicted LAI Month {month+1} - {region} ({start_year}-{end_year})",
                    pred_plot_path,
                    vmin=lower_bound, vmax=upper_bound)
            logging.info(f"Prediction plot for month {month+1} saved at: {pred_plot_path}")
            
            plot_data(true_data,
                    f"True LAI Month {month+1} - {region} ({start_year}-{end_year})",
                    true_plot_path,
                    vmin=lower_bound, vmax=upper_bound)
            logging.info(f"True value plot for month {month+1} saved at: {true_plot_path}")
            
            plot_data(corrected_data,
                    f"Bias-Corrected Predicted LAI Month {month+1} - {region} ({start_year}-{end_year})",
                    corrected_plot_path,
                    vmin=lower_bound, vmax=upper_bound)
            logging.info(f"Bias-corrected prediction plot for month {month+1} saved at: {corrected_plot_path}")
       
    # Plot predictions for classification
    if task_type == "classification":
        predicted_classes = y_pred.reshape(lat_dim, lon_dim)
        true_classes = y_test.reshape(lat_dim, lon_dim)
        cmap = plt.get_cmap('tab20')
        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_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_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":
            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)
            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}')
            if wandb_state: 
                wandb.log({
                    'accuracy': accuracy,
                    'f1_score': f1,
                    'precision': precision,
                    'recall': recall
                })
        
        elif task_type == "regression":
            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)
            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}')
    
            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)
            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}')
            if wandb_state: 
                wandb.log({
                    'mse': mse,
                    'rmse': rmse,
                    'mae': mae,
                    'r2_score': r2
                })
        gc.collect()

    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}")

    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 information."
        },
        "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":
        base_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
        )
        model = MultiOutputRegressor(base_model)
    elif task_type == "regression":
        base_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
            #callbacks=[XGBLoggingCallback(period=10)]

        )
        model = MultiOutputRegressor(base_model)
    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


# Get the model hyperparameters and sanitize them.
model_params = model.get_params()
sanitized_params = sanitize_params(model_params)

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': sanitized_params.get('n_estimators', None),
    'max_iter': sanitized_params.get('max_iter', None),
    'solver': sanitized_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': sanitized_params,
    '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()