import os
import netCDF4 as nc4
from logger import logger
from datetime import datetime, timedelta
import numpy as np
from math import ceil


def forecast_download(
    start_time, flight_duration, latitude, longitude, range_lat, range_lon
):
    # Start_time format is "dd.MM.yyyy HH:mm"
    # Flight_time is in hours
    # Check if the forecast folder exists

    if not os.path.exists("forecasts"):
        os.makedirs("forecasts")
        logger.info("Created forecasts directory.")

    download_date = datetime.now()
    # Normalize to the nearest previous hour divisible by 6
    hour_divisible_by_6 = download_date.hour - (download_date.hour % 6)

    # Create a new datetime object at the nearest previous hour divisible by 6
    nearest_hour = download_date.replace(
        hour=hour_divisible_by_6, minute=0, second=0, microsecond=0
    )

    # Ensure it is at least 6 hours ago
    if nearest_hour > download_date - timedelta(hours=6):
        # If not 6 hours ago, keep subtracting 6 hours until it is
        while nearest_hour > download_date - timedelta(hours=6):
            nearest_hour -= timedelta(hours=6)

    # Format the datetime object into the specified string format
    formatted_date = nearest_hour.strftime("%Y%m%d_%H")

    # Range latitude and longitude must be dividable by 0.25
    range_lon = round(range_lon * 4) / 4
    range_lat = round(range_lat * 4) / 4

    # Calculate the indices for the forecast data extraction
    start_time_index, end_time_index = calculate_indices(
        start_time, flight_duration, nearest_hour
    )
    location_string = (
        f"_{int(latitude)}_{int(longitude)}_{int(range_lat)}_{int(range_lon)}"
    )
    # Check if the forecast file exists
    forecast_filename = "forecasts/gfs_0p25_" + formatted_date + location_string + ".nc"
    print(f"Checking if forecast file exists: {forecast_filename}")
    logger.info(f"Checking if forecast file exists: {forecast_filename}")

    manager = NetCDFManager()
    try:

        if not os.path.exists(forecast_filename):
            # Find the nearest available forecast data
            forecast_filename, url = find_forecast(nearest_hour, location_string)

            if url is None:
                # Data already downloaded
                return forecast_filename
            else:
                print(
                    f"Downloading GFS forecast data for {forecast_filename} from {url}"
                )
                logger.info(
                    f"Downloading GFS forecast data for {forecast_filename} from {url}"
                )

                forecast_filename = download_GFS_forecast_data(
                    latitude,
                    longitude,
                    range_lat,
                    range_lon,
                    forecast_filename,
                    url,
                    start_time_index,
                    end_time_index,
                    manager,
                )
                print("Download complete. forecast_filename: ", forecast_filename)
                logger.info("Download complete.")
                return forecast_filename
        else:
            print("GFS forecast data already downloaded.")
            logger.info("GFS forecast data already downloaded.")
            return forecast_filename
    except Exception as e:
        print(f"Error downloading or processing GFS forecast data: {e}")
        logger.error(f"Error downloading or processing GFS forecast data: {e}")
        if os.path.exists(forecast_filename):
            manager.close_all()
            os.remove(forecast_filename)
            logger.info(
                f"Removed incomplete forecast file {forecast_filename} due to error."
            )
        raise


def find_forecast(nearest_hour, location_string):
    max_attempts = 3
    attempt = 0

    while attempt < max_attempts:
        date = nearest_hour.strftime("%Y%m%d")
        hour = nearest_hour.strftime("%H")
        url = (
            "https://nomads.ncep.noaa.gov/dods/gfs_0p25/gfs"
            + date
            + "/gfs_0p25_"
            + hour
            + "z"
        )
        try:
            nc4.Dataset(url)
            return forecast_filename, url
        except:
            logger.error(
                f"NOAA DODS Server error with url {url}. Attempt {attempt + 1} failed."
            )
            nearest_hour -= timedelta(hours=6)
            attempt += 1
            # Format the datetime object into the specified string format
            formatted_date = nearest_hour.strftime("%Y%m%d_%H")
            forecast_filename = (
                "forecasts/gfs_0p25_" + formatted_date + location_string + ".nc"
            )
            if os.path.exists(forecast_filename):
                return forecast_filename, None

            if attempt == max_attempts:
                raise Exception(
                    f"NOAA DODS Server error with url {url}. Data not downloaded."
                )


# Function to download GFS forecast data
def download_GFS_forecast_data(
    latitude,
    longitude,
    range_lat,
    range_lon,
    forecast_filename,
    url,
    start_time_index,
    end_time_index,
    manager,
):

    # Calculate min and max indexes
    # lat_min, lat_max, lon_min, lon_max = get_latitude_longitude_indices(
    #     latitude, longitude, range_lat, range_lon
    # )
    lat_min, lat_max, lon_min, lon_max = get_limits(
        longitude, latitude, range_lon, range_lat
    )
    # print(
    #     f"Latitude min: {lat_min},max: {lat_max}, Longitude min: {lon_min},max: {lon_max}"
    # )

    coords = ["time", "lat", "lon", "lev"]
    vars_out = ["vgrdprs", "ugrdprs", "hgtprs", "tmpprs", "dzdtprs"]

    try:
        nc_in = manager.open_file(url)
    except:
        raise Exception(f"NOAA DODS Server error with url {url}. Data not downloaded.")

    nc_out = manager.open_file(forecast_filename, "w")

    for name, dimension in nc_in.dimensions.items():
        nc_out.createDimension(
            name, len(dimension) if not dimension.isunlimited() else None
        )

    for name, variable in nc_in.variables.items():
        if name in coords:
            x = nc_out.createVariable(name, variable.datatype, variable.dimensions)
            v = nc_in.variables[name][:]
            nc_out.variables[name][:] = v
            logger.info(f"Downloaded {name}")
            print("Downloaded " + name)

        if name in vars_out:
            logger.info(f"Downloading {name}")
            print("Downloading " + name)
            x = nc_out.createVariable(
                name, variable.datatype, variable.dimensions, zlib=True
            )  # zlib for compression

            # Download only a chunk of the data
            for i in range(start_time_index, end_time_index + 1):
                for x in range(len(lat_max)):
                    print(f"lat min: {lat_min[x]}, lat_max {lat_max[x]}")
                    print(f"lon min: {lon_min[x]}, lon_max {lon_max[x]}")

                    data = nc_in.variables[name][
                        i,
                        :,
                        lat_min[x] : lat_max[x] + 1,
                        lon_min[x] : lon_max[x] + 1,
                    ]
                    nc_out.variables[name][
                        i,
                        :,
                        lat_min[x] : lat_max[x] + 1,
                        lon_min[x] : lon_max[x] + 1,
                    ] = data
                logger.info(
                    f"Downloaded and added to output file {name}; hour index - {i}; time - {i * 3}"
                )
                print(
                    "Downloaded and added to output file ",
                    name,
                    " hour index - ",
                    i,
                    " time - ",
                    i * 3,
                )
    manager.close_all()
    return forecast_filename


class NetCDFManager:
    def __init__(self):
        self.open_files = []

    def open_file(self, filename, mode="r"):
        nc_file = nc4.Dataset(filename, mode)
        self.open_files.append(nc_file)
        return nc_file

    def close_all(self):
        for nc_file in self.open_files:
            nc_file.close()
        self.open_files = []


# def get_latitude_longitude_indices(latitude, longitude, range_lat, range_lon):
#     """
#     Calculate the indices of latitude and longitude based on the given coordinates and range.

#     Parameters:
#     -----------
#     latitude: float
#         The latitude of the location.
#     longitude: float
#         The longitude of the location.
#     range_lat: float
#         The range of latitude.
#     range_lon: float
#         The range of longitude.

#     Returns:
#     --------
#     lat_min: int
#         The minimum latitude index.
#     lat_max: int
#         The maximum latitude index.
#     lon_min: int
#         The minimum longitude index.
#     lon_max: int
#         The maximum longitude index.
#     """
#     # Indexes for latitude and longitude
#     lat_i = getNearestLat(latitude, -90, 90.01)
#     lon_i = getNearestLon(longitude, 0, 360)
#     # Indexes for range latitude and longitude
#     range_lat_index = int(range_lon / 2 / 0.25)
#     range_lon_index = int(range_lat / 2 / 0.25)
#     # Maximal and minimal indexes for latitude and longitude
#     latitude_limit_index = 721
#     longitude_limit_index = 1440
#     # Calculate min and max indexes
#     lat_min = max(lat_i - range_lat_index, 0)
#     lat_max = min(lat_i + range_lat_index, latitude_limit_index - 1)
#     lon_min = max(lon_i - range_lon_index, 0)
#     lon_max = min(lon_i + range_lon_index, longitude_limit_index - 1)

#     # Create a continuous array of longitude indices
#     lon_indices = (
#         np.arange(lon_i - range_lon_index, lon_i + range_lon_index + 1) % 1440
#     ).astype(int)


#     print(
#         f"Latitude min: {lat_min},max: {lat_max}, Longitude min: {lon_min},max: {lon_max}"
#     )
#     return lat_min, lat_max, lon_min, lon_max


def get_latitude_longitude_indices(
    latitude, longitude, range_lat, range_lon, lat_size=721, lon_size=1440
):
    lat_i = get_nearest_lat(latitude, -90, 90.01)
    lon_i = get_nearest_lon(longitude, 0, 360)

    range_lat_index = int(range_lat / 2 / 0.25)
    range_lon_index = int(range_lon / 2 / 0.25)

    # Handle latitude indices with potential reflection at the poles
    lat_indices = np.arange(lat_i - range_lat_index, lat_i + range_lat_index + 1)
    lat_indices = np.clip(
        lat_indices, 0, lat_size - 1
    )  # Ensure indices stay within bounds

    # Reflecting over the poles
    lat_indices = np.where(lat_indices < 0, -lat_indices, lat_indices)
    lat_indices = np.where(
        lat_indices >= lat_size, 2 * lat_size - lat_indices - 2, lat_indices
    )

    # Handle longitude indices with wrap-around
    lon_indices = (
        np.arange(lon_i - range_lon_index, lon_i + range_lon_index + 1) % lon_size
    )

    return lat_indices, lon_indices


def closest(arr, k):
    """Given an ordered array and a value, determines the index of the closest item
    contained in the array.
    """
    return min(range(len(arr)), key=lambda i: abs(arr[i] - k))


def get_nearest_lat(lat, min_val, max_val, resolution=0.25):
    arr = np.arange(min_val, max_val, resolution)
    idx = np.abs(arr - lat).argmin()
    return idx


def get_nearest_lon(lon, min_val, max_val, resolution=0.25):
    lon = lon % 360  # Normalize longitude to be within 0-360
    arr = np.arange(min_val, max_val, resolution)
    idx = np.abs(arr - lon).argmin()
    return idx


# def getNearestLat(lat, min, max, res=0.25):
#     """Determines the nearest lattitude (to .25 degrees)"""
#     arr = np.arange(start=min, stop=max, step=res)
#     i = closest(arr, lat)
#     return i


# def getNearestLon(lon, min, max, res=0.25):
#     """Determines the nearest longitude (to .25 degrees)"""
#     lon = lon % 360  # convert from -180-180 to 0-360
#     arr = np.arange(start=min, stop=max, step=res)

#     i = closest(arr, lon)
#     return i


def find_model_time_index(time, model_start_time):
    """Calculate the index for a given time, based on the GFS model start time."""
    time_difference = time - model_start_time
    # Convert time difference to hours and then find the index (3 hours interval)
    return int((time_difference.total_seconds() // 3600) // 3)


def parse_time(time_str):
    """Convert the time string to datetime object."""
    return datetime.strptime(time_str, "%d.%m.%Y %H:%M")


def calculate_indices(flight_start_time_str, flight_duration, model_start_time):
    flight_start_time = parse_time(flight_start_time_str)
    flight_duration_hours = ceil(flight_duration)  # Round up to the nearest hour

    flight_end_time = flight_start_time + timedelta(hours=flight_duration_hours)

    # Normalize start time to the nearest previous 3-hour mark
    normalized_start_hour = flight_start_time.hour - (flight_start_time.hour % 3)
    normalized_start_time = flight_start_time.replace(
        hour=normalized_start_hour, minute=0, second=0, microsecond=0
    )

    # Normalize end time to the nearest upcoming 3-hour mark if it exceeds the hour boundary
    if (
        flight_end_time.minute > 0
        or flight_end_time.second > 0
        or flight_end_time.microsecond > 0
    ):
        normalized_end_time = flight_end_time + timedelta(
            hours=3 - (flight_end_time.hour % 3)
        )
    else:
        normalized_end_hour = flight_end_time.hour - (flight_end_time.hour % 3)
        normalized_end_time = flight_end_time.replace(
            hour=normalized_end_hour, minute=0, second=0, microsecond=0
        )

    # Calculate indices
    start_index = find_model_time_index(normalized_start_time, model_start_time)

    end_index = find_model_time_index(normalized_end_time, model_start_time)
    return start_index, end_index


def get_limits(lon_start, lat_start, lon_dist, lat_dist):
    lat_indices = get_grid_indices_from_location(
        lon_start, lat_start, lon_dist, lat_dist
    )
    x_limits_min, x_limits_max, y_limits_min, y_limits_max = (
        find_contiguous_groups_with_limits(lat_indices)
    )
    return x_limits_min, x_limits_max, y_limits_min, y_limits_max


def get_grid_indices_from_location(lon_start, lat_start, lon_dist, lat_dist):
    # Define the resolution
    lon_res = 0.25
    lat_res = 0.25

    # Calculate the number of points
    num_lon_points = 1440
    num_lat_points = 721

    # Normalize the starting longitude and latitude values
    lon_start = lon_start % 360
    lat_start = max(min(lat_start, 90), -90)

    # Convert the starting longitude and latitude to indices
    lon_start_idx = int(lon_start / lon_res)
    lat_start_idx = int((lat_start + 90) / lat_res)

    # Convert distances to index ranges
    lon_range = int(lon_dist / lon_res)
    lat_range = int(lat_dist / lat_res)

    # Calculate longitude indices considering wrapping
    lon_indices = [
        (lon_start_idx + i) % num_lon_points for i in range(-lon_range, lon_range + 1)
    ]

    # Calculate latitude indices and handle pole overflow
    lat_indices = set()
    for j in range(-lat_range, lat_range + 1):
        lat_idx = lat_start_idx + j
        if lat_idx < 0:
            # overflow = abs(lat_idx)
            overflow = abs(lat_idx) - 1
            lat_idx = overflow
            for lon_idx in lon_indices:
                lat_indices.add(
                    (lat_idx, (lon_idx + 720) % num_lon_points)
                )  # Wrap longitude 180 degrees
        elif lat_idx >= num_lat_points:
            # overflow = lat_idx - (num_lat_points - 1)
            overflow = lat_idx - num_lat_points
            lat_idx = (num_lat_points - 1) - overflow
            for lon_idx in lon_indices:
                lat_indices.add(
                    (lat_idx, (lon_idx + 720) % num_lon_points)
                )  # Wrap longitude 180 degrees
        else:
            for lon_idx in lon_indices:
                lat_indices.add((lat_idx, lon_idx))
    return lat_indices


def find_contiguous_groups_with_limits(list_lat_indices):
    # Convert the list to a numpy array
    arr = np.array(list(list_lat_indices))

    # Sort array by the second column
    arr = arr[arr[:, 1].argsort()]

    # Calculate the difference in the second column
    differences = np.diff(arr[:, 1])

    # Find the indices where the difference is greater than 1
    split_indices = np.where(differences > 1)[0] + 1

    # Split the array at the found indices
    groups = np.split(arr, split_indices)

    # Initialize the limits
    lat_limits_min = []
    lat_limits_max = []
    lon_limits_min = []
    lon_limits_max = []

    # Calculate limits for each group
    for group in groups:
        lat_limits_min.append(group[:, 0].min())
        lat_limits_max.append(group[:, 0].max())
        lon_limits_min.append(group[:, 1].min())
        lon_limits_max.append(group[:, 1].max())
    return lat_limits_min, lat_limits_max, lon_limits_min, lon_limits_max
