import argparse
import dask.array as da
import json
import joblib
import numpy as np
import os
from datetime import datetime
import logging
import matplotlib.pyplot as plt
import xarray as xr
import zarr
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
import xgboost as xgb  # Adding XGBoost
from scipy.ndimage import gaussian_filter
import gc

# --- Utility Functions ---
def denormalize_data(data, lower_bound, upper_bound):
    """Denormalize data back to its original range."""
    return data * (upper_bound - lower_bound) + lower_bound

def inverse_log_transform(data):
    """Reverse log transformation."""
    return np.expm1(data)

def plot_lai(data, title, filename, vmin=None, vmax=None, show_plot=True):
    """
    Plot LAI data with optional fixed color scale.

    Parameters:
        data (np.ndarray): The LAI data to be plotted.
        title (str): The plot title.
        filename (str): The path where the plot will be saved.
        vmin (float, optional): Minimum value for the color scale. If provided, forces the lower bound.
        vmax (float, optional): Maximum value for the color scale. If provided, forces the upper bound.
        show_plot (bool): If True, shows the plot after saving.
    """
    plt.figure(figsize=(10, 6))
    # Pass vmin and vmax to imshow so that if they are provided, the color scale is fixed.
    im = plt.imshow(np.flipud(data), cmap='viridis', origin='upper', vmin=vmin, vmax=vmax)
    plt.colorbar(im, label='Leaf Area Index (LAI)')
    plt.title(title)
    plt.xlabel('Longitude')
    plt.ylabel('Latitude')
    plt.savefig(filename)
    if show_plot:
        plt.show()
    plt.close()

def plot_error_map(error_data, title, filename, vmin=None, vmax=None, show_plot=True):
    """
    Plot error data with optional fixed color scale.

    Parameters:
        error_data (np.ndarray): The error data to be plotted.
        title (str): The plot title.
        filename (str): The path where the plot will be saved.
        vmin (float, optional): Minimum value for the color scale.
        vmax (float, optional): Maximum value for the color scale.
        show_plot (bool): If True, shows the plot after saving.
    """
    plt.figure(figsize=(10, 6))
    im = plt.imshow(np.flipud(error_data), cmap='coolwarm', origin='upper', vmin=vmin, vmax=vmax)
    plt.colorbar(im, label='Error (Prediction - Ground Truth)')
    plt.title(title)
    plt.xlabel('Longitude')
    plt.ylabel('Latitude')
    plt.savefig(filename)
    if show_plot:
        plt.show()
    plt.close()

def reshape_with_padding(data, target_shape):
    """
    Reshape a flattened or 1D array to the target shape with padding or trimming.

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

    Returns:
        np.ndarray: Reshaped array with padding or trimming applied.
    """
    data = data.flatten()  # Ensure the data is flattened first
    target_size = target_shape[0] * target_shape[1]
    current_size = data.size

    if current_size < target_size:
        # Pad with zeros to match the target size
        padding_size = target_size - current_size
        reshaped_data = np.pad(data, (0, padding_size), mode='constant', constant_values=0)
    elif current_size > target_size:
        # Trim excess values
        reshaped_data = data[:target_size]
    else:
        reshaped_data = data

    # Reshape to the desired target shape
    return reshaped_data.reshape(target_shape).astype(np.float32)

# Argument Parsing
parser = argparse.ArgumentParser(description="LAI Prediction and Visualization")

parser.add_argument("--start_year", type=int, required=True, help="Start year of the dataset")
parser.add_argument("--end_year", type=int, required=True, help="End year of the dataset")
parser.add_argument("--start_year_training", type=int, required=True, help="Test year")
parser.add_argument("--end_year_training", type=int, required=True, help="Test years")
parser.add_argument("--region", type=str, required=True, help="Region code (e.g., 'ES')")
parser.add_argument("--save_data_path", type=str, required=True, help="Path to saved data and models")
parser.add_argument("--output_data_path", type=str, required=True, help="Path to save data outputs")
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'"
)

args = parser.parse_args()


# 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"lai_prediction_log_{args.start_year}_{args.end_year}_{args.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()
])

# Set Paths and Parameters from Arguments
start_year = args.start_year
end_year = args.end_year
start_year_training = args.start_year_training
end_year_training = args.end_year_training
region = args.region
save_data_path = args.save_data_path
output_data_path = args.output_data_path
model_timestamp = args.model_timestamp
task_type = args.task_type

# Load metadata to get dimensions
metadata_path_input_data = os.path.join(save_data_path, f'metadata_inference_{start_year}_{end_year}_{region}.json')
with open(metadata_path_input_data, 'r') as f:
    metadata_input_data = json.load(f)

# Load metadata to get dimensions
metadata_path_training = os.path.join(save_data_path, f'xgboost_LAI_training_metadata_{start_year_training}_{end_year_training}_{region}_{model_timestamp}.json')
with open(metadata_path_training, 'r') as f:
    metadata_training = json.load(f)

# Input dtaa information
lat_dim = metadata_input_data['lat_dim']
lon_dim = metadata_input_data['lon_dim']
logging.info(f"  Latitude dimension: {lat_dim}")
logging.info(f"  Longitude dimension: {lon_dim}")

# Retrieve parameters from metadata
bias_correction = metadata_training['postprocessing']['bias_correction']['bias_value']
lower_bound = metadata_training['postprocessing']['clipping_range']['lower_bound']
upper_bound = metadata_training['postprocessing']['clipping_range']['upper_bound']
log_transform = metadata_training['ground_truth_preprocessing']['log_transform']
normalization_applied = metadata_training['ground_truth_preprocessing']['normalization']
autoregressive = metadata_training['postprocessing']['autoregression']['autoregression_applied']
autoregressive_smoothing = metadata_training['postprocessing']['autoregression']['autoregressive_smoothing_applied']
sigma = metadata_training['postprocessing']['autoregression'].get('sigma', 10)

# Log all retrieved metadata parameters
logging.info("Training metadata has been retrieved:")
logging.info(f"  Bias correction: {bias_correction}")
logging.info(f"  Clipping range: [{lower_bound}, {upper_bound}]")
logging.info(f"  Log transform applied: {log_transform}")
logging.info(f"  Normalization applied: {normalization_applied}")
logging.info(f"  Autoregressive feature applied: {autoregressive}")
logging.info(f"  Autoregressive smoothing applied: {autoregressive_smoothing}")
logging.info(f"  Smoothing sigma value: {sigma}")

# Ensure to use the correct model path and timestamp from the training phase
model_path = os.path.join(save_data_path, f'XGBoost_LAI_{start_year_training}_{end_year_training}_{region}_{model_timestamp}.joblib')

# Create output folder based on model details
output_folder_name = f"{region}_{start_year}_{end_year}_{model_timestamp}_LAI"
output_folder_path = os.path.join(output_data_path, output_folder_name)

# Create the directory if it doesn't exist
os.makedirs(output_folder_path, exist_ok=True)

# Log the folder creation
logging.info(f"Saving all outputs to folder: {output_folder_path}")

# Load the Trained Model
logging.info("Loading trained model.")
print(model_path)
model = joblib.load(model_path)
previous_year_predictions = None

# Iterate over each test year
for year in range(end_year, start_year - 1, -1):
    
    #set the ground truth to true before loop 
    logging.info(f"Processing data for the year: {year}")

    # Paths for loading LAI ground truth and features for predictions
    
    features_path = os.path.join(save_data_path, f'features_inference_flat_{year}_{region}.zarr')
    water_mask_path = os.path.join(save_data_path, f'water_mask_{region}_{year}.zarr')

    # Load processed data for the current year
    logging.info("Loading processed data.")
    logging.info(f"The feature path is: {features_path}")

    # Convert Dask arrays to NumPy arrays for predictions
    water_mask = da.from_zarr(water_mask_path)
    logging.info(f"Shape of water_mask (Dask): {water_mask.shape}")
    water_mask = water_mask.compute()
    logging.info(f"Shape of water_mask (NumPy): {water_mask.shape}")
    
    # Reshape the water mask to match predictions
    water_mask = reshape_with_padding(water_mask, (lat_dim, lon_dim))
    logging.info(f"Reshaped water mask to: {water_mask.shape}")

    features = da.from_zarr(features_path)
    logging.info(f"Shape of features (Dask): {features.shape}")
    features = features.compute()
    logging.info(f"Shape of features (NumPy): {features.shape}")

   # If we have predictions from the previous year, replace the placeholder
    if autoregressive and previous_year_predictions is not None:
        logging.info("Updating placeholder with the previous year's predictions.")
        try:
            # Ensure `previous_year_predictions` is reshaped correctly
            # Update the placeholder in the features (assuming the autoregressive feature occupies the first 12 columns)
            features[:, :12] = AR_flat
            logging.info(f"Updated autoregressive feature with shape: {AR_flat.shape}")
    
        except ValueError as e:
            logging.error(f"Error updating placeholder with previous year's predictions: {e}")
            raise

    # Make predictions
    logging.info("Making predictions.")
    predictions = model.predict(features)
    logging.info(f"Shape of predictions (Dask): {predictions.shape}")
    # predictions = reshape_with_padding(predictions, (lat_dim, lon_dim))
    # logging.info(f"Reshaped predictions to: {predictions.shape}")
    
    # Determine the number of months (outputs)
    num_months = predictions.shape[1]
    yearly_preds = []  # Empty list to hold predictions for each month

    
    # Process each month separately
    for month in range(num_months):
        logging.info(f"Processing month {month+1} for year {year}.")
        
        # Extract predictions for the current month
        month_pred = predictions[:, month]
        
        # Reshape the flat predictions to the spatial grid
        month_pred = reshape_with_padding(month_pred, (lat_dim, lon_dim))

        # Reverse ground truth preprocessing
        if log_transform:
            logging.info("Applying inverse log transformation.")
            month_pred = inverse_log_transform(month_pred)
        if normalization_applied:
            logging.info("Denormalizing predictions.")
            month_pred = denormalize_data(month_pred, lower_bound, upper_bound)

        # Reshape predictions for comparison
        month_pred = month_pred.reshape((lat_dim, lon_dim))
        logging.info(f"Shape of predictions (Dask): {month_pred.shape}")
        
        # Bias Correction
        # simple bias correction 
        month_pred_corrected = month_pred + abs(bias_correction)
        logging.info(f"Bias Correction: {abs(bias_correction)}")
        
        # Apply water mask 
        water_mask = water_mask.astype(np.float32)
        month_pred = month_pred * water_mask
        
        month_pred = np.clip(month_pred, lower_bound, upper_bound)
        logging.info("Finished bias correction. Predictions clipped to range.")
        
        # Append each monthly prediction
        yearly_preds.append(month_pred)  

        # Save predictions for the current year
        month_pred_file_path = os.path.join(output_folder_path, f"predicted_LAI_{year}_{month+1}_{region}_{start_year}_{end_year}_{model_timestamp}.zarr")
        zarr.save(month_pred_file_path, month_pred)
        logging.info(f"Predictions for month {month+1} of year {year} saved at: {month_pred_file_path}")
        
        # Optionally, log or plot the results
        predicted_plot_path = os.path.join(output_folder_path, f"predicted_LAI_{year}_{month+1}_{region}.png")
        plot_lai(month_pred,
            f"Predicted LAI {month+1} {year} - {region}",
            predicted_plot_path,
            vmin=lower_bound,   # Fixed lower bound (e.g., 0)
            vmax=upper_bound)   # Fixed upper bound (e.g., 7)
        
        # Calculate error metrics if ground truth is available
        ground_truth_path = os.path.join(save_data_path, f'LAI_{year}_{region}.zarr')
        if os.path.exists(ground_truth_path):
            try:
                ground_truth_ds = xr.open_zarr(ground_truth_path)
                ground_truth_lai = ground_truth_ds['LAI'].values
                logging.info(f"Loaded ground truth LAI as a dataset.")
                ground_truth_ds.close()  # Close the dataset to free up resources
            except (zarr.errors.ContainsArrayError, KeyError):
                logging.info("Falling back to loading as a Zarr array.")
                ground_truth_lai = da.from_zarr(ground_truth_path).compute()  # Directly load Zarr array
                logging.info(f"Shape of LAI ground truth (from Zarr array): {ground_truth_lai.shape}")

            # # Reshape ground truth data using the same function as training
            # logging.info(f"Original ground truth shape: {ground_truth_lai.shape}")
            # ground_truth_lai = reshape_with_padding(ground_truth_lai, (lat_dim, lon_dim))
            # logging.info(f"Reshaped ground truth LAI to: {ground_truth_lai.shape}")

            # Extract the ground truth for the current month from the stack
            gt_month = reshape_with_padding(ground_truth_lai[:, month], (lat_dim, lon_dim))
            logging.info(f"Ground truth for month {month+1} reshaped to: {gt_month.shape}")
    
            # # Final check to confirm consistency
            # if ground_truth_lai.shape != (lat_dim, lon_dim):
            #     logging.error(f"Unexpected shape after reshaping: {ground_truth_lai.shape}.")
            #     raise ValueError(f"Ground truth LAI for year {year} does not have the expected shape ({lat_dim}, {lon_dim}).")

                
            mse = mean_squared_error(gt_month, month_pred)
            mae = mean_absolute_error(gt_month, month_pred)
            r2 = r2_score(gt_month, month_pred)
            
            logging.info(f"Year {year}, Month {month+1} - MSE: {mse:.4f}, MAE: {mae:.4f}, R²: {r2:.4f}")
            ground_truth_file_path = os.path.join(output_folder_path, f"ground_truth_LAI_{year}_month{month+1}_{region}_{start_year}_{end_year}_{model_timestamp}.png")

            plot_lai(gt_month,
                    f"Ground Truth LAI {year}, Month {month+1} - {region}",
                    ground_truth_file_path,
                    vmin=lower_bound,
                    vmax=upper_bound)
            # Plot error map of ground truth and predictions
            error_map = month_pred - gt_month  # Difference between predictions and ground truth
            absolute_error_map = np.abs(error_map)  # Absolute error
            # Save and plot error maps
            error_file_path = os.path.join(output_folder_path, f"error_map_{year}__month{month+1}_{region}_{start_year}_{end_year}_{model_timestamp}.png")
            absolute_error_file_path = os.path.join(output_folder_path, f"absolute_error_map_{year}_month{month+1}_{region}_{start_year}_{end_year}_{model_timestamp}.png")
            # Plot error map (raw difference)
            plot_error_map(error_map, f"Error Map {year}, Month {month+1} - {region}",
                            error_file_path,
                            vmin=-upper_bound,
                            vmax=upper_bound)
            # Plot absolute error map
            plot_error_map(absolute_error_map, f"Absolute Error Map {year}, Month {month+1} - {region}",
                        absolute_error_file_path,
                            vmin=lower_bound,
                            vmax=upper_bound) 
        else:
            logging.info(f"No ground truth available for year {year}. Skipping error metrics.")

    # After processing all months for the current year:
    # Stack the monthly predictions along a new axis (shape: (12, lat, lon))
    stacked_year_preds = np.stack(yearly_preds, axis=0)
    # Prepare the auto_regressive data: 
    if not autoregressive:        
        previous_year_predictions = np.zeros((lat_dim, lon_dim), dtype=np.float32)
        logging.info("Autoregressive feature is disabled. Overwriting the previous_year_prediction with Zero")
    else :
    # Update the variable to hold predictions for the next year's placeholder
        if autoregressive_smoothing:
            previous_year_predictions = gaussian_filter(stacked_year_preds, sigma=10)
            logging.info("Applied Gaussian smoothing to autoregressive to last year prediction.")
        else: 
            previous_year_predictions = stacked_year_preds
            logging.info("Prediction is passed as autoregressive to last year prediction without any smoothing")

    # Reshape to (lat*lon, 12) to match training autoregressive feature shape
    AR_flat = previous_year_predictions.reshape(12, -1).T  # shape: (lat*lon, 12)
    
    # Clean up memory 
    del features, water_mask, predictions  # Remove references to large arrays if they're no longer needed
    gc.collect()  # Force garbage collection to free memory
        
logging.info(f"All outputs have been saved to: {output_folder_path}")
logging.info("LAI prediction and visualization completed successfully.")
