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
import sys 

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

# 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("--use_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_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
use_db = args.use_db
db_path = args.db_path

result = read_db(use_db, db_path, region, train_start_year, train_end_year)

if result:
    model_timestamp = result["model_timestamp"]
    accuracy = result.get("accuracy")
    f1_score = result.get("f1_score")
    precision = result.get("precision")
    recall = result.get("recall")
    class_mapping = result.get("class_mapping")

# 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'XGBoost_{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)

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

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

    # Load features data for the current year
    features_path = os.path.join(save_data_path, f'features_{year}_{region}_inference.zarr')
    features = da.from_zarr(features_path).compute()

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

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

    # Reshape predictions to match spatial dimensions
    predictions = predictions.reshape((lat_dim, lon_dim))

    # Vectorized mapping of predictions to original class values
    predictions = np.vectorize(index_to_class.get)(predictions)

    # Extract the land-sea mask (assuming the mask is in column 14)
    land_sea_mask = features[:, 14].reshape((lat_dim, lon_dim))

    # Apply the water-land mask to the predictions
    predictions_adjusted = apply_land_sea_mask_inference(predictions, land_sea_mask, class_to_index)

    # 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 with the default colormap
    plot_data(predictions_adjusted, f"Predicted LCCS Classes {year} - {region} (Default)", predicted_file_path_default)
    
    # Plot the predictions with the custom colormap
    plot_data_with_custom_colormap(predictions_adjusted, f"Predicted LCCS Classes {year} - {region} (Custom)", predicted_file_path_custom)
    
    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}")
    
    # 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}")

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

