import numpy as np
import xarray as xr
import os
import pandas as pd
import fsspec
import s3fs

from sklearn.neighbors import BallTree
from tqdm import tqdm
from .utils import create_folder_indice

def extract_goes_production(time_choice, polar_data, parallel_index,channels,goes_aws_url_folder):

    """ 
    Extracts GOES data for a specific time from the polar data and returns the images along with valid latitudes, longitudes, and times.

    Args:
        time_choice (str): The time for which to extract GOES data, in 'YYYY-MM-DD HH:MM:SS' format.
        polar_data (xarray.Dataset): Polar data containing latitude and longitude information. Used to create a grid of valid latitudes and longitudes.
        parallel_index (int): Index for parallel processing, used to identify the specific GOES data to extract, generated by the `index_parallel` function.
        channels (list): List of GOES channels to extract.
        goes_aws_url_folder (str): AWS URL folder for the GOES data, default is "noaa-goes16/ABI-L2-CMIPF".
    
    Returns:
        images (np.ndarray): Array of extracted GOES images for the specified time.
        valid_lats (np.ndarray): Array of valid latitudes corresponding to the GOES images
        valid_lons (np.ndarray): Array of valid longitudes corresponding to the GOES images
        valid_times (np.ndarray): Array of valid times corresponding to the GOES images.

    """
    time_formated = (
        np.datetime64(time_choice).astype("datetime64[ns]").astype("float64")
    )

    longrid, latgrid = np.meshgrid(polar_data["longitude"], polar_data["latitude"])
    lon_array = longrid.flatten()
    lat_array = latgrid.flatten()
    time_array = np.full_like(lon_array, time_formated)

    valid_lons = lon_array
    valid_lats = lat_array
    valid_times = time_array

    print('INFO : Extracting GOES data')
    images = extract_goes_inference(np.datetime64(time_choice), parallel_index,channels,goes_aws_url_folder)


    images = np.expand_dims(images, axis=1)


    return images, valid_lats, valid_lons, valid_times

def get_indices(lat_grid, lon_grid, Goeslat, Goeslon, radius=0.125):
    """
    Finds the corresponding GOES row and column indices for each scatterometer point
    using a BallTree for efficiency, and then filtering points to form a square bounding box.

    Args:
        lat_grid (numpy.ndarray): 2D array of latitudes from the scatterometer data.
        lon_grid (numpy.ndarray): 2D array of longitudes from the scatterometer data.
        Goeslat (numpy.ndarray): 2D array of latitudes from the GOES data.
        Goeslon (numpy.ndarray): 2D array of longitudes from the GOES data.
        radius (float): Radius in degrees to define the bounding box around each scatterometer point.
    Returns:
        indices_array (numpy.ndarray): 2D array of tuples, where each tuple contains the row and column indices of the corresponding GOES pixel for each scatterometer point.
        

    """

    print("INFO : Calculating indices")
    # Flatten GOES data
    Goeslat_flat = Goeslat.flatten()
    Goeslon_flat = Goeslon.flatten()
    goes_points = np.column_stack((Goeslat_flat, Goeslon_flat))

    # Build BallTree with haversine distance
    goes_points_rad = np.radians(goes_points)
    goes_tree = BallTree(goes_points_rad, metric="haversine")

    # Flatten scatter grids
    lat_flat = lat_grid.flatten()
    lon_flat = lon_grid.flatten()
    scatter_points = np.column_stack((lat_flat, lon_flat))
    scatter_points_rad = np.radians(scatter_points)

    # Radius for broad-phase query: diagonal of the bounding box
    # Square box ±radius: diagonal = radius * sqrt(2)
    diag_radius = radius * np.sqrt(2)
    diag_radius_rad = np.radians(diag_radius)

    indices_array = np.empty(lat_flat.shape, dtype=object)
    goes_shape = Goeslat.shape

    for i, (lat_val, lon_val) in enumerate(zip(lat_flat, lon_flat)):
        # Broad-phase: query all points within diagonal distance
        candidate_indices = goes_tree.query_radius(
            np.array([scatter_points_rad[i]]), r=diag_radius_rad
        )[0]

        if candidate_indices.size == 0:
            # No points found, store empty
            indices_array[i] = (np.array([], dtype=int), np.array([], dtype=int))
            continue

        # Post-filter candidates to keep only those in the bounding box
        lat_min = lat_val - radius
        lat_max = lat_val + radius
        lon_min = lon_val - radius
        lon_max = lon_val + radius

        cand_lats = Goeslat_flat[candidate_indices]
        cand_lons = Goeslon_flat[candidate_indices]

        mask = (
            (cand_lats >= lat_min)
            & (cand_lats <= lat_max)
            & (cand_lons >= lon_min)
            & (cand_lons <= lon_max)
        )

        final_indices = candidate_indices[mask]

        # Convert these flat indices back to row,col
        rows, cols = np.unravel_index(final_indices, goes_shape)
        indices_array[i] = (rows, cols)

    # Reshape indices_array to the original shape
    indices_array = indices_array.reshape(lat_grid.shape)
    return indices_array



def calculate_degrees(file_id):
    """This function calculates the latitude and longitude of the GOES ABI fixed grid projection. 
    This function comes from NOAA/NESDIS/STAR. (2025). Latitude and longitude remapping of GOES-R ABI imagery using Python . Atmospheric Composition Science Team. Retrieved from https://www.star.nesdis.noaa.gov/atmospheric-composition-training/python_abi_lat_lon.php

    Args:
        file_id (xarray.Dataset): The xarray dataset containing the GOES ABI fixed grid projection variables.

    Returns:
        abi_lat (numpy.ndarray): The latitude of the GOES ABI fixed grid projection.
        abi_lon (numpy.ndarray): The longitude of the GOES ABI fixed grid projection.
    

    """

    # Read in GOES ABI fixed grid projection variables and constants
    x_coordinate_1d = file_id.variables["x"][:]  # E/W scanning angle in radians
    y_coordinate_1d = file_id.variables["y"][:]  # N/S elevation angle in radians
    projection_info = file_id.goes_imager_projection
    lon_origin = projection_info.longitude_of_projection_origin
    H = projection_info.perspective_point_height + projection_info.semi_major_axis
    r_eq = projection_info.semi_major_axis
    r_pol = projection_info.semi_minor_axis

    # Create 2D coordinate matrices from 1D coordinate vectors
    x_coordinate_2d, y_coordinate_2d = np.meshgrid(x_coordinate_1d, y_coordinate_1d)

    # Equations to calculate latitude and longitude
    lambda_0 = (lon_origin * np.pi) / 180.0
    a_var = np.power(np.sin(x_coordinate_2d), 2.0) + (
        np.power(np.cos(x_coordinate_2d), 2.0)
        * (
            np.power(np.cos(y_coordinate_2d), 2.0)
            + (
                ((r_eq * r_eq) / (r_pol * r_pol))
                * np.power(np.sin(y_coordinate_2d), 2.0)
            )
        )
    )
    b_var = -2.0 * H * np.cos(x_coordinate_2d) * np.cos(y_coordinate_2d)
    c_var = (H**2.0) - (r_eq**2.0)
    r_s = (-1.0 * b_var - np.sqrt((b_var**2) - (4.0 * a_var * c_var))) / (2.0 * a_var)
    s_x = r_s * np.cos(x_coordinate_2d) * np.cos(y_coordinate_2d)
    s_y = -r_s * np.sin(x_coordinate_2d)
    s_z = r_s * np.cos(x_coordinate_2d) * np.sin(y_coordinate_2d)

    # Ignore numpy errors for sqrt of negative number; occurs for GOES-16 ABI CONUS sector data
    np.seterr(all="ignore")

    abi_lat = (180.0 / np.pi) * (
        np.arctan(
            ((r_eq * r_eq) / (r_pol * r_pol))
            * ((s_z / np.sqrt(((H - s_x) * (H - s_x)) + (s_y * s_y))))
        )
    )
    abi_lon = (lambda_0 - np.arctan(s_y / (H - s_x))) * (180.0 / np.pi)

    print("INFO : Latitude and longitude calculated")
    return abi_lat, abi_lon


def index_parallel(ds, ScatterDataset):
    """
    Finds the corresponding GOES row and column indices for the entire scatterometer dataset.

    Args:
        ScatterDataset: xarray Dataset containing scatterometer data.
        scatter_name: Name for the output file.
        output_path: Path to save the output file.

    Returns:
        parallel_indice_values: 2D array of tuples containing GOES row and column indices corresponding to scatterometer data.
    """

    create_folder_indice("satellite_indices")

    ds_spatial_resolution = ds.spatial_resolution
    ds_spatial_resolution.replace(" ", "_")

    name_str = f"lat_{ScatterDataset.latitude.min().values}_{ScatterDataset.latitude.max().values}_lon_{ScatterDataset.longitude.min().values}_{ScatterDataset.longitude.max().values}_res_{ds_spatial_resolution}"
    name_str = name_str.replace(".", "_")
    if os.path.exists(
        f"./satellite_indices/{ds_spatial_resolution}_index.npy"
    ):
        parallel_index = np.load(
            f"./satellite_indices/{ds_spatial_resolution}_index.npy",
            allow_pickle=True,
        )

        return parallel_index

    else:
        print(
            "INFO : Satellite index file not found, creating new index file. This might take a while."
        )

        # Extract scatterometer latitudes and longitudes
        Latitudes_Scatter = ScatterDataset["latitude"].values
        Longitudes_Scatter = ScatterDataset["longitude"].values

        # Create a meshgrid of scatterometer coordinates
        lon_grid, lat_grid = np.meshgrid(Longitudes_Scatter, Latitudes_Scatter)

        # Extract GOES latitudes and longitudes
        Goeslat, Goeslon = calculate_degrees(ds)
        Goeslat[np.isnan(Goeslat)] = 999
        Goeslon[np.isnan(Goeslon)] = 999
        # Use the optimized get_indices function
        parallel_indice_values = get_indices(lat_grid, lon_grid, Goeslat, Goeslon)

        # Save the indices array
        np.save(
            f"./satellite_indices/{ds_spatial_resolution}_index.npy",
            parallel_indice_values,
        )

        return parallel_indice_values


def get_goes_url(time, goes_aws_url_folder='noaa-goes16/ABI-L2-CMIPF', goes_channel="C01"):
    """
    This function gets the nearest GOES-16 files from the time given.
    The function returns a list of urls to the files.
    The function uses the s3fs library to access the AWS GOES-16 data.

    Args:
        time (numpy.datetime[ns]): The time of the scatterometer data.
        goes_aws_url_folder (str): The folder in the AWS S3 bucket where the GOES-16 data is stored.
        goes_channel (list): The channel of interest.

    Returns:
        urls (list): A list of urls to the GOES-16 files.


    """
    date_c = time.astype("datetime64[ns]")
    date = pd.to_datetime(date_c)
    date_str = date.strftime("%Y/%j/%H")
    min = int(date.strftime("%M"))
    min_range = [(min + i) % 60 for i in range(-6, 7)]
    min_range_str = [f"{x:02d}" for x in min_range]
    fs = s3fs.S3FileSystem(anon=True)
    # get the nearest goes file from time

    urls = []
    channel = goes_channel
    path = f"{goes_aws_url_folder}/{date_str}"
    files = fs.ls(path)
    filter_channel = [x for x in files if channel in x]
    if len(filter_channel) == 0:
        print(f"INFO :No file found for {channel} on day {date_str}, skipping file")
        return
    file = [x for x in filter_channel if x[73:75] in min_range_str]
    if len(file) == 0:
        print(
            f"INFO :No file found for {channel} on day {date_str} for minute {min}, skipping file"
        )
        return np.zeros(len(goes_channel))
    urls.append("s3://" + file[0])

    return urls


def get_image(ds, parallel_index, lat_grd, lon_grd, lat_search, lon_search,goes_image_size=128):

    """
    This function retrieves a trainable GOES image for a given latitude and longitude from a GOES16 `.nc` file.

    Args:
        ds (xarray.Dataset): The xarray dataset containing the GOES data.
        parallel_index (numpy.ndarray): The precomputed indices for GOES pixels corresponding to scatterometer lat/lon.
        lat_grd (numpy.ndarray): The latitude grid of the scatterometer data.
        lon_grd (numpy.ndarray): The longitude grid of the scatterometer data.
        lat_search (float): The latitude to search for in the GOES data.
        lon_search (float): The longitude to search for in the GOES data.
        goes_image_size (int): The size of the output image. Default is 128.

    Returns:
        padded_image (xarray.DataArray): A padded xarray DataArray containing the GOES image centered around the specified lat/lon.

    """
    index_row = np.where(
        lat_grd == lat_search,
    )
    index_column = np.where(lon_grd == lon_search)

    rows_goes = parallel_index[index_row[0][0], index_column[0][0]][0]
    columns_goes = parallel_index[index_row[0][0], index_column[0][0]][1]

    if rows_goes.size == 0 or columns_goes.size == 0:
        return None

    pixels_from_center = (goes_image_size-1) // 2
    mean_row = rows_goes.mean().astype(int)
    min_row = mean_row - pixels_from_center
    max_row = mean_row + pixels_from_center

    mean_col = columns_goes.mean().astype(int)
    min_col = mean_col - pixels_from_center
    max_col = mean_col + pixels_from_center

    if "CMI" in ds: # If using GOES-16 L2 processed data
        image = ds.CMI[min_row:max_row, min_col:max_col].values

    elif "Rad" in ds: #If using GOES-16 L1b data
        image = ds.Rad[min_row:max_row, min_col:max_col].values

    # debug
    # print(min_row,'= min_row', max_row,'= max_row', min_col, '= min_col', max_col, '= max_col')
    target_size = (goes_image_size, goes_image_size)

    padded_image = np.pad(
        image,
        (
            (
                (target_size[0] - image.shape[0]) // 2,
                (target_size[0] - image.shape[0] + 1) // 2,
            ),
            (
                (target_size[1] - image.shape[1]) // 2,
                (target_size[1] - image.shape[1] + 1) // 2,
            ),
        ),
        constant_values=0,
    )

    padded_image = xr.DataArray(padded_image, dims=("x", "y"))
    return padded_image

def extract_goes(
    observation_times,
    observation_lats,
    observation_lons,
    scatterometer_data_path,
    goes_aws_url_folder,
    goes_channel="C01",
    goes_image_size=128,
    verbose=True,
):
    """
    This function extracts GOES images for the given observation times, latitudes, and longitudes.
    It retrieves the GOES data from the specified AWS S3 bucket and processes it to create images of the specified size.

    Args:
        observation_times (numpy.ndarray): The times of observation of the scatterometer data. 
        observation_lats (numpy.ndarray): The latitudes of the scatterometer data.
        observation_lons (numpy.ndarray): The longitudes of the scatterometer data.
        scatterometer_data_path (str): The path to the scatterometer data directory.
        goes_aws_url_folder (str): The folder in the AWS S3 bucket where the GOES data is stored.
        goes_channel (str): The channel of interest. Default is "C01".
        goes_image_size (int): The size of the output images. Default is 128.
        verbose (bool): If True, prints progress information.
    
    Returns:
        images (numpy.ndarray): A 4D numpy array of shape (num_observations, num_channels, goes_image_size, goes_image_size) containing the extracted GOES images.

    """

    for file in os.listdir(scatterometer_data_path):
        if file.endswith(".nc"):
            polar = xr.open_dataset(
                os.path.join(scatterometer_data_path, file),
                engine="h5netcdf",
                drop_variables=["DQF"],
            )
            break
        
        else:
            print('WARNING : No .nc file found in the scatterometer data path, please check the path')

    template_scatter = polar.isel(time=0)
    lat_grd, lon_grd = (
        template_scatter["latitude"].values,
        template_scatter["longitude"].values,
    )

    fs = fsspec.filesystem("s3", anon=True, default_block_size=512 * 1024**1024)

    values, counts = np.unique(observation_times, return_counts=True)

    all_urls = []  # getting unique URLS
    for value in values:
        urls = get_goes_url(value, goes_aws_url_folder,goes_channel)
        all_urls.append(urls)

    values_url, indices_url, counts_url = np.unique(
        all_urls, return_index=True, return_counts=True, axis=0
    )
    # Sort indices to "unsort" the URLs
    sorted_indices = sorted(range(len(indices_url)), key=lambda k: indices_url[k])
    values_url = [all_urls[indices_url[i]] for i in sorted_indices]

    # Reorder counts_url using the same sorted indices
    counts_url = [counts_url[i] for i in sorted_indices]

    compressed_urls = values_url
    compressed_counts = []
    start_idx = 0

    for size in counts_url:
        group_sum = counts[start_idx : start_idx + size].sum()
        compressed_counts.append(group_sum)
        start_idx += size

    width = goes_image_size
    height = goes_image_size

    images = np.zeros([len(observation_times), 1 , width, height], dtype=np.float32)

    total_idx = 0
    for unique_idx, unique_urls in tqdm(
        enumerate(compressed_urls),
        desc="INFO : Retrieving and processing GOES data",
        total=len(compressed_urls),
        disable=not verbose,
    ):

    
        for CH_idx, url_CH in enumerate(unique_urls):

            if url_CH == 0:
                images[total_idx, CH_idx] = np.zeros([width, height])
                continue

            with fs.open(url_CH, mode="rb") as f:

                ds = xr.open_dataset(
                    f, engine="h5netcdf", drop_variables=["DQF"]
                )  # this is the bottleneck

                parallel_index = index_parallel(
                    ds,
                    template_scatter,
                )
                for i in range(compressed_counts[unique_idx]):
                    images[total_idx + i, CH_idx] = get_image(
                        ds=ds,
                        parallel_index=parallel_index,
                        lat_grd=lat_grd,
                        lon_grd=lon_grd,
                        lat_search=observation_lats[total_idx + i],
                        lon_search=observation_lons[total_idx + i],
                        goes_image_size=goes_image_size,
                    )

        total_idx += compressed_counts[unique_idx]

    if verbose:
        print(
            f"INFO : Extracted {len(observation_times)} images from {len(compressed_urls)} GOES files."
        )
    return images




def goes_index(parallel_index, lat_grd, lon_grd, lat_search, lon_search):
    """
    This function retrieves the indices of the GOES image corresponding to a given latitude and longitude. This is an archived function. Current implementation decides on extent based on chosen image size.

    Args:
        parallel_index (numpy.ndarray): The precomputed indices for GOES pixels corresponding to scatterometer lat/lon.
        lat_grd (numpy.ndarray): The latitude grid of the scatterometer data.
        lon_grd (numpy.ndarray): The longitude grid of the scatterometer data.
        lat_search (float): The latitude to search for in the GOES data.
        lon_search (float): The longitude to search for in the GOES data.

    Returns:
        min_row (int): The minimum row index of the GOES image.
        max_row (int): The maximum row index of the GOES image.
        min_col (int): The minimum column index of the GOES image.
        max_col (int): The maximum column index of the GOES image.
    """

    index_row = np.where(lat_grd == lat_search)
    index_column = np.where(lon_grd == lon_search)

    rows_goes = parallel_index[index_row[0][0], index_column[0][0]][0]
    columns_goes = parallel_index[index_row[0][0], index_column[0][0]][1]

    if rows_goes.size == 0 or columns_goes.size == 0:
        return None

    min_row = rows_goes.min()
    max_row = rows_goes.max()

    min_col = columns_goes.min()
    max_col = columns_goes.max()

    return min_row, max_row, min_col, max_col


def extract_goes_inference(date_time, parallel_index,channels="C01",goes_aws_url_folder= 'noaa-goes16/ABI-L2-CMIPF'):
    """
    This function extracts GOES images for a given date_time and parallel_index. (whole GOES slice, used for inference which differs from images used in training that have a matched orbit with scatterometers.)
    It retrieves the GOES data from the specified AWS S3 bucket and processes it to create
    images of the specified size (128x128).

    Args:
        date_time (numpy.datetime64): The time of the GOES data.
        parallel_index (numpy.ndarray): The precomputed indices for GOES pixels corresponding to scatterometer lat/lon.
        channels (str or list): The channel(s) of interest. Default is "C01".
        goes_aws_url_folder (str): The folder in the AWS S3 bucket where the GOES data is stored. Default is 'noaa-goes16/ABI-L2-CMIPF'.
        
    Returns:
        images (list): A list of numpy arrays containing the extracted GOES images of shape (128, 128).
    """

    # ignore divide by zero errors which occur when the GOES data can't form a 128x128 image
    np.seterr(invalid='ignore', divide='ignore')

    fs = fsspec.filesystem("s3", anon=True, default_block_size=512 * 1024**1024)
    urls = get_goes_url(date_time, goes_aws_url_folder='noaa-goes16/ABI-L2-CMIPF', goes_channel= channels)
    with fs.open(urls[0], mode="rb") as f:
        print("INFO : Reading file:", urls[0])
        goes_image = xr.open_dataset(f)
        goes_image = goes_image.rename({"x": "x_index", "y": "y_index"})

        # Assign the index coordinates (if not already done)
        goes_image = goes_image.assign_coords(
            x_index=np.arange(goes_image.sizes["x_index"]),
            y_index=np.arange(goes_image.sizes["y_index"]),
        )

        images = []
        goes_image.load()
        print("INFO : Extracting images")
        for i in range(parallel_index.shape[0]):
            for j in range(parallel_index.shape[1]):
                try:
                    x_mean = parallel_index[i][j][1].mean().astype(int)
                    x_min = x_mean - 63
                    x_max = x_mean + 63
                    y_mean = parallel_index[i][j][0].mean().astype(int)
                    y_min = y_mean - 63
                    y_max = y_mean + 63
                    image = goes_image.CMI.sel(
                        x_index=slice(x_min, x_max), y_index=slice(y_min, y_max)
                    )

                    target_size = (128, 128)

                    padded_image = np.pad(
                        image,
                        (
                            (
                                (target_size[0] - image.shape[0]) // 2,
                                (target_size[0] - image.shape[0] + 1) // 2,
                            ),
                            (
                                (target_size[1] - image.shape[1]) // 2,
                                (target_size[1] - image.shape[1] + 1) // 2,
                            ),
                        ),
                        constant_values=0,
                    )

                except:
                    images.append(np.zeros((128, 128)))

                    continue
                images.append(padded_image)

        return images
