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 for Inference")

parser.add_argument("--train_start_year", type=int, required=True, help="Start year of the training dataset")
parser.add_argument("--train_end_year", type=int, required=True, help="End year of the training dataset")
parser.add_argument("--inference_years", type=int, nargs='+', required=True, help="Years for which to run inference (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 the 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_inference_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
train_start_year = args.train_start_year
train_end_year = args.train_end_year
inference_years = args.inference_years
region = args.region
save_data_path = args.save_data_path
model_timestamp = args.model_timestamp

# Load training metadata to get dimensions
metadata_path = os.path.join(save_data_path, f'metadata_{train_start_year}_{train_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_{train_start_year}_{train_end_year}_{region}_{model_timestamp}.joblib')

# Create output folder based on model details
output_folder_name = f"{region}_{train_start_year}_{train_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_{train_start_year}_{train_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 inference year
for year in inference_years:
    logging.info(f"Processing data for the year: {year}")

    # Set paths for loading feature data for each year (no ground truth for inference)
    features_path = os.path.join(save_data_path, f'features_{year}_{region}_inference.zarr')

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

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

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

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

    # 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))

    # File paths for predictions
    predicted_file_path_default = os.path.join(output_folder_path, f"predicted_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")
    
    # Plot the predictions for the current year with the default colormap
    plot_data(predictions, f"Predicted LCCS Classes {year} - {region} (Default)", predicted_file_path_default)
    
    # Plot the predictions 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)
    
    # Log the file paths
    logging.info(f"Predicted LCCS classes saved (default) to: {predicted_file_path_default}")
    logging.info(f"Predicted LCCS classes saved (custom) to: {predicted_file_path_custom}")

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