import xarray as xr
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import os
import logging
import argparse
from datetime import datetime
import zarr
import shutil
import json
import sys  # To get the size of the objects
import gc
import numcodecs
import scipy  
from numcodecs import Blosc
gc.collect()

# Define a fast compressor: LZ4 is very fast with a low compression level
fast_compressor = Blosc(cname='lz4', clevel=1, shuffle=Blosc.BITSHUFFLE)

def plot_data(data, title, filename):
    plt.figure(figsize=(10, 6))
    plt.imshow(np.flipud(data), cmap='viridis', origin='upper')  # Flip the data upside down to match the standard orientation
    plt.colorbar(label='Value')
    plt.title(title)
    plt.xlabel('Longitude')
    plt.ylabel('Latitude')
    plt.savefig(filename)
    plt.close()

def get_size(obj, name):
    size = sys.getsizeof(obj)
    logging.info(f"Size of {name}: {size} bytes")
    return size

# --- NEW HELPER FUNCTION FOR ALIGNMENT ---
def align_to_ground_truth(ds, gt_lon, gt_lat):
    # Use interpolation to force ds to have the same lon/lat coordinates as the ground truth.
    ds_aligned = ds.interp(lon=gt_lon, lat=gt_lat, method='nearest')
    return ds_aligned

# # (Optional) You can remove or comment out the old fix_feature_length function.
# def fix_feature_length(feature_array, target_length):
#     """
#     (Deprecated) Ensure that the 1D feature_array has exactly target_length elements.
#     Padding/trimming flattened arrays is no longer recommended as it destroys spatial alignment.
#     """
#     current_length = feature_array.size
#     if current_length < target_length:
#         padding = target_length - current_length
#         fixed_array = np.pad(feature_array, (0, padding), mode='constant', constant_values=0)
#         logging.info(f"Padded feature from {current_length} to {target_length} elements.")
#     elif current_length > target_length:
#         fixed_array = feature_array[:target_length]
#         logging.info(f"Trimmed feature from {current_length} to {target_length} elements.")
#     else:
#         fixed_array = feature_array
#     return fixed_array

# Parse command-line arguments
parser = argparse.ArgumentParser(description="Process land use data and generate yearly datasets.")
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("--month", type=int, default=7, help="Month used for the ground truth dataset in CONFESS_LAI file name (e.g., 7 for July)")
parser.add_argument("--region", type=str, required=True, help="Region is the window that is going to be processed)")
parser.add_argument("--include_luh2", action='store_true', help="Include LUH2 dataset in the processed data")
parser.add_argument("--include_hilda", action='store_true', help="Include HILDA+ dataset in the processed data")
parser.add_argument("--input_data_override", type=str, help="Path to override the default input data directory")
parser.add_argument("--inspection", action='store_true', help="Enable inspection mode with additional plotting")
parser.add_argument("--overwrite", action='store_true', help="Overwrite existing processed data if it exists")
parser.add_argument("--dtype_ground_truth", type=str, default="float32", help="Data type for ground truth data (e.g., int16, float32)")
parser.add_argument("--dtype_luh2", type=str, default="float32", help="Data type for LUH2 data (e.g., int16, float32)")
parser.add_argument("--dtype_hilda", type=str, default="float32", help="Data type for HILDA+ data (e.g., int16, float32)")
parser.add_argument("--dtype_esa_cci", type=str, default="float32", help="Data type for ESA CCI data (e.g., int16, float32)")
parser.add_argument("--autoregressive", action='store_true', help="Include autoregressive feature from ground truth data (year +1)")
parser.add_argument("--include_esa_cci", action='store_true', help="Include ESA CCI dataset in the processed data")
parser.add_argument("--create_watermask", action='store_true', help="Create water mask - It needs HILDA+ to be included")
parser.add_argument("--save_data_path", required=True, help="Path to the place to save the data")
parser.add_argument("--input_data_path", required=True, help="Path to the input data")
parser.add_argument("--region_data_path", required=True, help="Path to the region json")
parser.add_argument(
    "--preprocessing_goal", 
    type=str, 
    required=True, 
    choices=["training", "inference"], 
    help="Specify the goal of preprocessing: 'training' or 'inference'."
)

args = parser.parse_args()
start_year = args.start_year
end_year = args.end_year
ground_truth_month = args.month
region = args.region
include_luh2 = args.include_luh2
include_hilda = args.include_hilda
include_esa_cci = args.include_esa_cci
input_data_override = args.input_data_override
inspection_mode = args.inspection
overwrite_existing = args.overwrite
dtype_ground_truth = args.dtype_ground_truth
dtype_luh2 = args.dtype_luh2
dtype_hilda = args.dtype_hilda
dtype_esa_cci = args.dtype_esa_cci
autoregressive = args.autoregressive
create_watermask = args.create_watermask
save_data_path = args.save_data_path
default_data_path = args.input_data_path
regions_file_path = args.region_data_path
preprocessing_goal = args.preprocessing_goal 

# Configure logging
timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')

# Create a directory for logs
current_path = os.getcwd()
logs_path = os.path.join(current_path, "Logs")
os.makedirs(logs_path, exist_ok=True)

log_file_name = f"preprocessing_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("Preprocessing started.")
logging.info(f"Arguments received: start_year={start_year}, end_year={end_year}, region={region}, include_luh2={include_luh2}, include_hilda={include_hilda}, input_data_override={input_data_override}, inspection_mode={inspection_mode}, overwrite_existing={overwrite_existing}, dtype_ground_truth={dtype_ground_truth}, dtype_luh2={dtype_luh2}, dtype_hilda={dtype_hilda}, ground_truth_month={ground_truth_month}")

# Define default paths if input_data_override is not provided
input_data_path = input_data_override if input_data_override else default_data_path
logging.info(f"Using input data path: {input_data_path}")

# Check if the regions file exists
if not os.path.exists(regions_file_path):
    logging.error(f"Regions JSON file not found at: {regions_file_path}")
    sys.exit(1)

# Load the region information from the JSON file
with open(regions_file_path, 'r') as regions_file:
    regions_data = json.load(regions_file)

# Check if the specified region exists in the regions data
if region not in regions_data:
    logging.error(f"Region {region} not found in regions JSON file.")
    sys.exit(1)

# Extract latitude and longitude slices for the given region
region_lon_slice = regions_data[region]["lon_slice"]
region_lat_slice = regions_data[region]["lat_slice"]

logging.info(f"Set slicing for region {region}: lon_slice={region_lon_slice}, lat_slice={region_lat_slice}")

# Create a main directory for plots if it doesn't exist
plots_main_path = os.path.join(save_data_path, 'plots')
os.makedirs(plots_main_path, exist_ok=True)
logging.info(f"Plots main directory set at: {plots_main_path}")

# Create a region-specific subdirectory for plots
plots_path = os.path.join(plots_main_path, f'{region}_{start_year}_{end_year}_{preprocessing_goal}')
os.makedirs(plots_path, exist_ok=True)
logging.info(f"Plots subdirectory set at: {plots_path}")

# Initialize a dictionary to store metadata
metadata = {
    "start_year": start_year,
    "end_year": end_year,
    "region": region,
    "include_luh2": include_luh2,
    "include_hilda": include_hilda,
    "include_esa_cci": include_esa_cci,
    "autoregressive": autoregressive,
    "data_path": input_data_path,
    "processing_date": datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
    "dtype_ground_truth": dtype_ground_truth,
    "dtype_luh2": dtype_luh2,
    "dtype_hilda": dtype_hilda,
    "dtype_esa_cci": dtype_esa_cci,
    "ground_truth_month": ground_truth_month
}

# Initialize a variable to hold the last available LAI dataset. This will be used in case there is no LAI ground truth
last_available_LAI = None

# Backward iteration
for year in range(end_year, start_year - 1, -1):
    logging.info(f"Processing year: {year}")
    
    # Paths to the LUH2 datasets
    luh2_data_paths = [os.path.join(input_data_path, f"luh2class/{dataset}_{year}_0.0089deg.nc") for dataset in [
        "primf", "primn", "secdf", "secdn", "urban", "c3ann", "c4ann", "c3per", "c4per", "c3nfx", "pastr", "range", "secmb", "secma"
    ]]
    logging.info(f"LUH2 data paths: {luh2_data_paths}")

    # Path to the HILDA+ dataset
    hilda_dataset_path = os.path.join(input_data_path, f"hilda/hildaplus_GLOB-1-0-f_states_{year}_0.0089deg.nc")
    logging.info(f"HILDA+ dataset path: {hilda_dataset_path}")

    # Path to the Land Cover dataset    
    esa_cci_path = os.path.join(input_data_path, f"lccsclass/lccsclass_{year}_0.0089deg.nc")
    logging.info(f"Loading Land Cover data from: {esa_cci_path}")

    ### Process Ground Truth Dataset ###
    if preprocessing_goal == "training":
        # Load the ground truth dataset (using the parameterized month)
        ground_truth_path = os.path.join(input_data_path, f"confess_lai/CONFESS_LAI_{year}{ground_truth_month:02d}15_1km.nc")
        logging.info(f"Loading ground truth data from: {ground_truth_path}")
    elif preprocessing_goal == "inference":
        # READ FIXED CONFESS LAI example to extract necessary metadata 
        fixed_year = 2015
        ground_truth_path = os.path.join(input_data_path, f"confess_lai/CONFESS_LAI_{fixed_year}{ground_truth_month:02d}15_1km.nc")
        logging.info(f"Loading ground truth data from: {ground_truth_path}")

    # Open the ground truth dataset without slicing
    ground_truth_ds = xr.open_dataset(ground_truth_path)
    # Check coordinate names and rename if necessary
    if 'latitude' in ground_truth_ds.coords:
        ground_truth_ds = ground_truth_ds.rename({'latitude': 'lat'})
    if 'longitude' in ground_truth_ds.coords:
        ground_truth_ds = ground_truth_ds.rename({'longitude': 'lon'})

    # Reverse latitude and longitude if necessary to match the orientation of other datasets
    if ground_truth_ds['lat'][0] > ground_truth_ds['lat'][-1]:
        ground_truth_ds = ground_truth_ds.reindex(lat=ground_truth_ds['lat'][::-1])
        logging.info("Reversed latitude axis for ground truth data.")

    if ground_truth_ds['lon'][0] > ground_truth_ds['lon'][-1]:
        ground_truth_ds = ground_truth_ds.reindex(lon=ground_truth_ds['lon'][::-1])
        logging.info("Reversed longitude axis for ground truth data.")

    # Apply slicing using region slices (these define the desired spatial window)
    lat_slice_gt = slice(region_lat_slice[0], region_lat_slice[1])
    lon_slice_gt = slice(region_lon_slice[0], region_lon_slice[1])
    logging.info(f"Adjusted slicing for Ground Truth dataset: lon_slice={lon_slice_gt}, lat_slice={lat_slice_gt}")
    ground_truth_ds = ground_truth_ds.sel(lat=lat_slice_gt, lon=lon_slice_gt)
    
    # Fill NaNs and convert to specified dtype
    ground_truth_ds['LAI'] = ground_truth_ds['LAI'].fillna(0).astype(dtype_ground_truth)
    logging.info(f"Ground truth dataset loaded and sliced with dimensions: {ground_truth_ds.dims}")
    ground_truth_ds['LAI'] = ground_truth_ds['LAI'].astype(dtype_ground_truth)
    get_size(ground_truth_ds['LAI'].values, 'Ground Truth Data')

    # Get actual dimensions for reshaping using sizes
    lat_dim = ground_truth_ds.sizes['lat']
    lon_dim = ground_truth_ds.sizes['lon']
    logging.info(f"lat_dim={lat_dim}, lon_dim={lon_dim}")
    
    expected_flat_size = lat_dim * lon_dim
    actual_flat_size = ground_truth_ds['LAI'].isel(time=0).values.flatten().size
    logging.info(f"DEBUG: Expected flattened LAI size: {expected_flat_size}, Actual flattened LAI size: {actual_flat_size}")

    # Save these dimensions to a file
    dims_path = os.path.join(save_data_path, f'dims_{start_year}_{end_year}_{region}.npy')
    np.save(dims_path, np.array([lat_dim, lon_dim]))
    logging.info(f"Saved dimensions to: {dims_path}")

    # Update metadata with dimensions
    metadata["lat_dim"] = int(lat_dim)
    metadata["lon_dim"] = int(lon_dim)

    # Transpose to ensure dimensions are in the correct order
    ground_truth_ds['LAI'] = ground_truth_ds['LAI'].transpose('lat', 'lon', 'time')
    logging.info("Ground truth data transposed.")

    # Select the first time step
    LAI_first_time = ground_truth_ds['LAI'].isel(time=0)
    logging.info(f"Selected the first time step for LAI data with shape: {LAI_first_time.shape}")
    logging.info(f"DEBUG: Flattened first time step LAI size: {LAI_first_time.values.flatten().size}")

    # IMPORTANT: Save the ground truth coordinate arrays for alignment of other datasets.
    gt_lon = ground_truth_ds.lon.values
    gt_lat = ground_truth_ds.lat.values

    # Plot the ground truth data only if inspection mode is enabled
    if inspection_mode and preprocessing_goal == "training":
        logging.info("Plotting ground truth data for inspection mode.")
        plot_data(
            LAI_first_time.values,
            f"LAI - {region} {year}",
            os.path.join(plots_path, f"LAI_{region}_{year}_before.png")
        )
    gc.collect()  # Explicit garbage collection

    # Initialize list for features
    input_features = []
    # Initialize a dictionary to store the feature order
    feature_order = {}
    # Initialize index counter for features
    feature_index = 0
        
    # Process autoregressive feature
    next_year = year + 1
    next_year_ground_truth_path = os.path.join(input_data_path, f"confess_lai/CONFESS_LAI_{next_year}{ground_truth_month:02d}15_1km.nc")

    if os.path.exists(next_year_ground_truth_path):
        logging.info(f"Loading autoregressive ground truth data for year +1 from: {next_year_ground_truth_path}")

        next_year_ground_truth_ds = xr.open_dataset(next_year_ground_truth_path)

        if 'latitude' in next_year_ground_truth_ds.coords:
            next_year_ground_truth_ds = next_year_ground_truth_ds.rename({'latitude': 'lat'})
        if 'longitude' in next_year_ground_truth_ds.coords:
            next_year_ground_truth_ds = next_year_ground_truth_ds.rename({'longitude': 'lon'})

        if next_year_ground_truth_ds['lat'][0] > next_year_ground_truth_ds['lat'][-1]:
            next_year_ground_truth_ds = next_year_ground_truth_ds.reindex(lat=next_year_ground_truth_ds['lat'][::-1])
            logging.info("Reversed latitude axis for autoregressive ground truth data.")

        if next_year_ground_truth_ds['lon'][0] > next_year_ground_truth_ds['lon'][-1]:
            next_year_ground_truth_ds = next_year_ground_truth_ds.reindex(lon=next_year_ground_truth_ds['lon'][::-1])
            logging.info("Reversed longitude axis for autoregressive ground truth data.")

        next_year_ground_truth_ds = next_year_ground_truth_ds.sel(lat=lat_slice_gt, lon=lon_slice_gt)
        next_year_ground_truth_ds['LAI'] = next_year_ground_truth_ds['LAI'].fillna(0).astype(dtype_ground_truth)

        next_year_LAI_first_time = next_year_ground_truth_ds['LAI'].isel(time=0)
        
        # Save a copy of the last available ground truth data
        last_available_LAI = next_year_LAI_first_time.values

        # Append the flattened data to input_features (alignment is already ensured by slicing)
        input_features.append(next_year_LAI_first_time.values)
        logging.info("Last known state of LAI (year +1) added to input features.")
        
        if inspection_mode:
            plot_data(
                next_year_LAI_first_time.values,
                f"Last known state of LAI - {region} {next_year}",
                os.path.join(plots_path, f"Last_known_LAI_{region}_{next_year}.png")
            )
    else:
        logging.error(f"Ground truth for year {next_year} not found. Reusing the last available LAI state.")
        if last_available_LAI is not None:
            input_features.append(last_available_LAI)
            logging.info("Reused the last available LAI state for input features.")
        else:
            logging.error("No previous LAI data available to reuse. Exiting.")
            sys.exit(1)
        
    # Update the feature list: the first feature is always the autoregressive LAI state.
    feature_order[feature_index] = "Autoregressive (Last known LAI state)"
    feature_index += 1
    
    # Close the ground truth datasets to free memory
    ground_truth_ds.close()
    if 'next_year_ground_truth_ds' in locals():
        next_year_ground_truth_ds.close()
    gc.collect() 
                
    ### Process LUH2 Dataset ###
    if include_luh2:
        luh2_ds = xr.open_mfdataset(luh2_data_paths, combine='by_coords')

        if 'latitude' in luh2_ds.coords:
            luh2_ds = luh2_ds.rename({'latitude': 'lat'})
        if 'longitude' in luh2_ds.coords:
            luh2_ds = luh2_ds.rename({'longitude': 'lon'})

        # Instead of slicing by region manually, align LUH2 to ground truth coordinates.
        luh2_ds = align_to_ground_truth(luh2_ds, gt_lon, gt_lat)
        logging.info(f"LUH2 datasets aligned to ground truth coordinates with dimensions: {luh2_ds.dims}")

        for var in luh2_ds.data_vars:
            luh2_ds[var] = luh2_ds[var].astype(dtype_luh2)
        logging.info(f"LUH2 data converted to dtype: {dtype_luh2}")
        get_size(luh2_ds, 'LUH2 Data')

        # Transpose each variable if necessary (they should now match ground truth grid)
        for var in luh2_ds.data_vars:
            luh2_ds[var] = luh2_ds[var].transpose('lat', 'lon')
            logging.info(f"DEBUG: LUH2 variable '{var}' shape after transposition: {luh2_ds[var].shape}")
        logging.info("LUH2 data transposed.")

        if inspection_mode:
            for var in luh2_ds.data_vars:
                logging.info(f"Plotting LUH2 data for variable: {var}")
                plot_data(
                    luh2_ds[var].values,
                    f"LUH2 {var} - {region} {year}",
                    os.path.join(plots_path, f"luh2_{var}_{region}_{year}_before.png")
                )

        for var in luh2_ds.data_vars:
            input_features.append(luh2_ds[var].values)
            feature_order[feature_index] = f"LUH2 feature: {var}"
            feature_index += 1
        logging.info("LUH2 features appended.")
        luh2_ds.close()
        gc.collect()

    ### Process HILDA+ Dataset ###
    if include_hilda or create_watermask:
        hilda_ds = xr.open_dataset(hilda_dataset_path, decode_times=False)

        if 'latitude' in hilda_ds.coords:
            hilda_ds = hilda_ds.rename({'latitude': 'lat'})
        if 'longitude' in hilda_ds.coords:
            hilda_ds = hilda_ds.rename({'longitude': 'lon'})

        # Align HILDA+ dataset to ground truth coordinates.
        hilda_ds = align_to_ground_truth(hilda_ds, gt_lon, gt_lat)
        logging.info(f"HILDA+ dataset aligned to ground truth with dimensions: {hilda_ds.dims}")

        hilda_ds['LULC_states'] = hilda_ds['LULC_states'].astype(dtype_hilda)
        logging.info(f"HILDA+ data converted to dtype: {dtype_hilda}")
        get_size(hilda_ds['LULC_states'].values, 'HILDA+ Data')

        if create_watermask:
            # Create water mask: ocean (code 0) as water, non-water as 1.
            water_mask = xr.where(hilda_ds['LULC_states'] == 0, 0, 1)
            logging.info("Water mask created: ocean=0; non-water=1.")

            water_mask_first_time = water_mask.isel(time=0).transpose('lat', 'lon')
            logging.info(f"Selected first time step for water mask with shape: {water_mask_first_time.shape}")

            water_mask_path = os.path.join(save_data_path, f'water_mask_{region}_{year}.zarr')
            chunk_size = (max(water_mask_first_time.values.flatten().size // 10, 1),)
            zWaterMask = zarr.open_array(
                water_mask_path, mode='w',
                shape=water_mask_first_time.values.flatten().shape,
                chunks=chunk_size,
                dtype=water_mask_first_time.values.dtype,
                compressor=fast_compressor
            )
            zWaterMask[:] = water_mask_first_time.values.flatten()
            logging.info(f"Water mask saved to: {water_mask_path}")

            if inspection_mode:
                logging.info("Plotting water mask for inspection mode.")
                plot_data(
                    water_mask_first_time.values,
                    f"Water Mask - {region} {year}",
                    os.path.join(plots_path, f"water_mask_{region}_{year}.png")
                )
        
        if include_hilda:
            hilda_first_time = hilda_ds['LULC_states'].isel(time=0).transpose('lat', 'lon')
            logging.info(f"Selected first time step for HILDA+ data with shape: {hilda_first_time.shape}")

            if inspection_mode:
                logging.info("Plotting HILDA+ data for inspection mode.")
                plot_data(
                    hilda_first_time.values,
                    f"HILDA+ LULC States - {region} {year}",
                    os.path.join(plots_path, f"hilda_lulc_states_{region}_{year}_before.png")
                )

            input_features.append(hilda_first_time.values)
            feature_order[feature_index] = "HILDA+ Land use"
            feature_index += 1
            logging.info("HILDA+ features appended.")
        hilda_ds.close()
        gc.collect()

    ### Process ESA CCI Dataset ###
    if include_esa_cci:
        esa_cci_ds = xr.open_dataset(esa_cci_path)

        if 'latitude' in esa_cci_ds.coords:
            esa_cci_ds = esa_cci_ds.rename({'latitude': 'lat'})
        if 'longitude' in esa_cci_ds.coords:
            esa_cci_ds = esa_cci_ds.rename({'longitude': 'lon'})

        if esa_cci_ds['lat'][0] > esa_cci_ds['lat'][-1]:
            esa_cci_ds = esa_cci_ds.reindex(lat=esa_cci_ds['lat'][::-1])
            logging.info("Reversed latitude axis for ESA CCI data.")

        if esa_cci_ds['lon'][0] > esa_cci_ds['lon'][-1]:
            esa_cci_ds = esa_cci_ds.reindex(lon=esa_cci_ds['lon'][::-1])
            logging.info("Reversed longitude axis for ESA CCI data.")

        # Align ESA CCI dataset to ground truth coordinates.
        esa_cci_ds = align_to_ground_truth(esa_cci_ds, gt_lon, gt_lat)
        logging.info(f"ESA CCI dataset aligned to ground truth with dimensions: {esa_cci_ds.dims}")

        esa_cci_ds['lccs_class'] = esa_cci_ds['lccs_class'].astype(dtype_esa_cci)
        logging.info(f"ESA CCI data converted to dtype: {dtype_esa_cci}")

        esa_cci_ds['lccs_class'] = esa_cci_ds['lccs_class'].transpose('lat', 'lon', 'time')
        logging.info("ESA CCI data transposed.")

        esa_cci_first_time = esa_cci_ds['lccs_class'].isel(time=0)
        logging.info(f"Selected first time step for ESA CCI data with shape: {esa_cci_first_time.shape}")
        
        input_features.append(esa_cci_first_time.values)
        feature_order[feature_index] = "ESA CCI Land Cover"
        feature_index += 1
        logging.info("ESA CCI feature added to input features.")

        if inspection_mode:
            logging.info("Plotting ESA CCI data for inspection mode.")
            plot_data(
                esa_cci_first_time.values,
                f"ESA CCI Land Cover - {region} {year}",
                os.path.join(plots_path, f"esa_cci_{region}_{year}.png")
            )

        esa_cci_ds.close()
        gc.collect()
    
    # With all datasets aligned to the ground truth grid, they now share the same spatial shape.
    if input_features:
        # Instead of fixing feature lengths by padding/trimming, simply flatten and stack.
        stacked_features = np.vstack([f.flatten() for f in input_features]).T
        logging.info(f"Stacked feature shape for year {year}: {stacked_features.shape}")
    if input_features:
        # Directly flatten and stack – no padding is needed.
        stacked_features = np.vstack([f.flatten() for f in input_features]).T
        logging.info(f"Stacked feature shape for year {year}: {stacked_features.shape}")

        
        # Save water mask (which is already aligned)
        water_mask_flat = water_mask_first_time.values.flatten() if create_watermask else None

        # --- Save Water Mask with Fast Compressor ---
        if water_mask_flat is not None:
            water_mask_path = os.path.join(save_data_path, f'water_mask_{region}_{year}.zarr')
            chunk_size = (max(water_mask_flat.size // 10, 1),)
            zWaterMask = zarr.open_array(
                water_mask_path, mode='w',
                shape=water_mask_flat.shape,
                chunks=chunk_size,
                dtype=water_mask_flat.dtype,
                compressor=fast_compressor
            )
            zWaterMask[:] = water_mask_flat
            logging.info(f"Aligned water mask saved to: {water_mask_path}")

        # --- Save Data for Training or Inference ---
        if preprocessing_goal == "training":
            LAI = LAI_first_time.values.flatten()  # Already aligned with ground truth
            LAI_year_path = os.path.join(save_data_path, f'LAI_{year}_{region}.zarr')
            zLAI = zarr.open_array(
                LAI_year_path, mode='w',
                shape=LAI.shape,
                chunks=(max(LAI.size // 10, 1),),
                dtype=LAI.dtype,
                compressor=fast_compressor
            )
            zLAI[:] = LAI
            logging.info(f"Ground truth for year {year} saved at: {LAI_year_path}")

            features_year_path = os.path.join(save_data_path, f'features_training_{year}_{region}.zarr')
            feature_chunk_size = (max(stacked_features.shape[0] // 10, 1), stacked_features.shape[1])
            zFeatures = zarr.open_array(
                features_year_path, mode='w',
                shape=stacked_features.shape,
                chunks=feature_chunk_size,
                dtype=stacked_features.dtype,
                compressor=fast_compressor
            )
            zFeatures[:] = stacked_features
            logging.info(f"Training features data for year {year} saved at: {features_year_path}")

        if preprocessing_goal == "inference":
            features_year_path = os.path.join(save_data_path, f'features_inference_{year}_{region}.zarr')
            feature_chunk_size = (max(stacked_features.shape[0] // 10, 1), stacked_features.shape[1])
            zFeatures = zarr.open_array(
                features_year_path, mode='w',
                shape=stacked_features.shape,
                chunks=feature_chunk_size,
                dtype=stacked_features.dtype,
                compressor=fast_compressor
            )
            zFeatures[:] = stacked_features
            logging.info(f"Inference features data for year {year} saved at: {features_year_path}")
    else:
        logging.warning(f"No features were created for year {year}.")

    if input_features:
        logging.info(f"Stacked feature shape for year {year}: {stacked_features.shape}")
    gc.collect()

# Save metadata including dimensions, included datasets, and other relevant details
metadata_path = os.path.join(save_data_path, f'metadata_{preprocessing_goal}_{start_year}_{end_year}_{region}.json')
metadata["feature_order"] = feature_order
with open(metadata_path, 'w') as f:
    json.dump(metadata, f)
logging.info(f"Metadata saved to: {metadata_path}")

logging.info("Preprocessing completed successfully.")