import os
import glob
import re
import argparse
import pickle
from typing import List, Tuple, Optional, Dict

import numpy as np
from sklearn.preprocessing import StandardScaler  # Assuming scalers are sklearn types
from tensorflow.keras.models import load_model  # Assuming models are Keras models
from tqdm import tqdm
import logging

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

def load_pickle(file_path: str):
    """Load a pickle file."""
    try:
        with open(file_path, 'rb') as f:
            return pickle.load(f)
    except Exception as e:
        logging.error(f"Failed to load pickle file {file_path}: {e}")
        return None

def save_pickle(obj, file_path: str):
    """Save an object to a pickle file."""
    try:
        with open(file_path, 'wb') as f:
            pickle.dump(obj, f)
        logging.info(f"Saved pickle file to {file_path}")
    except Exception as e:
        logging.error(f"Failed to save pickle file to {file_path}: {e}")

def main(data_input_dir: str, scaler_dir: str, model_dir: str, data_output_dir: str):
    """
    Main function to process GRIB2 input features, perform predictions, post-process, and save results.

    Args:
        data_input_dir (str): Directory containing input feature pickle files.
        scaler_dir (str): Directory containing scaler pickle files.
        model_dir (str): Directory containing Keras model files.
        data_output_dir (str): Directory to save the output pickle files.
    """
    # Ensure output directory exists
    os.makedirs(data_output_dir, exist_ok=True)

    # Find all files matching *_input_features_*_data.pkl in data_input_dir
    pattern = os.path.join(data_input_dir, '*_input_features_*_data.pkl')
    files = glob.glob(pattern)

    # Group files by country and input format
    countries = {}
    for file in files:
        basename = os.path.basename(file)
        match = re.match(r'(\w+)_input_features_(aifs|ifs)_data\.pkl', basename)
        if match:
            country_name = match.group(1).lower()
            input_format = match.group(2)
            if country_name not in countries:
                countries[country_name] = {}
            countries[country_name][input_format] = file
        else:
            logging.warning(f"File {basename} does not match the expected pattern and will be skipped.")

    if not countries:
        logging.error("No matching input feature files found. Exiting.")
        return

    logging.info(f"Found input feature files for countries: {list(countries.keys())}")

    # Process each country
    for country, formats in countries.items():
        logging.info(f"\nProcessing country: {country}")

        # Check for both aifs and ifs files
        if 'aifs' not in formats or 'ifs' not in formats:
            logging.warning(f"Missing 'aifs' or 'ifs' input features for country {country}. Skipping.")
            continue

        aifs_file = formats['aifs']
        ifs_file = formats['ifs']

        # Load input features
        logging.info(f"Loading input features for {country}: aifs and ifs.")
        aifs_input_features = load_pickle(aifs_file)
        ifs_input_features = load_pickle(ifs_file)

        if aifs_input_features is None or ifs_input_features is None:
            logging.error(f"Failed to load input features for country {country}. Skipping.")
            continue

        # Load scalers
        scaler_X_path = os.path.join(scaler_dir, f"scaler_X_{country}.pkl")
        scaler_y_path = os.path.join(scaler_dir, f"scaler_y_{country}.pkl")

        print(scaler_X_path)

        if not os.path.exists(scaler_X_path) or not os.path.exists(scaler_y_path):
            logging.warning(f"Missing scaler files for country {country}. Skipping.")
            continue

        logging.info(f"Loading scalers for {country}.")
        scaler_X = load_pickle(scaler_X_path)
        scaler_y = load_pickle(scaler_y_path)

        if scaler_X is None or scaler_y is None:
            logging.error(f"Failed to load scalers for country {country}. Skipping.")
            continue

        # Normalize inputs
        logging.info(f"Normalizing inputs for {country}.")
        try:
            aifs_normalized = scaler_X.transform(aifs_input_features)
            ifs_normalized = scaler_X.transform(ifs_input_features[:, :-2])
        except Exception as e:
            logging.error(f"Normalization failed for country {country}: {e}")
            continue

        # Load model
        model_path = os.path.join(model_dir, f"{country}_model.keras")
        if not os.path.exists(model_path):
            logging.warning(f"Model file {model_path} does not exist for country {country}. Skipping.")
            continue

        logging.info(f"Loading model for {country}.")
        try:
            model = load_model(model_path)
        except Exception as e:
            logging.error(f"Failed to load model for country {country}: {e}")
            continue

        # Predict
        logging.info(f"Running predictions for {country}.")
        try:
            inference_ifs = model.predict(ifs_normalized)
            inference_aifs = model.predict(aifs_normalized)
        except Exception as e:
            logging.error(f"Prediction failed for country {country}: {e}")
            continue

        # Denormalize outputs
        logging.info(f"Denormalizing predictions for {country}.")
        try:
            inference_aifs_denorm = scaler_y.inverse_transform(inference_aifs)
            inference_ifs_denorm = scaler_y.inverse_transform(inference_ifs)
        except Exception as e:
            logging.error(f"Denormalization failed for country {country}: {e}")
            continue

        # Calculate average sum values for AIFS
        logging.info(f"Calculating average sum values for AIFS in {country}.")
        num_aifs_files = len(msl_data_aifs)
        if num_aifs_files == 0:
            logging.warning(f"No AIFS 'msl' data entries found for country {country}. Skipping AIFS post-processing.")
            average_entire_aifs = None
        else:
            total_aifs_samples = inference_aifs_denorm.shape[0]
            block_size_aifs = total_aifs_samples // num_aifs_files
            sum_values_aifs = []
            for i in range(0, total_aifs_samples, block_size_aifs):
                block = inference_aifs_denorm[i:i + block_size_aifs]
                if block.size == 0:
                    continue
                sum_block = 0
                for val in block:
                    # Assuming val has at least two elements
                    sum_block += np.sqrt(val[0]**2 + val[1]**2)
                avg_sum = sum_block / len(block)
                sum_values_aifs.append(avg_sum)
            average_entire_aifs = np.array(sum_values_aifs)
            logging.info(f"Average Sum Values for AIFS in {country}: {average_entire_aifs}")

        # Calculate average sum values for IFS
        logging.info(f"Calculating average sum values for IFS in {country}.")
        num_ifs_files = len(msl_data_ifs)
        if num_ifs_files == 0:
            logging.warning(f"No IFS 'msl' data entries found for country {country}. Skipping IFS post-processing.")
            average_entire_ifs = None
        else:
            total_ifs_samples = inference_ifs_denorm.shape[0]
            block_size_ifs = total_ifs_samples // num_ifs_files
            sum_values_ifs = []
            for i in range(0, total_ifs_samples, block_size_ifs):
                block = inference_ifs_denorm[i:i + block_size_ifs]
                if block.size == 0:
                    continue
                sum_block = 0
                for val in block:
                    # Assuming val has at least two elements
                    sum_block += np.sqrt(val[0]**2 + val[1]**2)
                avg_sum = sum_block / len(block)
                sum_values_ifs.append(avg_sum)
            average_entire_ifs = np.array(sum_values_ifs)
            logging.info(f"Average Sum Values for IFS in {country}: {average_entire_ifs}")

        # Save the averages as pickle files
        if average_entire_aifs is not None:
            aifs_output_path = os.path.join(data_output_dir, f"average_entire_{country}_from_aifs_100m.pkl")
            save_pickle(average_entire_aifs, aifs_output_path)
        else:
            logging.info(f"No average_entire_aifs_100m to save for {country}.")

        if average_entire_ifs is not None:
            ifs_output_path = os.path.join(data_output_dir, f"average_entire_{country}_from_ifs_100m.pkl")
            save_pickle(average_entire_ifs, ifs_output_path)
        else:
            logging.info(f"No average_entire_ifs_100m to save for {country}.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Perform inference on normalized input features using pre-trained models.")
    parser.add_argument(
        "--data_input_dir",
        type=str,
        required=True,
        help="Path to the data input directory containing input feature pickle files."
    )
    parser.add_argument(
        "--scaler_dir",
        type=str,
        required=True,
        help="Path to the scaler directory containing scaler_X_{country}.pkl and scaler_y_{country}.pkl."
    )
    parser.add_argument(
        "--model_dir",
        type=str,
        required=True,
        help="Path to the model directory containing {country}_model.keras files."
    )
    parser.add_argument(
        "--data_output_dir",
        type=str,
        required=True,
        help="Path to the data output directory where results will be saved."
    )

    args = parser.parse_args()
    main(args.data_input_dir, args.scaler_dir, args.model_dir, args.data_output_dir)2024-10-10 14:21:54,612 - INFO - Found input feature files for countries: ['france', 'spain', 'germany']
2024-10-10 14:21:54,612 - INFO -
Processing country: france
2024-10-10 14:21:54,612 - INFO - Loading input features for france: aifs and ifs.
scalers/scaler_X_france.pkl
2024-10-10 14:21:54,644 - INFO - Loading scalers for france.
2024-10-10 14:21:54,646 - INFO - Normalizing inputs for france.
2024-10-10 14:21:54,671 - ERROR - Normalization failed for country france: list indices must be integers or slices, not tuple
2024-10-10 14:21:54,671 - INFO -
Processing country: spain
2024-10-10 14:21:54,671 - INFO - Loading input features for spain: aifs and ifs.
scalers/scaler_X_spain.pkl
2024-10-10 14:21:54,697 - INFO - Loading scalers for spain.
2024-10-10 14:21:54,698 - INFO - Normalizing inputs for spain.
2024-10-10 14:21:54,715 - ERROR - Normalization failed for country spain: list indices must be integers or slices, not tuple
2024-10-10 14:21:54,715 - INFO -
Processing country: germany
2024-10-10 14:21:54,715 - INFO - Loading input features for germany: aifs and ifs.
scalers/scaler_X_germany.pkl
2024-10-10 14:21:54,737 - INFO - Loading scalers for germany.
2024-10-10 14:21:54,737 - INFO - Normalizing inputs for germany.
2024-10-10 14:21:54,753 - ERROR - Normalization failed for country germany: list indices must be integers or slices, not tuple please fix ERROR