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 matplotlib.colors as mcolors

def plot_data(data, title, filename, show_plot=True):
    plt.figure(figsize=(10, 6))
    plt.imshow(data, cmap='viridis', interpolation='nearest')
    plt.colorbar(label='Value')
    plt.title(title)
    plt.savefig(filename)
    if show_plot:
        plt.show()
    plt.close()
    
def plot_data_with_custom_colormap(data, title, filename, show_plot=True):
    # Define a custom color map from your default_color_map
    default_color_map = {
        0: '#ffff64', 10: '#ffff64', 11: '#ffff00', 12: '#aaf0f0', 20: '#dcf064', 30: '#c8c864', 40: '#006400',
        50: '#00a000', 60: '#00a000', 61: '#aac800', 62: '#003c00', 70: '#003c00', 71: '#005000', 72: '#285000',
        80: '#285000', 81: '#286400', 82: '#788200', 90: '#8ca000', 100: '#be9600', 110: '#966400', 120: '#966400',
        121: '#966400', 122: '#ffb432', 130: '#ffdcd2', 140: '#ffebaf', 150: '#ffc864', 151: '#ffd278', 152: '#ffebaf',
        153: '#00785a', 160: '#009678', 170: '#00dc82', 180: '#c31400', 190: '#fff5d7', 200: '#dcdcdc', 201: '#fff5d7',
        202: '#0046c8', 210: '#ffffff', 220: '#ffffff'
    }
    
    # Create a ListedColormap based on your default color map
    cmap = mcolors.ListedColormap([default_color_map[key] for key in sorted(default_color_map.keys())])
    
    # Normalize data to map to the correct color
    norm = mcolors.BoundaryNorm(sorted(default_color_map.keys()), cmap.N)

    # Plot using the custom colormap
    plt.figure(figsize=(10, 6))
    plt.imshow(data, cmap=cmap, norm=norm, interpolation='nearest')
    plt.colorbar(label='Land Cover Class')
    plt.title(title)
    plt.savefig(filename)
    if show_plot:
        plt.show()
    plt.close()


# Argument Parsing
parser = argparse.ArgumentParser(description="Land Use 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("--test_years", type=int, nargs='+', required=True, help="Test years (multiple allowed)")
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 save data and model outputs")
parser.add_argument("--model_timestamp", type=str, required=True, help="Model timestamp from training phase")

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"prediction_log_{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
test_years = args.test_years
region = args.region
save_data_path = args.save_data_path
model_timestamp = args.model_timestamp

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

lat_dim = metadata['lat_dim']
lon_dim = metadata['lon_dim']

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

# Create output folder based on model details
output_folder_name = f"{region}_{start_year}_{end_year}_{model_timestamp}_{os.path.basename(model_path).split('.')[0]}"
output_folder_path = os.path.join(save_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 Class-to-Index Mapping
class_to_index_path = os.path.join(save_data_path, f'class_to_index_{start_year}_{end_year}_{region}_{model_timestamp}.json')
with open(class_to_index_path, 'r') as f:
    class_to_index = json.load(f)

# Convert string keys back to float if needed
class_to_index = {float(k): v for k, v in class_to_index.items()}
index_to_class = {v: k for k, v in class_to_index.items()}  # Reverse the mapping

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

# Iterate over each test year
for year in test_years:
    logging.info(f"Processing data for the year: {year}")

    # Set paths for loading data for each year
    lccs_class_path = os.path.join(save_data_path, f'lccs_class_{year}_{region}.zarr')
    features_path = os.path.join(save_data_path, f'features_{year}_{region}.zarr')

    # Load processed data for the current year
    logging.info("Loading processed data.")
    lccs_class = da.from_zarr(lccs_class_path)
    features = da.from_zarr(features_path)

    # Log the shapes of the loaded data
    logging.info(f"Shape of lccs_class (Dask): {lccs_class.shape}")
    logging.info(f"Shape of features (Dask): {features.shape}")

    # Convert Dask arrays to NumPy arrays
    lccs_class = lccs_class.compute()
    features = features.compute()

    # Log the shapes after computation
    logging.info(f"Shape of lccs_class (NumPy): {lccs_class.shape}")
    logging.info(f"Shape of features (NumPy): {features.shape}")

    # Check if lccs_class has one dimension and reshape using lat_dim and lon_dim
    if len(lccs_class.shape) == 1 and lccs_class.shape[0] == lat_dim * lon_dim:
        lccs_class = lccs_class.reshape((lat_dim, lon_dim))
    else:
        logging.error(f"Unexpected shape for lccs_class: {lccs_class.shape}. Expected a 1D array with size {lat_dim * lon_dim}.")
        raise ValueError(f"lccs_class for year {year} does not have the expected shape (lat_dim * lon_dim).")

    # Make predictions
    logging.info("Making predictions.")
    predictions = model.predict(features)

    # Adjust Labels Back to Original for Plotting
    predictions = np.array([index_to_class.get(idx, np.nan) for idx in predictions])
    predictions = predictions.reshape((lat_dim, lon_dim))

    # Check Data Length and Reshape for Plotting
    expected_size_per_year = lat_dim * lon_dim

    if len(predictions.flatten()) == expected_size_per_year and len(lccs_class.flatten()) == expected_size_per_year:
        # File paths for predictions and ground truth
        predicted_file_path_default = os.path.join(output_folder_path, f"predicted_lccs_classes_{year}_{region}_{model_timestamp}_default.png")
        ground_truth_file_path_default = os.path.join(output_folder_path, f"ground_truth_lccs_classes_{year}_{region}_{model_timestamp}_default.png")
        predicted_file_path_custom = os.path.join(output_folder_path, f"predicted_lccs_classes_{year}_{region}_{model_timestamp}_custom.png")
        ground_truth_file_path_custom = os.path.join(output_folder_path, f"ground_truth_lccs_classes_{year}_{region}_{model_timestamp}_custom.png")
        
        # Plot the predictions and ground truth for the current year with the default colormap
        plot_data(predictions, f"Predicted LCCS Classes {year} - {region} (Default)", predicted_file_path_default)
        plot_data(lccs_class, f"Ground Truth LCCS Classes {year} - {region} (Default)", ground_truth_file_path_default)
        
        # Plot the predictions and ground truth for the current year with the custom colormap
        plot_data_with_custom_colormap(predictions, f"Predicted LCCS Classes {year} - {region} (Custom)", predicted_file_path_custom)
        plot_data_with_custom_colormap(lccs_class, f"Ground Truth LCCS Classes {year} - {region} (Custom)", ground_truth_file_path_custom)
        
        # Log the file paths
        logging.info(f"Predicted LCCS classes saved (default) to: {predicted_file_path_default}")
        logging.info(f"Ground truth LCCS classes saved (default) to: {ground_truth_file_path_default}")
        logging.info(f"Predicted LCCS classes saved (custom) to: {predicted_file_path_custom}")
        logging.info(f"Ground truth LCCS classes saved (custom) to: {ground_truth_file_path_custom}")

        # # Save the predicted data in Zarr format
        # logging.info("Saving predictions in Zarr format.")

        # # Convert predictions to Dask array for efficient storage
        # predictions_dask = da.from_array(predictions, chunks=(lat_dim, lon_dim))

        # # Path to save the predicted Zarr data
        # predicted_zarr_path = os.path.join(output_folder_path, f"predicted_lccs_classes_{year}_{region}_{model_timestamp}.zarr")

        # # Save the predictions in Zarr format
        # predictions_dask.to_zarr(predicted_zarr_path)

        # logging.info(f"Predicted LCCS classes saved in Zarr format to: {predicted_zarr_path}")
    else:
        logging.error(f"Size mismatch for year {year}: Cannot reshape to the desired dimensions.")
        raise ValueError(f"Cannot reshape arrays of size {len(predictions.flatten())} and {len(lccs_class.flatten())} into shape ({lat_dim}, {lon_dim}) for year {year}")

logging.info(f"All outputs have been saved to: {output_folder_path}")
logging.info("Prediction and visualization completed successfully.")
