import xarray as xr
import numpy as np
import os
import logging
import argparse
from datetime import datetime
import zarr
import json
import gc
import matplotlib.pyplot as plt
import re
from glob import glob

# Import compression modules
import numcodecs
from numcodecs import Blosc

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

gc.collect()

def extract_timestamp(file_path):
    """
    Extracts the timestamp from the file path.

    Parameters:
    - file_path: Path of the file.

    Returns:
    - Extracted timestamp as a string, or None if no timestamp is found.
    """
    match = re.search(r'_(\d{8}_\d{6})_LAI', file_path)
    if match:
        return match.group(1)
    return None

def construct_prediction_path_with_wildcards(base_dir, window, start_year, end_year):
    """
    Constructs the path with a wildcard for the timestamp.

    Parameters:
    - base_dir: Base directory containing all model outputs.
    - window: Window ID (e.g., W021).
    - start_year: Start year for the range.
    - end_year: End year for the range.

    Returns:
    - Wildcard path for prediction files.
    """
    return os.path.join(base_dir, f"{window}_{start_year}_{end_year}_*_LAI")

def plot_global_dataset(data_array, year, output_dir):
    """
    Plot the global dataset and save the figure.

    Parameters:
    - data_array: xarray.DataArray containing the stitched predictions.
    - year: The year of the data being plotted.
    - output_dir: Directory where the plot will be saved.
    """
    plt.figure(figsize=(12, 6))
    plt.title(f"Stitched Global LAI Predictions for {year}")
    plt.pcolormesh(data_array['lon'], data_array['lat'], data_array, cmap="viridis", shading="auto")
    plt.colorbar(label="LAI Value")
    plt.xlabel("Longitude")
    plt.ylabel("Latitude")

    # Save the figure
    plot_file = os.path.join(output_dir, f"stitched_LAI_{year}.png")
    plt.savefig(plot_file, dpi=300, bbox_inches="tight")
    plt.close()
    logging.info(f"Saved plot for year {year} at {plot_file}")
    
# Parse command-line arguments
parser = argparse.ArgumentParser(description="Combine specific window predictions into a stitched dataset.")
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("--start_year_training", type=int, required=True, help="Training year")
parser.add_argument("--end_year_training", type=int, required=True, help="Training years")
parser.add_argument("--start_year_inference", type=int, required=True, help="Inference year")
parser.add_argument("--end_year_inference", type=int, required=True, help="Inference years")
parser.add_argument("--windows", type=str, nargs='+', required=True, help="List of window IDs to process (e.g., W004 W005)")
parser.add_argument("--predictions_dir", type=str, required=True, help="Directory containing all window predictions")
parser.add_argument("--output_dir", type=str, required=True, help="Directory to save combined datasets")
parser.add_argument("--region", type=str, required=True, help="Region code (e.g., 'ES')")
parser.add_argument("--regions_file", type=str, required=True, help="Path to the regions JSON file")
parser.add_argument("--input_data_path", required=True, help="Path to the input data")
parser.add_argument("--preprocessing_goal", type=str, required=True, help="Goal of preprocessing (e.g., training or inference)")
parser.add_argument("--save_data_path", type=str, required=True, help="Path to save data outputs")
parser.add_argument("--window_timestamps", type=str, required=False, help="Thescript could accept a an exact time stamp")
args = parser.parse_args()
args = parser.parse_args()

# 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
start_year_inference = args.start_year_inference
end_year_inference = args.end_year_inference
windows = args.windows
predictions_dir = args.predictions_dir
output_dir = args.output_dir
regions_file = args.regions_file
region = args.region
input_data_path = args.input_data_path

# Configure logging
timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
logs_path = os.path.join(os.getcwd(), "Logs")
os.makedirs(logs_path, exist_ok=True)

logging.basicConfig(
    level=logging.INFO, 
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[
        logging.FileHandler(os.path.join(logs_path, f"window_combiner_{timestamp}.log")),
        logging.StreamHandler()
    ]
)

# Log all the passed arguments
logging.info("Starting the script with the following parameters:")
logging.info(f"Start Year: {start_year}")
logging.info(f"End Year: {end_year}")
logging.info(f"Start Year (Training): {start_year_training}")
logging.info(f"End Year (Training): {end_year_training}")
logging.info(f"Start Year (Inference): {start_year_inference}")
logging.info(f"End Year (Inference): {end_year_inference}")
logging.info(f"Windows to process: {windows}")
logging.info(f"Predictions Directory: {predictions_dir}")
logging.info(f"Output Directory: {output_dir}")
logging.info(f"Region Code: {region}")
logging.info(f"Regions File: {regions_file}")
logging.info(f"Input Data Path: {input_data_path}")

# Check existence of required files and directories
logging.info("Verifying the existence of required files and directories...")

# Check regions file
if not os.path.exists(regions_file):
    logging.error(f"Regions file not found at: {regions_file}")
    raise FileNotFoundError(f"Regions file not found at: {regions_file}")
logging.info(f"Regions file exists: {regions_file}")

# Check predictions directory
if not os.path.isdir(predictions_dir):
    logging.error(f"Predictions directory not found or is not a directory: {predictions_dir}")
    raise NotADirectoryError(f"Predictions directory not found or is not a directory: {predictions_dir}")
logging.info(f"Predictions directory exists: {predictions_dir}")

# Check input data path
if not os.path.isdir(input_data_path):
    logging.error(f"Input data path not found or is not a directory: {input_data_path}")
    raise NotADirectoryError(f"Input data path not found or is not a directory: {input_data_path}")
logging.info(f"Input data path exists: {input_data_path}")

# Load and validate regions JSON
logging.info("Loading and validating regions JSON file...")
try:
    with open(regions_file, "r") as f:
        regions_data = json.load(f)
    if not isinstance(regions_data, dict):
        raise ValueError("Regions file does not contain a valid JSON object.")
    logging.info(f"Regions JSON file loaded successfully with {len(regions_data)} regions.")
except json.JSONDecodeError:
    logging.error("Failed to parse regions JSON file.")
    raise

### Process Ground Truth Dataset ###
# Load the ground truth dataset
year = 2015
ground_truth_path = os.path.join(input_data_path, f"confess_lai/CONFESS_LAI_{year}0715_1km.nc")
logging.info(f"Loading ground truth data from: {ground_truth_path}")

if not os.path.exists(ground_truth_path):
    logging.error(f"Ground truth file not found at: {ground_truth_path}")
    raise FileNotFoundError(f"Ground truth file not found at: {ground_truth_path}")


# Open the ground truth dataset
try:
    ground_truth_ds = xr.open_dataset(ground_truth_path)
    logging.info("Ground truth dataset loaded successfully.")
    if "lat" not in ground_truth_ds or "lon" not in ground_truth_ds:
        logging.error("Ground truth dataset missing latitude or longitude variables.")
        raise KeyError("Latitude or longitude missing in ground truth dataset.")
    logging.info(f"Ground truth dataset contains {len(ground_truth_ds['lat'])} latitudes and {len(ground_truth_ds['lon'])} longitudes.")
except Exception as e:
    logging.error(f"Failed to load ground truth dataset: {e}")
    raise

# Check prediction files and select the most recent one
logging.info("Checking prediction files and selecting the most recent timestamp for each window...")
for year in range(start_year, end_year + 1):
    logging.info(f"Checking predictions for year {year}...")
    for window in windows:
        # Construct the directory path dynamically with wildcards
        prediction_dir = construct_prediction_path_with_wildcards(predictions_dir, window, start_year_inference, end_year_inference)
        prediction_pattern = os.path.join(prediction_dir, f"predicted_LAI_{year}_{window}_*.zarr")

        # Use glob to find matching files
        matching_files = glob(prediction_pattern)

        if matching_files:
            if len(matching_files) > 1:
                # Extract timestamps and sort files
                timestamps_and_files = [
                    (extract_timestamp(file), file) for file in matching_files if extract_timestamp(file)
                ]
                if not timestamps_and_files:
                    logging.warning(f"No valid timestamps found for multiple files for window {window} and year {year}. Skipping...")
                    continue

                # Sort by timestamp in descending order and pick the latest
                timestamps_and_files.sort(key=lambda x: x[0], reverse=True)
                selected_file = timestamps_and_files[0][1]
                logging.info(f"Selected the most recent prediction file for window {window}, year {year}: {selected_file}")
            else:
                # Single match, use it directly
                selected_file = matching_files[0]
                logging.info(f"Matched prediction file for {window}, year {year}: {selected_file}")
        else:
            logging.warning(f"No prediction file found matching pattern: {prediction_pattern}")
            
# Extract latitude and longitude
if "lat" in ground_truth_ds and "lon" in ground_truth_ds:
    latitudes = ground_truth_ds["lat"].values
    longitudes = ground_truth_ds["lon"].values
    logging.info(f"Ground truth data contains {len(latitudes)} latitudes and {len(longitudes)} longitudes.")
else:
    logging.error("Latitude and longitude variables not found in the ground truth data.")
    raise KeyError("Missing latitude or longitude in the ground truth data.")

# Initialize the stitched dataset grid
logging.info("Creating placeholder for the stitched dataset.")

# Determine the bounds of the stitched grid
combined_lats = []
combined_lons = []

logging.info("Creating placeholder aligned with the ground truth dataset.")

combined_lats = ground_truth_ds["lat"].values
combined_lons = ground_truth_ds["lon"].values

# Create an empty placeholder for predictions with NaN values
stitched_predictions = np.full((len(combined_lats), len(combined_lons)), np.nan, dtype=np.float32)

logging.info(f"Created placeholder aligned with ground truth dimensions: {stitched_predictions.shape}.")


# Create an empty placeholder for predictions with NaN values
stitched_predictions = np.full((len(combined_lats), len(combined_lons)), np.nan, dtype=np.float32)

logging.info(f"Created placeholder with dimensions {stitched_predictions.shape}.")

# Process each window and load the predictions
logging.info("Loading prediction files and updating the stitched dataset.")
for year in range(start_year, end_year + 1):
    logging.info(f"Processing year {year}...")

    for window in windows:
        # Check if the window exists in the regions file
        if window not in regions_data:
            logging.warning(f"Window {window} not found in regions file. Skipping...")
            continue

        # Extract region details
        details = regions_data[window]
        lon_min, lon_max = details["lon_slice"]
        lat_min, lat_max = details["lat_slice"]

        logging.info(f"Processing window {window} with details: {details}")

        # Construct prediction file path and pattern
        prediction_dir = construct_prediction_path_with_wildcards(
            predictions_dir, window, start_year_inference, end_year_inference
        )
        prediction_pattern = os.path.join(
            prediction_dir, f"predicted_LAI_{year}_{window}_*.zarr"
        )

        # Find matching files
        matching_files = glob(prediction_pattern)
        if not matching_files:
            logging.warning(f"No prediction files found for window {window}, year {year}. Pattern: {prediction_pattern}")
            continue

        # Select the most recent file if multiple matches are found
        selected_file = max(matching_files, key=extract_timestamp)
        logging.info(f"Selected prediction file for window {window}, year {year}: {selected_file}")

        # Load and process the prediction file
        try:
            # Open Zarr data as unstructured
            zarr_data = zarr.open(selected_file, mode="r")

            # Extract prediction array
            window_predictions = zarr_data[:]

            # Calculate prediction resolution dynamically
            pred_resolution_lat = abs(combined_lats[1] - combined_lats[0])  # Use the ground truth resolution
            pred_resolution_lon = abs(combined_lons[1] - combined_lons[0])

            # Generate indices for the combined grid
            lat_indices = np.where((combined_lats >= lat_min) & (combined_lats <= lat_max))[0]
            lon_indices = np.where((combined_lons >= lon_min) & (combined_lons <= lon_max))[0]

            # Validate the shape of the predictions
            if window_predictions.shape == (len(lat_indices), len(lon_indices)):
                
                # Flip the predictions vertically
                window_predictions = np.flipud(window_predictions)

                # Update the stitched dataset
                stitched_predictions[np.ix_(lat_indices, lon_indices)] = window_predictions
                logging.info(f"Updated stitched dataset for window {window}, year {year}.")
            else:
                logging.error(f"Dimension mismatch for window {window}, year {year}. "
                              f"Prediction shape: {window_predictions.shape}, "
                              f"Grid indices: lat ({len(lat_indices)}), lon ({len(lon_indices)})")

        except Exception as e:
            logging.error(f"Failed to load or process file {selected_file} for window {window}, year {year}: {e}")

    # Save combined predictions to a Zarr file
    stitched_dataset = xr.DataArray(
        stitched_predictions, coords=[combined_lats, combined_lons], dims=["lat", "lon"]
    )
    # Set a name for the DataArray so the encoding applies (do not change min/max)
    stitched_dataset.name = "LAI"
    os.makedirs(output_dir, exist_ok=True)
    output_path = os.path.join(output_dir, f"stitched_LAI_{year}.zarr")
    
    # Define a reasonable chunk shape (roughly 1/10th of the grid in each dimension)
    chunk_shape = (max(len(combined_lats) // 10, 1), max(len(combined_lons) // 10, 1))
    encoding = {"LAI": {"compressor": fast_compressor, "chunks": chunk_shape}}
    
    stitched_dataset.to_zarr(output_path, mode="w", encoding=encoding)
    logging.info(f"Saved stitched predictions for year {year} at {output_path}")

    # Plot the global dataset
    plot_global_dataset(stitched_dataset, year, output_dir)
    gc.collect()

logging.info("Window stitching process completed successfully.")