# **************************************************************************** #
#                                                                              #
#                                                         :::      ::::::::    #
#    tools.py                                           :+:      :+:    :+:    #
#                                                     +:+ +:+         +:+      #
#    By: daniloceano <danilo.oceano@gmail.com>      +#+  +:+       +#+         #
#                                                 +#+#+#+#+#+   +#+            #
#    Created: 2023/12/19 17:33:03 by daniloceano       #+#    #+#              #
#    Updated: 2025/02/26 11:00:00 by daniloceano      ###   ########.fr        #
#                                                                              #
# **************************************************************************** #

"""
Utility functions for the Lorenz Energy Cycle Toolkit.

This module contains general utility functions including logging initialization,
coordinate conversion, CDS API data retrieval, and helper functions for
finding extremum coordinates.
"""

import argparse
import logging
import math
import os
import tempfile

import numpy as np
import pandas as pd
import xarray as xr


def initialize_logging(results_subdirectory, args):
    """
    Initializes the logging configuration for the application.

    Args:
        results_subdirectory (str): Directory path to save the log file.
        args (object): The argparse object containing the command line arguments.
    """

    verbose = True if args.verbosity else False

    # Set root logger to higher severity level (INFO or ERROR)
    root_log_level = logging.ERROR if not verbose else logging.INFO
    logging.basicConfig(
        level=root_log_level, format="%(asctime)s - %(levelname)s - %(message)s"
    )

    # Create a separate logger for the application
    app_logger = logging.getLogger("lorenzcycletoolkit")
    app_log_level = logging.DEBUG if verbose else logging.INFO
    app_logger.setLevel(app_log_level)
    # Prevent the logger from propagating messages to the root logger
    app_logger.propagate = False

    # Create file handler for saving logs
    log_file_name = f'log.{os.path.basename(args.infile).split(".")[0]}'
    log_file = os.path.join(results_subdirectory, log_file_name)
    file_handler = logging.FileHandler(log_file, mode="w")
    file_handler.setLevel(app_log_level)
    file_formatter = logging.Formatter(
        "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
    )
    file_handler.setFormatter(file_formatter)
    app_logger.addHandler(file_handler)

    # Create a console handler for app logger
    console_handler = logging.StreamHandler()
    console_handler.setLevel(app_log_level)
    console_handler.setFormatter(file_formatter)
    app_logger.addHandler(console_handler)

    return app_logger


def convert_longitude_range(df: xr.Dataset, lon_indexer: str) -> xr.Dataset:
    """
    Convert longitude range from 0-360 degrees to -180 to 180 degrees.

    This function modifies the longitude coordinates of the provided xarray Dataset
    and sorts the data based on these updated longitudes.

    Args:
        df (xr.Dataset): The dataset containing longitude coordinates.
        lon_indexer (str): The name of the longitude coordinate in the dataset.

    Returns:
        xr.Dataset: The dataset with updated longitude coordinates.
    """
    df.coords[lon_indexer] = (df.coords[lon_indexer] + 180) % 360 - 180
    df = df.sortby(df[lon_indexer])
    return df


def find_extremum_coordinates(
    data: xr.DataArray, lat: xr.DataArray, lon: xr.DataArray, variable: str
) -> tuple:
    """
    Finds the indices of extremum values for a given variable in a dataset.

    Args:
        data (xr.DataArray): DataArray containing the data for calculation.
        lat (xr.DataArray): DataArray containing the latitudes.
        lon (xr.DataArray): DataArray containing the longitudes.
        variable (str): Name of the variable to find extremum indices.

    Returns:
        tuple: Tuple containing latitude and longitude of the extremum value.
    """
    lat_values = lat.values
    lon_values = lon.values
    min_lat = lat_values.min()

    data = np.array(data)  # Ensure data is a numpy array

    if variable == "min_max_zeta":
        index = np.unravel_index(
            data.argmin() if min_lat < 0 else data.argmax(), data.shape
        )
    elif variable in ["min_hgt", "max_wind"]:
        index = np.unravel_index(
            data.argmin() if variable == "min_hgt" else data.argmax(), data.shape
        )
    else:
        logging.error(f"Invalid variable specified: {variable}")
        raise ValueError(f"Invalid variable specified: {variable}")

    return lat_values[index[0]], lon_values[index[1]]


def get_cdsapi_data(
    args: argparse.Namespace, track: pd.DataFrame, app_logger: logging.Logger
) -> str:
    """
    Retrieve ERA5 pressure level data from the CDS API based on track information.

    This function downloads ERA5 reanalysis data day-by-day to avoid issues with
    large file downloads, then concatenates the daily files into a single output file.
    Temporary daily files are automatically deleted after concatenation.

    Args:
        args (argparse.Namespace): Command-line arguments containing the output file path.
        track (pd.DataFrame): DataFrame with 'Lat', 'Lon' columns and datetime index.
        app_logger (logging.Logger): Logger for the application.

    Returns:
        str: Path to the downloaded NetCDF file.

    Raises:
        FileNotFoundError: If temporary files are not created during download.
        Exception: For other errors during data retrieval or concatenation.
    """
    import math
    import tempfile

    import cdsapi

    # Extract bounding box (lat/lon limits) from track
    min_lat, max_lat = float(track["Lat"].min()), float(track["Lat"].max())
    min_lon, max_lon = float(track["Lon"].min()), float(track["Lon"].max())

    # Apply a 15-degree buffer and round to nearest integer
    buffered_max_lat = math.ceil(max_lat + 15)
    buffered_min_lon = math.floor(min_lon - 15)
    buffered_min_lat = math.floor(min_lat - 15)
    buffered_max_lon = math.ceil(max_lon + 15)

    # Define the area for the request [North, West, South, East]
    area = [buffered_max_lat, buffered_min_lon, buffered_min_lat, buffered_max_lon]

    # Pressure levels in hPa
    pressure_levels = [
        "1", "2", "3", "5", "7", "10", "20", "30", "50", "70",
        "100", "125", "150", "175", "200", "225", "250", "300",
        "350", "400", "450", "500", "550", "600", "650", "700",
        "750", "775", "800", "825", "850", "875", "900", "925",
        "950", "975", "1000",
    ]

    # Variables to retrieve
    variables = [
        "u_component_of_wind",
        "v_component_of_wind",
        "temperature",
        "vertical_velocity",
        "geopotential",
    ]

    # Convert track index to DatetimeIndex and find the date range
    track_datetime_index = pd.DatetimeIndex(track.index)
    first_timestamp = track_datetime_index.min()
    last_timestamp = track_datetime_index.max()

    # Get all unique dates needed for download
    dates = pd.date_range(
        start=first_timestamp.date(),
        end=last_timestamp.date(),
        freq='D'
    )

    # Use time_resolution from args (default is 3 hours)
    time_step = args.time_resolution
    app_logger.debug(f"⏱️ Using time resolution from args: {time_step} hour(s)")

    # Log track file bounds and requested data bounds
    app_logger.debug(
        f"📍 Track File Limits: lon range: [{min_lon:.2f}, {max_lon:.2f}], "
        f"lat range: [{min_lat:.2f}, {max_lat:.2f}]"
    )
    app_logger.debug(
        f"🗺️ Buffered Data Bounds: lon range: [{buffered_min_lon}, {buffered_max_lon}], "
        f"lat range: [{buffered_min_lat}, {buffered_max_lat}]"
    )
    app_logger.debug(
        f"📅 Track period: {first_timestamp} to {last_timestamp}"
    )
    app_logger.info(
        f"📥 Will download {len(dates)} day(s) of data: {dates[0].date()} to {dates[-1].date()}"
    )

    # Create temporary directory for daily files
    temp_dir = tempfile.mkdtemp(prefix="cdsapi_daily_")
    app_logger.debug(f"📁 Created temporary directory: {temp_dir}")
    
    # Download data day by day
    app_logger.info(f"🌐 Starting ERA5 data download from CDS API...")
    app_logger.info(f"📦 Expected files to download: {len(dates)}")
    for i, date in enumerate(dates, 1):
        app_logger.info(f"  📅 Day {i}: {date.date()}")
    
    daily_files = []
    client = cdsapi.Client(timeout=600, retry_max=500)
    
    try:
        for i, date in enumerate(dates, 1):
            year = date.strftime('%Y')
            month = date.strftime('%m')
            day = date.strftime('%d')
            
            # Determine which hours to download for this day
            # For the first day, start from the first timestamp hour
            # For the last day, end at the last timestamp hour
            # For middle days, download all hours
            if date.date() == first_timestamp.date():
                # First day: start from first timestamp hour
                start_hour = first_timestamp.hour
                # Round down to nearest time_step
                start_hour = (start_hour // time_step) * time_step
                if date.date() == last_timestamp.date():
                    # Same day: both start and end on same day
                    end_hour = last_timestamp.hour
                    # Round up to include the last timestamp - go one step beyond
                    end_hour = ((end_hour // time_step) + 1) * time_step
                    if end_hour <= start_hour:
                        end_hour = start_hour + time_step
                    # Cap at 24 hours (end of day)
                    if end_hour > 24:
                        end_hour = 24
                else:
                    # First day but not last: download until end of day
                    end_hour = 24
            elif date.date() == last_timestamp.date():
                # Last day (not first): start from beginning of day
                start_hour = 0
                end_hour = last_timestamp.hour
                # Round up to nearest time_step to INCLUDE the last timestamp
                # We need to go one step further to include the last hour in the range
                end_hour = ((end_hour // time_step) + 1) * time_step
                # If last timestamp is at midnight (hour 0), download that hour
                if end_hour == 0:
                    end_hour = time_step
                # Cap at 24 hours (end of day)
                if end_hour > 24:
                    end_hour = 24
            else:
                # Middle day: download all hours
                start_hour = 0
                end_hour = 24 
            
            # Generate times list for this specific day
            times = [f"{hour:02d}:00" for hour in range(start_hour, end_hour, time_step)]
            
            if not times:
                app_logger.warning(f"⚠️ No times to download for {year}-{month}-{day}, skipping...")
                continue
            
            # Create temporary file for this day
            temp_file = os.path.join(temp_dir, f"era5_{year}{month}{day}.nc")
            daily_files.append(temp_file)
            
            app_logger.info(f"⬇️ Downloading day {i}/{len(dates)}: {year}-{month}-{day}")
            app_logger.debug(f"  ⏰ Times: {times[0]} to {times[-1]} ({len(times)} timesteps)")
            
            # Prepare request for single day
            request = {
                "product_type": "reanalysis",
                "format": "netcdf",
                "variable": variables,
                "pressure_level": pressure_levels,
                "year": year,
                "month": month,
                "day": day,
                "time": times,
                "area": area,
            }
            
            # Download data for this day
            try:
                client.retrieve(
                    "reanalysis-era5-pressure-levels",
                    request,
                    temp_file,
                )
                
                # Verify file was created
                if not os.path.exists(temp_file):
                    app_logger.error("❌ CDS API file creation failed!")
                    app_logger.error("\n" + "="*70)
                    app_logger.error("☁️  CDS API DOWNLOAD ERROR")
                    app_logger.error("="*70)
                    app_logger.error(f"Expected file: {temp_file}")
                    app_logger.error(f"Date: {year}-{month}-{day}")
                    app_logger.error("\n💡 Possible causes:")
                    app_logger.error("   1. CDS API request succeeded but file wasn't written")
                    app_logger.error("   2. Disk space full")
                    app_logger.error("   3. Permission issues in temp directory")
                    app_logger.error("\n🔧 Developer Info:")
                    app_logger.error(f"   Request: {request}")
                    app_logger.error(f"   Temp dir: {temp_dir}")
                    app_logger.error("="*70 + "\n")
                    raise FileNotFoundError(
                        f"CDS API file not created at expected path: {temp_file}"
                    )
                
                file_size = os.path.getsize(temp_file)
                app_logger.info(f"  ✓ Downloaded successfully: {file_size / (1024**2):.2f} MB")
                
            except FileNotFoundError:
                raise  # Re-raise the FileNotFoundError we just created
            except KeyError as e:
                app_logger.error(f"❌ CDS API authentication or configuration error!")
                app_logger.error("\n" + "="*70)
                app_logger.error("🔑 CDS API AUTHENTICATION ERROR")
                app_logger.error("="*70)
                app_logger.error(f"Date: {year}-{month}-{day}")
                app_logger.error(f"Error: {e}")
                app_logger.error("\n💡 User Solutions:")
                app_logger.error("   1. Check if ~/.cdsapirc file exists and is configured:")
                app_logger.error("      url: https://cds.climate.copernicus.eu/api/v2")
                app_logger.error("      key: YOUR_UID:YOUR_API_KEY")
                app_logger.error("   2. Register at: https://cds.climate.copernicus.eu/")
                app_logger.error("   3. Get your API key from your profile page")
                app_logger.error("\n🔧 Developer Info:")
                app_logger.error(f"   Exception: {type(e).__name__}: {str(e)}")
                app_logger.error("="*70 + "\n")
                raise
            except Exception as e:
                app_logger.error(f"❌ Error downloading CDS API data for {year}-{month}-{day}!")
                app_logger.error("\n" + "="*70)
                app_logger.error("☁️  CDS API REQUEST ERROR")
                app_logger.error("="*70)
                app_logger.error(f"Date: {year}-{month}-{day}")
                app_logger.error(f"Error type: {type(e).__name__}")
                app_logger.error(f"Error message: {str(e)}")
                app_logger.error("\n💡 Common causes:")
                app_logger.error("   1. Network connectivity issues")
                app_logger.error("   2. CDS API service is down or overloaded")
                app_logger.error("   3. Invalid request parameters")
                app_logger.error("   4. API quota exceeded")
                app_logger.error("   5. Timeout (large requests)")
                app_logger.error("\n🔍 Check CDS API status:")
                app_logger.error("   https://cds.climate.copernicus.eu/live/queue")
                app_logger.error("\n🔧 Developer Info:")
                app_logger.error(f"   Request parameters: {request}")
                app_logger.error(f"   Exception: {type(e).__name__}")
                app_logger.error(f"   Message: {str(e)}")
                app_logger.error("="*70 + "\n")
                raise
        
        app_logger.info("=" * 60)
        app_logger.info("✅ All daily files downloaded successfully!")
        app_logger.info(f"📦 Total files: {len(daily_files)}")
        app_logger.info("=" * 60)
        
        # Concatenate daily files
        app_logger.info("🔗 Concatenating daily files into single dataset...")
        datasets = []
        for idx, daily_file in enumerate(daily_files, 1):
            app_logger.debug(f"  📂 Loading file {idx}/{len(daily_files)}: {os.path.basename(daily_file)}")
            ds = xr.open_dataset(daily_file)
            datasets.append(ds)
        
        # Concatenate along time dimension
        app_logger.info("🔀 Merging datasets along time dimension...")
        combined_ds = xr.concat(datasets, dim='valid_time')
        
        # Save concatenated dataset
        app_logger.info(f"💾 Saving final dataset to: {args.infile}")
        combined_ds.to_netcdf(args.infile)
        combined_ds.close()
        
        # Close all datasets
        for ds in datasets:
            ds.close()
        
        # Verify the final file was created
        if not os.path.exists(args.infile):
            app_logger.error("❌ Final CDS API file was not created!")
            app_logger.error("\n" + "="*70)
            app_logger.error("📦 FILE CONCATENATION ERROR")
            app_logger.error("="*70)
            app_logger.error(f"Expected file: {args.infile}")
            app_logger.error(f"Daily files downloaded: {len(daily_files)}")
            app_logger.error("\n💡 Possible causes:")
            app_logger.error("   1. Error during dataset concatenation")
            app_logger.error("   2. Disk space full")
            app_logger.error("   3. Permission issues")
            app_logger.error("\n🔧 Developer Info:")
            app_logger.error(f"   Target file: {args.infile}")
            app_logger.error(f"   Daily files: {len(daily_files)}")
            app_logger.error(f"   Temp directory: {temp_dir}")
            app_logger.error("="*70 + "\n")
            raise FileNotFoundError(
                f"CDS API file not created at expected path: {args.infile}. "
                f"Check disk space and permissions."
            )
        
        final_file_size = os.path.getsize(args.infile)
        app_logger.info("=" * 60)
        app_logger.info("🎉 DATA DOWNLOAD COMPLETED SUCCESSFULLY!")
        app_logger.info(f"📄 Final file: {args.infile}")
        app_logger.info(f"📊 File size: {final_file_size / (1024**2):.2f} MB")
        app_logger.info(f"📅 Time coverage: {first_timestamp} to {last_timestamp}")
        app_logger.info("=" * 60)
        
    finally:
        # Clean up temporary files and directory
        app_logger.debug("🧹 Cleaning up temporary files...")
        for temp_file in daily_files:
            try:
                if os.path.exists(temp_file):
                    os.remove(temp_file)
                    app_logger.debug(f"🗑️ Deleted temporary file: {temp_file}")
            except Exception as e:
                app_logger.warning(f"⚠️ Could not delete temporary file {temp_file}: {e}")
        
        try:
            os.rmdir(temp_dir)
            app_logger.debug(f"🗑️ Deleted temporary directory: {temp_dir}")
        except Exception as e:
            app_logger.warning(f"⚠️ Could not delete temporary directory {temp_dir}: {e}")
    
    return args.infile
