import argparse
import dask.array as da
import json
import joblib
import numpy as np
import os
import sys
from datetime import datetime
import logging

# Start modulirzing the code 
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../source')))
from database import get_latest_model
from plotting import plot_data, plot_data_with_custom_colormap
from data_processing import apply_land_sea_mask


# 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=False, 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=False, help="Model timestamp from training phase")
parser.add_argument("--read_db", action="store_true", help="Read model details from the database")
parser.add_argument("--db_path", type=str, required=False, help="Path to the SQLite database")

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 or Database
if args.read_db:
    logging.info("Fetching model details from the database.")
    model_info = get_latest_model(args.db_path, args.region, args.start_year, args.end_year)
    if model_info:
        start_year = model_info["start_year"]
        end_year = model_info["end_year"]
        region = model_info["region"]
        model_timestamp = model_info["model_timestamp"]
        test_years = model_info["test_years"]
    else:
        raise ValueError("No model found in the database for the given region and years.")
else:
    start_year = args.start_year
    end_year = args.end_year
    region = args.region
    test_years = args.test_years
    model_timestamp = args.model_timestamp

# Load metadata to get dimensions
metadata_path = os.path.join(args.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(args.save_data_path, f'XGBoost_{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(args.save_data_path, output_folder_name)

os.makedirs(output_folder_path, exist_ok=True)
logging.info(f"Saving all outputs to folder: {output_folder_path}")

# Load Class-to-Index Mapping
class_to_index_path = os.path.join(args.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)

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()}

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

    lccs_class_path = os.path.join(args.save_data_path, f'lccs_class_{year}_{region}.zarr')
    features_path = os.path.join(args.save_data_path, f'features_{year}_{region}.zarr')

    logging.info("Loading processed data.")
    lccs_class = da.from_zarr(lccs_class_path)
    features = da.from_zarr(features_path)

    logging.info(f"Shape of lccs_class (Dask): {lccs_class.shape}")
    logging.info(f"Shape of features (Dask): {features.shape}")

    lccs_class = lccs_class.compute()
    features = features.compute()

    logging.info(f"Shape of lccs_class (NumPy): {lccs_class.shape}")
    logging.info(f"Shape of features (NumPy): {features.shape}")

    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}.")
        raise ValueError(f"lccs_class for year {year} does not have the expected shape (lat_dim * lon_dim).")

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

    # Apply the water-land mask to the predictions (if applicable)
    predictions_adjusted = apply_land_sea_mask(predictions, features, class_to_index)

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

    # Save the predicted data in Zarr format
    logging.info(f"Saving predictions for the year {year} in Zarr format.")

    # Convert predictions to Dask array for efficient storage
    predictions_dask = da.from_array(predictions_adjusted, 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, overwrite=True)

    logging.info(f"Predicted LCCS classes for the year {year} saved in Zarr format to: {predicted_zarr_path}")

    expected_size_per_year = lat_dim * lon_dim

    if len(predictions_adjusted.flatten()) == expected_size_per_year and len(lccs_class.flatten()) == expected_size_per_year:
        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_data(predictions_adjusted, 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_data_with_custom_colormap(predictions_adjusted, 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)

        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}")
    else:
        logging.error(f"Size mismatch for year {year}: Cannot reshape to the desired dimensions.")
        raise ValueError(f"Cannot reshape arrays for year {year}")

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