# SPDX-License-Identifier: GPL-3.0-or-later
# FileType: SOURCE
# SPDX-FileCopyrightText: 2025 GFZ Helmholtz Centre for Geosciences
# SPDX-FileCopyrightText: 2025, "Eva Boergens" at GFZ Potsdam
"""Grid utils module."""

from logging import Logger
from typing import Union

import numpy as np
from scipy.spatial import cKDTree

from .utils import get_constant
from ..datamodels.grids import Grid2DObject, Grid3DObject


def standardize_lat_lon_grid(logger=Logger,
                             longitudes=None,
                             latitudes=None,
                             grid: Union[Grid2DObject, Grid3DObject, None] = None,
                             target_lon_format: str = "[-180, 180]",
                             target_lat_order: str = "ascending"
                             ):
    """
    Check and reformat longitude, latitude, and optionally associated grids.

    This function supports Grid2DObject (lat, lon) and Grid3DObject (time, lat, lon) and ensures their
    longitude and latitude definitions follow the desired conventions.

    Parameters
    ----------
    logger: Logger, logger object to log the error messages
    longitudes : array-like
    1D array of longitude values (can be in 0–360 or -180–180 format)
    latitudes : array-like
    1D array of latitude values (can be in 90 to -90 or -90 to 90 format)
    grid : Grid2DObject or Grid3DObject, optional
    The grid of values aligned with lat/lon. If None, only coordinates are standardized.
    target_lon_format : str, optional
    Desired longitude format: "[-180, 180]" (default) or "[0, 360]"
    target_lat_order : str, optional
    Desired latitude order:
    - "ascending" → from -90 to 90 (default)
    - "descending" → from 90 to -90

    Returns
    -------
    tuple
        If grid is None:
            (longitudes_fixed, latitudes_fixed)
        If grid is Grid2DObject:
            Grid2DObject with updated coords and grid
        If grid is Grid3DObject:
            Grid3DObject with updated coords and grid

    Notes
    -----
    - Longitude values are wrapped if needed to match the target format.
    - Grid is rearranged accordingly along longitude and/or latitude axes.
    - Latitude arrays are flipped if their order does not match the target.

    """
    longitudes = np.array(longitudes)
    latitudes = np.array(latitudes)

    grid_fixed = None
    dates = None
    if isinstance(grid, Grid2DObject):
        grid_fixed = np.array(grid.grid)
    elif isinstance(grid, Grid3DObject):
        grid_fixed = np.array(grid.grid)
        dates = grid.dates

    # --- Fix longitude range ---
    if target_lon_format == "[0, 360]":
        longitudes_fixed = np.mod(longitudes, 360)
    elif target_lon_format == "[-180, 180]":
        longitudes_fixed = ((longitudes + 180) % 360) - 180
    else:
        message = "Invalid target_lon_format. Use '[0, 360]' or '[-180, 180]'."
        logger.error(message)
        raise ValueError(message)

    # Sort longitudes and rearrange grid accordingly
    sort_idx = np.argsort(longitudes_fixed)
    longitudes_fixed = longitudes_fixed[sort_idx]
    if grid_fixed is not None:
        if grid_fixed.ndim == 2:
            grid_fixed = grid_fixed[:, sort_idx]
        elif grid_fixed.ndim == 3:
            grid_fixed = grid_fixed[:, :, sort_idx]

    # --- Fix latitude order ---
    lat_is_descending = latitudes[0] > latitudes[-1]
    if target_lat_order == "ascending" and lat_is_descending:
        latitudes_fixed = np.flip(latitudes, axis=0)
        if grid_fixed is not None:
            grid_fixed = np.flip(grid_fixed, axis=1 if grid_fixed.ndim == 3 else 0)
    elif target_lat_order == "descending" and not lat_is_descending:
        latitudes_fixed = np.flip(latitudes, axis=0)
        if grid_fixed is not None:
            grid_fixed = np.flip(grid_fixed, axis=1 if grid_fixed.ndim == 3 else 0)
    else:
        latitudes_fixed = latitudes

    # --- Return result ---
    if grid is None:
        return longitudes_fixed, latitudes_fixed
    if isinstance(grid, Grid2DObject):
        return Grid2DObject(logger, grid_fixed, longitudes_fixed, latitudes_fixed)
    if isinstance(grid, Grid3DObject):
        return Grid3DObject(logger, grid_fixed, dates, longitudes_fixed, latitudes_fixed)
    message = "grid must be Grid2DObject, Grid3DObject, or None"
    logger.error(message)
    raise TypeError(message)


def land_ocean_mask_buffer(logger: Logger, lo_mask: Grid2DObject, buffer: float) -> Grid2DObject:
    """
    Produce buffered land-ocean mask.

    Parameters
    ----------
    logger: Logger, logger object to log the error messages
    lo_mask (Grid2DObject): Mask of the land (1 or True) and ocean (0 or False)
    buffer (float): Buffer (in km) around the coast, >=0

    Returns
    -------
    lo_masked_buffered (Grid2DObject): Land ocean mask with extended land of buffer size,
                                       land (1 or True) and ocean (0 or False)

    """
    if buffer < 0:
        message = f'Buffer set to {buffer}, but needs to be >=0.'
        logger.error(message)
        raise ValueError(message)

    # --- Prepare coordinate mesh ---
    lon2d, lat2d = np.meshgrid(lo_mask.lon, lo_mask.lat)

    # --- Extract land and ocean coordinates ---
    land_idx = np.where(lo_mask.grid == 1)
    ocean_idx = np.where(lo_mask.grid == 0)
    land_points = np.column_stack((lat2d[land_idx], lon2d[land_idx]))
    ocean_points = np.column_stack((lat2d[ocean_idx], lon2d[ocean_idx]))

    # --- Convert to 3D Cartesian coordinates (for great-circle distance) ---
    R = get_constant('earths_radius_iers') / 1000  # Earth radius in km

    def _sph2xyz(lat, lon):
        deg2rad = get_constant('deg_2_rad')
        lat, lon = lat * deg2rad, lon * deg2rad
        x = R * np.cos(lat) * np.cos(lon)
        y = R * np.cos(lat) * np.sin(lon)
        z = R * np.sin(lat)
        return np.column_stack((x, y, z))

    land_xyz = _sph2xyz(land_points[:, 0], land_points[:, 1])
    ocean_xyz = _sph2xyz(ocean_points[:, 0], ocean_points[:, 1])

    # --- Build KD-tree for land points in 3D space ---
    tree = cKDTree(land_xyz)

    # --- Query nearest land distance for each ocean point ---
    distances, _ = tree.query(ocean_xyz, k=1)

    # Convert chord distance (in km) to surface distance
    # (Great-circle distance = 2 * R * asin(chord/(2R)))
    great_circle_km = 2 * R * np.arcsin(np.clip(distances / (2 * R), 0, 1))

    # --- Define buffer width ---
    ocean_near_land = great_circle_km <= buffer

    # --- Create buffered mask ---
    lo_masked_buffered = lo_mask.copy(deep=True)
    lo_masked_buffered.grid[ocean_idx[0][ocean_near_land], ocean_idx[1][ocean_near_land]] = 1

    return lo_masked_buffered


def global_weighted_sum(grid: Grid2DObject,
                        radius: float = None) -> float:
    """
    Compute the global area weighted sum of input grid.

    Parameters
    ----------
    grid (Grid2DObject: Grid which is summed up weighted by area
    radius (float): Reference radius to use with the area function. Default set to earths_radius from constants

    Returns
    -------
    weighted sum (float)
    """
    if radius is None:
        radius = get_constant('earths_radius_iers')
    area = grid.get_grid_area(radius)
    weighted_sum = np.nansum(area * grid.grid)
    return weighted_sum


def ellipsoid_radius(lat: float) -> float:
    """
    Compute distance to the center of the ellipsoid for a given geocentric latitude.

    Parameters
    ----------
    lat : float
        Geocentric latitude [radians].

    Returns
    -------
    r : float
        Distance from ellipsoid center [m]
    """
    radius = get_constant("earths_radius_iers")
    flattening = get_constant("earths_flattening_iers")

    radius_equator = radius * np.cos(lat)
    radius_polar = (radius - radius * flattening) * np.sin(lat)

    r = np.sqrt(radius_equator ** 2 + radius_polar ** 2)

    return r


def periodic_lon_index(i: int, nlon: int, is_global: bool) -> int | None:
    """
    Periodic longitude indexing for regular grids.

    Parameters
    ----------
    i : int
        Candidate index (can be outside [0, nlon-1]).
    nlon : int
        Number of longitude samples.
    is_global : bool
        If True, wrap module nlon. If False, return None for out-of-range.

    Returns
    -------
    int | None
        Wrapped index or None if regional and out of bounds.
    """
    if is_global:
        return i % nlon
    if 0 <= i < nlon:
        return i
    return None


def spherical_distance(
    lat1_deg: float, lon1_deg: float, lat2_deg: float, lon2_deg: float, radius_m: float
) -> tuple[float, float]:
    """
    Great-circle distance on a sphere (SI).

    Parameters
    ----------
    lat1_deg, lon1_deg, lat2_deg, lon2_deg : float
        Coordinates in degrees.
    radius_m : float
        Sphere radius [m].

    Returns
    -------
    (distance_m, cos_central_angle) : (float, float)
        Distance in meters and cos(central angle) (clamped).
    """
    # degrees -> radians
    lat1 = np.deg2rad(lat1_deg)
    lon1 = np.deg2rad(lon1_deg)
    lat2 = np.deg2rad(lat2_deg)
    lon2 = np.deg2rad(lon2_deg)

    rinvangle = (
        np.sin(lat1) * np.sin(lat2)
        + np.cos(lat1) * np.cos(lat2) * np.cos(lon1 - lon2)
    )
    rinvangle = float(np.clip(rinvangle, -1.0, 1.0))
    return radius_m * float(np.arccos(rinvangle)), rinvangle


def vincenty_distance(
    lat1_deg: float,
    lon1_deg: float,
    lat2_deg: float,
    lon2_deg: float,
    radius_equator: float | None = None,
    radius_polar: float | None = None,
    tol: float = 1e-12,
    max_iter: int = 200,
) -> float:
    """
    Ellipsoidal geodesic distance via Vincenty's inverse method (meters).

    Parameters
    ----------
    lat1_deg, lon1_deg : float
        Latitude and longitude of point 1 in degrees.
    lat2_deg, lon2_deg : float
        Latitude and longitude of point 2 in degrees.
    radius_equator, radius_polar : float | None
        Semi-major (equatorial) and semi-minor (polar) axes in meters. If None, use constants.
    tol : float
        Convergence tolerance (radians) for the longitude difference.
    max_iter : int
        Maximum iterations.

    Returns
    -------
    float
        Geodesic distance in meters.
    """
    if (lat1_deg == lat2_deg) and (lon1_deg == lon2_deg):
        return 0.0

    if radius_equator is None:
        radius_equator = float(get_constant("earths_radius_iers"))
    if radius_polar is None:
        f_const = float(get_constant("earths_flattening_iers"))
        radius_polar = radius_equator * (1.0 - f_const)

    f = (radius_equator - radius_polar) / radius_equator

    deg_to_rad = np.pi / 180.0
    phi1 = lat1_deg * deg_to_rad
    phi2 = lat2_deg * deg_to_rad
    L = (lon2_deg - lon1_deg) * deg_to_rad

    U1 = np.arctan((1.0 - f) * np.tan(phi1))
    U2 = np.arctan((1.0 - f) * np.tan(phi2))
    sinU1, cosU1 = np.sin(U1), np.cos(U1)
    sinU2, cosU2 = np.sin(U2), np.cos(U2)

    lam = L
    sin_sigma = cos_sigma = sigma = sin_alpha = cos2_alpha = cos2sigma = np.nan

    for _ in range(max_iter):
        sinlam = np.sin(lam)
        coslam = np.cos(lam)

        sin_sigma = np.sqrt(
            (cosU2 * sinlam) ** 2
            + (cosU1 * sinU2 - sinU1 * cosU2 * coslam) ** 2
        )
        if sin_sigma == 0.0:
            return 0.0

        cos_sigma = sinU1 * sinU2 + cosU1 * cosU2 * coslam
        sigma = np.arctan2(sin_sigma, cos_sigma)

        sin_alpha = (cosU1 * cosU2 * sinlam) / sin_sigma
        cos2_alpha = 1.0 - sin_alpha * sin_alpha

        if cos2_alpha == 0.0:
            cos2sigma = 0.0
        else:
            cos2sigma = cos_sigma - 2.0 * sinU1 * sinU2 / cos2_alpha

        C = (f / 16.0) * cos2_alpha * (4.0 + f * (4.0 - 3.0 * cos2_alpha))
        lam_prev = lam
        lam = (
            L
            + (1.0 - C) * f * sin_alpha
            * (sigma + C * sin_sigma * (cos2sigma + C * cos_sigma * (-1.0 + 2.0 * (cos2sigma ** 2))))
        )
        if abs(lam - lam_prev) <= tol:
            break
    # The 'else' belongs to the for-loop: it executes only if no 'break' occurred
    # (i.e., the iteration failed to converge within max_iter).
    else:
        # Fallback: haversine on a reasonable mean radius (meters)
        mean_radius = (2.0 * radius_equator + radius_polar) / 3.0
        dphi = (lat2_deg - lat1_deg) * deg_to_rad
        dlam = (lon2_deg - lon1_deg) * deg_to_rad
        h = np.sin(dphi / 2.0) ** 2 + np.cos(phi1) * np.cos(phi2) * np.sin(dlam / 2.0) ** 2
        return float(2.0 * mean_radius * np.arcsin(np.sqrt(h)))

    u2 = cos2_alpha * (radius_equator * radius_equator - radius_polar * radius_polar) / (radius_polar * radius_polar)
    A = 1.0 + (u2 / 16384.0) * (4096.0 + u2 * (-768.0 + u2 * (320.0 - 175.0 * u2)))
    B = (u2 / 1024.0) * (256.0 + u2 * (-128.0 + u2 * (74.0 - 47.0 * u2)))

    delta_sigma = (
        B * sin_sigma * (
            cos2sigma
            + 0.25 * B * (
                cos_sigma * (-1.0 + 2.0 * (cos2sigma ** 2))
                - (B / 6.0) * cos2sigma * (-3.0 + 4.0 * (sin_sigma ** 2)) * (-3.0 + 4.0 * (cos2sigma ** 2))
            )
        )
    )

    return float(radius_polar * A * (sigma - delta_sigma))


def anisotropy_scaled_distance_m(
    lat_src_deg: float, lon_src_deg: float, lat_tgt_deg: float, lon_tgt_deg: float,
    radius_m: float, flattening: float,
) -> float:
    """
    Spherical great-circle distance scaled for ellipsoidal anisotropy (SI).

    Parameters
    ----------
    lat_src_deg, lon_src_deg, lat_tgt_deg, lon_tgt_deg : float
        Coordinates in degrees.
    radius_m : float
        Local spherical-equivalent radius at source [m].
    flattening : float
        flattening = (radius_equator - radius_polar)/radius_equator (dimensionless).

    Returns
    -------
    float
        Anisotropy-scaled path length [m].
    """
    # Great-circle on a sphere with the provided radius (meters).
    dist_m, _ = spherical_distance(lat_src_deg, lon_src_deg, lat_tgt_deg, lon_tgt_deg, radius_m)
    if dist_m == 0.0:
        return 0.0

    deg2rad = np.pi / 180.0
    lat_s = lat_src_deg * deg2rad
    lon_s = lon_src_deg * deg2rad
    lat_t = lat_tgt_deg * deg2rad
    lon_t = lon_tgt_deg * deg2rad

    s = dist_m / max(radius_m, 1e-16)
    if s == 0.0:
        return dist_m

    rinv = (
        np.sin(lat_t) * np.cos(lat_s)
        - np.cos(lat_t) * np.sin(lat_s) * np.cos(lon_s - lon_t)
    ) / np.sin(s)
    rinv = float(np.clip(rinv, -1.0, 1.0))
    azimut = float(np.arccos(rinv))

    g = 1.0 / ((1.0 - flattening) ** 2) - 1.0
    dist_scale = 1.0 / np.sqrt(1.0 + g * (np.cos(azimut) ** 2))
    return float(dist_m * dist_scale)
