"""Core utilities, types, and internal functions for cyclogeostrophy computations."""

from typing import NamedTuple

import jax
import jax.numpy as jnp
from jaxtyping import Float

from ..geostrophy import geostrophy
from ..utils import geometry, kinematics, operators, sanitize


# =============================================================================
# Types
# =============================================================================

class CyclogeostrophySetup(NamedTuple):
    """
    Contains precomputed values needed for cyclogeostrophic calculations.

    Attributes
    ----------
    is_land : Float[jax.Array, "lat lon"]
        Land mask where `True`/`1` indicates land
    ug_u : Float[jax.Array, "lat lon"]
        U component of geostrophic velocity on U grid
    vg_v : Float[jax.Array, "lat lon"]
        V component of geostrophic velocity on V grid
    lat_u : Float[jax.Array, "lat lon"]
        Latitudes of the U grid
    lon_u : Float[jax.Array, "lat lon"]
        Longitudes of the U grid
    lat_v : Float[jax.Array, "lat lon"]
        Latitudes of the V grid
    lon_v : Float[jax.Array, "lat lon"]
        Longitudes of the V grid
    dx_u : Float[jax.Array, "lat lon"]
        Spatial steps in meters along x on U grid
    dy_u : Float[jax.Array, "lat lon"]
        Spatial steps in meters along y on U grid
    dx_v : Float[jax.Array, "lat lon"]
        Spatial steps in meters along x on V grid
    dy_v : Float[jax.Array, "lat lon"]
        Spatial steps in meters along y on V grid
    coriolis_factor_u : Float[jax.Array, "lat lon"]
        Coriolis factor on U grid
    coriolis_factor_v : Float[jax.Array, "lat lon"]
        Coriolis factor on V grid
    coriolis_factor_t : Float[jax.Array, "lat lon"]
        Coriolis factor on T grid
    """

    is_land: Float[jax.Array, "lat lon"]
    ug_u: Float[jax.Array, "lat lon"]
    vg_v: Float[jax.Array, "lat lon"]
    lat_u: Float[jax.Array, "lat lon"]
    lon_u: Float[jax.Array, "lat lon"]
    lat_v: Float[jax.Array, "lat lon"]
    lon_v: Float[jax.Array, "lat lon"]
    dx_u: Float[jax.Array, "lat lon"]
    dy_u: Float[jax.Array, "lat lon"]
    dx_v: Float[jax.Array, "lat lon"]
    dy_v: Float[jax.Array, "lat lon"]
    coriolis_factor_u: Float[jax.Array, "lat lon"]
    coriolis_factor_v: Float[jax.Array, "lat lon"]
    coriolis_factor_t: Float[jax.Array, "lat lon"]


class CyclogeostrophyResult(NamedTuple):
    """
    Result of cyclogeostrophic velocity computation.

    This NamedTuple provides named access to results, avoiding positional unpacking errors.
    All fields except ``ucg`` and ``vcg`` are optional and depend on the
    ``return_*`` flags passed to the computation function.

    Attributes
    ----------
    ucg : Float[jax.Array, "lat lon"]
        U component of cyclogeostrophic velocity on U grid
    vcg : Float[jax.Array, "lat lon"]
        V component of cyclogeostrophic velocity on V grid
    ug : Float[jax.Array, "lat lon"] | None
        U component of geostrophic velocity on U grid (if ``return_geos=True``)
    vg : Float[jax.Array, "lat lon"] | None
        V component of geostrophic velocity on V grid (if ``return_geos=True``)
    lat_u : Float[jax.Array, "lat lon"] | None
        Latitudes of U grid (if ``return_grids=True``)
    lon_u : Float[jax.Array, "lat lon"] | None
        Longitudes of U grid (if ``return_grids=True``)
    lat_v : Float[jax.Array, "lat lon"] | None
        Latitudes of V grid (if ``return_grids=True``)
    lon_v : Float[jax.Array, "lat lon"] | None
        Longitudes of V grid (if ``return_grids=True``)
    losses : Float[jax.Array, "n_it"] | None
        Cyclogeostrophic imbalance over iterations (if ``return_losses=True``)
    """

    ucg: Float[jax.Array, "lat lon"]
    vcg: Float[jax.Array, "lat lon"]
    ug: Float[jax.Array, "lat lon"] | None = None
    vg: Float[jax.Array, "lat lon"] | None = None
    lat_u: Float[jax.Array, "lat lon"] | None = None
    lon_u: Float[jax.Array, "lat lon"] | None = None
    lat_v: Float[jax.Array, "lat lon"] | None = None
    lon_v: Float[jax.Array, "lat lon"] | None = None
    losses: Float[jax.Array, "n_it"] | None = None


# =============================================================================
# Setup and Result Assembly
# =============================================================================

def setup_cyclogeostrophy(
    ssh_t: Float[jax.Array, "lat lon"],
    lat_t: Float[jax.Array, "lat lon"],
    lon_t: Float[jax.Array, "lat lon"],
    mask: Float[jax.Array, "lat lon"] = None
) -> CyclogeostrophySetup:
    """
    Computes all preliminary values needed for cyclogeostrophic calculations.

    This includes: land mask, geostrophic velocities, U/V grids, spatial steps,
    and Coriolis factors on all grids.

    Parameters
    ----------
    ssh_t : Float[jax.Array, "lat lon"]
        SSH field on T grid
    lat_t : Float[jax.Array, "lat lon"]
        Latitudes of T grid
    lon_t : Float[jax.Array, "lat lon"]
        Longitudes of T grid
    mask : Float[jax.Array, "lat lon"], optional
        Land mask where `1`/`True` is land. If None, inferred from ssh_t nan values.

    Returns
    -------
    CyclogeostrophySetup
        Named tuple containing all precomputed values
    """
    is_land = sanitize.init_land_mask(ssh_t, mask)

    geos_results = geostrophy(
        ssh_t, lat_t, lon_t, is_land, return_grids=True
    )

    ug_u = geos_results.ug
    vg_v = geos_results.vg
    lat_u = geos_results.lat_u
    lon_u = geos_results.lon_u
    lat_v = geos_results.lat_v
    lon_v = geos_results.lon_v

    dx_u, dy_u = geometry.compute_spatial_step(lat_u, lon_u)
    dx_v, dy_v = geometry.compute_spatial_step(lat_v, lon_v)

    coriolis_factor_u = geometry.compute_coriolis_factor(lat_u)
    coriolis_factor_v = geometry.compute_coriolis_factor(lat_v)
    coriolis_factor_t = geometry.compute_coriolis_factor(lat_t)

    return CyclogeostrophySetup(
        is_land=is_land,
        ug_u=ug_u,
        vg_v=vg_v,
        lat_u=lat_u,
        lon_u=lon_u,
        lat_v=lat_v,
        lon_v=lon_v,
        dx_u=dx_u,
        dy_u=dy_u,
        dx_v=dx_v,
        dy_v=dy_v,
        coriolis_factor_u=coriolis_factor_u,
        coriolis_factor_v=coriolis_factor_v,
        coriolis_factor_t=coriolis_factor_t,
    )


def assemble_result(
    ucg_u: Float[jax.Array, "lat lon"],
    vcg_v: Float[jax.Array, "lat lon"],
    setup: CyclogeostrophySetup,
    return_geos: bool = False,
    return_grids: bool = True,
    return_losses: bool = False,
    losses: Float[jax.Array, "n_it"] = None,
) -> CyclogeostrophyResult:
    """
    Assembles the final result, sanitizing outputs and including optional fields.

    Parameters
    ----------
    ucg_u : Float[jax.Array, "lat lon"]
        U component of cyclogeostrophic velocity
    vcg_v : Float[jax.Array, "lat lon"]
        V component of cyclogeostrophic velocity
    setup : CyclogeostrophySetup
        Precomputed setup values
    return_geos : bool, optional
        Include geostrophic velocities in result
    return_grids : bool, optional
        Include U/V grid coordinates in result
    return_losses : bool, optional
        Include losses in result
    losses : Float[jax.Array, "n_it"], optional
        Loss values from iterative methods

    Returns
    -------
    CyclogeostrophyResult
        Named tuple with computed velocities and optional fields
    """
    ucg_u = sanitize.sanitize_data(ucg_u, jnp.nan, setup.is_land)
    vcg_v = sanitize.sanitize_data(vcg_v, jnp.nan, setup.is_land)

    return CyclogeostrophyResult(
        ucg=ucg_u,
        vcg=vcg_v,
        ug=setup.ug_u if return_geos else None,
        vg=setup.vg_v if return_geos else None,
        lat_u=setup.lat_u if return_grids else None,
        lon_u=setup.lon_u if return_grids else None,
        lat_v=setup.lat_v if return_grids else None,
        lon_v=setup.lon_v if return_grids else None,
        losses=losses if return_losses else None,
    )


# =============================================================================
# Public API Functions
# =============================================================================

def cyclogeostrophic_loss(
    ug: Float[jax.Array, "lat lon"],
    vg: Float[jax.Array, "lat lon"],
    ucg: Float[jax.Array, "lat lon"],
    vcg: Float[jax.Array, "lat lon"],
    lat_t: Float[jax.Array, "lat lon"] = None,
    lon_t: Float[jax.Array, "lat lon"] = None,
    lat_u: Float[jax.Array, "lat lon"] = None,
    lon_u: Float[jax.Array, "lat lon"] = None,
    lat_v: Float[jax.Array, "lat lon"] = None,
    lon_v: Float[jax.Array, "lat lon"] = None,
    mask: Float[jax.Array, "lat lon"] = None,
    vel_on_uv: bool = True
) -> Float[jax.Array, ""]:
    """
    Computes the cyclogeostrophic imbalance loss from a geostrophic SSC velocity field and a cyclogeostrophic SSC velocity field.

    The velocity fields can be provided either on the U and V grids (``vel_on_uv=True``) or on the T grid (``vel_on_uv=False``).

    Parameters
    ----------
    ug : Float[jax.Array, "lat lon"]
        U component of the geostrophic SSC velocity field
    vg : Float[jax.Array, "lat lon"]
        V component of the geostrophic SSC velocity field
    ucg : Float[jax.Array, "lat lon"]
        U component of the cyclogeostrophic SSC velocity field
    vcg : Float[jax.Array, "lat lon"]
        V component of the cyclogeostrophic SSC velocity field
    lat_t : Float[jax.Array, "lat lon"], optional
        Latitudes of the T grid.
        If ``lat_u``, ``lon_u``, ``lat_v``, and ``lon_v`` are not provided, ``lat_t`` and ``lon_t`` must be provided to compute them.
        Defaults to `None`
    lon_t : Float[jax.Array, "lat lon"], optional
        Longitudes of the T grid.
        If ``lat_u``, ``lon_u``, ``lat_v``, and ``lon_v`` are not provided, ``lat_t`` and ``lon_t`` must be provided to compute them.
        Defaults to `None`
    lat_u : Float[jax.Array, "lat lon"], optional
        Latitudes of the U grid.
        Defaults to `None`
    lon_u : Float[jax.Array, "lat lon"], optional
        Longitudes of the U grid.
        Defaults to `None`
    lat_v : Float[jax.Array, "lat lon"], optional
        Latitudes of the V grid.
        Defaults to `None`
    lon_v : Float[jax.Array, "lat lon"], optional
        Longitudes of the V grid.
        Defaults to `None`
    mask : Float[jax.Array, "lat lon"], optional
        Mask defining the marine area of the spatial domain; `1` or `True` stands for masked (i.e. land).
        If not provided, inferred from ``ug`` `nan` values.
        Defaults to `None`
    vel_on_uv : bool, optional
        If `True`, ``ucg`` and ``vcg`` are on the U and V grids.
        If `False`, they are on the T grid.
        Defaults to `True`

    Returns
    -------
    loss : Float[jax.Array, ""]
        Cyclogeostrophic imbalance loss
    """
    if mask is None:
        mask = sanitize.init_land_mask(ug)

    if not vel_on_uv:
        ug = operators.interpolation(ug, mask, axis=1, padding="right")
        vg = operators.interpolation(vg, mask, axis=0, padding="right")
        ucg = operators.interpolation(ucg, mask, axis=1, padding="right")
        vcg = operators.interpolation(vcg, mask, axis=0, padding="right")

    if lat_u is None or lon_u is None or lat_v is None or lon_v is None:
        if lat_t is None or lon_t is None:
            raise ValueError("Either lat_t and lon_t, or lat_u, lon_u, lat_v, and lon_v must be provided")
        lat_u, lon_u, lat_v, lon_v = geometry.compute_uv_grids(lat_t, lon_t)

    dx_u, dy_u = geometry.compute_spatial_step(lat_u, lon_u)
    dx_v, dy_v = geometry.compute_spatial_step(lat_v, lon_v)
    coriolis_factor_u = geometry.compute_coriolis_factor(lat_u)
    coriolis_factor_v = geometry.compute_coriolis_factor(lat_v)

    return _cyclogeostrophic_loss(
        ug, vg, ucg, vcg,
        dx_u, dx_v, dy_u, dy_v,
        coriolis_factor_u, coriolis_factor_v,
        mask
    )


def cyclogeostrophic_imbalance(
    ug: Float[jax.Array, "lat lon"],
    vg: Float[jax.Array, "lat lon"],
    ucg: Float[jax.Array, "lat lon"],
    vcg: Float[jax.Array, "lat lon"],
    lat_t: Float[jax.Array, "lat lon"] = None,
    lon_t: Float[jax.Array, "lat lon"] = None,
    lat_u: Float[jax.Array, "lat lon"] = None,
    lon_u: Float[jax.Array, "lat lon"] = None,
    lat_v: Float[jax.Array, "lat lon"] = None,
    lon_v: Float[jax.Array, "lat lon"] = None,
    mask: Float[jax.Array, "lat lon"] = None,
    vel_on_uv: bool = True
) -> tuple[Float[jax.Array, "lat lon"], Float[jax.Array, "lat lon"]]:
    """
    Computes the cyclogeostrophic imbalance of a 2d velocity field, on a C-grid (following NEMO convention).

    The velocity fields can be provided either on the U and V grids (``vel_on_uv=True``) or on the T grid (``vel_on_uv=False``).

    Parameters
    ----------
    ug : Float[jax.Array, "lat lon"]
        U component of the geostrophic velocity field
    vg : Float[jax.Array, "lat lon"]
        V component of the geostrophic velocity field
    ucg : Float[jax.Array, "lat lon"]
        U component of the cyclogeostrophic velocity field
    vcg : Float[jax.Array, "lat lon"]
        V component of the cyclogeostrophic velocity field
    lat_t : Float[jax.Array, "lat lon"], optional
        Latitudes of the T grid.
        If ``lat_u``, ``lon_u``, ``lat_v``, and ``lon_v`` are not provided, ``lat_t`` and ``lon_t`` must be provided to compute them.
        Defaults to `None`
    lon_t : Float[jax.Array, "lat lon"], optional
        Longitudes of the T grid.
        If ``lat_u``, ``lon_u``, ``lat_v``, and ``lon_v`` are not provided, ``lat_t`` and ``lon_t`` must be provided to compute them.
        Defaults to `None`
    lat_u : Float[jax.Array, "lat lon"], optional
        Latitudes of the U grid.
        Defaults to `None`
    lon_u : Float[jax.Array, "lat lon"], optional
        Longitudes of the U grid.
        Defaults to `None`
    lat_v : Float[jax.Array, "lat lon"], optional
        Latitudes of the V grid.
        Defaults to `None`
    lon_v : Float[jax.Array, "lat lon"], optional
        Longitudes of the V grid.
        Defaults to `None`
    mask : Float[jax.Array, "lat lon"], optional
        Mask defining the marine area of the spatial domain; `1` or `True` stands for masked (i.e. land).
        Defaults to `None`
    vel_on_uv : bool, optional
        If `True`, the velocity components are assumed to be located on the U and V grids.
        If `False`, they are on the T grid.
        Defaults to `True`

    Returns
    -------
    u_imbalance_u : Float[jax.Array, "lat lon"]
        U component of the cyclogeostrophic imbalance, on the U grid
    v_imbalance_v : Float[jax.Array, "lat lon"]
        V component of the cyclogeostrophic imbalance, on the V grid
    """
    if not vel_on_uv:
        ug_u = operators.interpolation(ug, mask, axis=1, padding="right")
        vg_v = operators.interpolation(vg, mask, axis=0, padding="right")
        ucg_u = operators.interpolation(ucg, mask, axis=1, padding="right")
        vcg_v = operators.interpolation(vcg, mask, axis=0, padding="right")
    else:
        ug_u = ug
        vg_v = vg
        ucg_u = ucg
        vcg_v = vcg

    if lat_u is None or lon_u is None or lat_v is None or lon_v is None:
        if lat_t is None or lon_t is None:
            raise ValueError("Either lat_t and lon_t, or lat_u, lon_u, lat_v, and lon_v must be provided")
        lat_u, lon_u, lat_v, lon_v = geometry.compute_uv_grids(lat_t, lon_t)

    dx_u, dy_u = geometry.compute_spatial_step(lat_u, lon_u)
    dx_v, dy_v = geometry.compute_spatial_step(lat_v, lon_v)
    coriolis_factor_u = geometry.compute_coriolis_factor(lat_u)
    coriolis_factor_v = geometry.compute_coriolis_factor(lat_v)

    return _cyclogeostrophic_imbalance(
        ug_u, vg_v, ucg_u, vcg_v,
        dx_u, dx_v, dy_u, dy_v, coriolis_factor_u, coriolis_factor_v, mask
    )


def radius_of_curvature(
    u: Float[jax.Array, "lat lon"],
    v: Float[jax.Array, "lat lon"],
    lat_t: Float[jax.Array, "lat lon"] = None,
    lon_t: Float[jax.Array, "lat lon"] = None,
    lat_u: Float[jax.Array, "lat lon"] = None,
    lon_u: Float[jax.Array, "lat lon"] = None,
    lat_v: Float[jax.Array, "lat lon"] = None,
    lon_v: Float[jax.Array, "lat lon"] = None,
    mask: Float[jax.Array, "lat lon"] = None,
    vel_on_uv: bool = True
) -> Float[jax.Array, "lat lon"]:
    """
    Computes the radius of curvature of a 2d velocity field, on a C-grid (following NEMO convention).

    The velocity field can be provided either on the U and V grids (``vel_on_uv=True``) or on the T grid (``vel_on_uv=False``).

    Parameters
    ----------
    u : Float[jax.Array, "lat lon"]
        U component of the velocity field
    v : Float[jax.Array, "lat lon"]
        V component of the velocity field
    lat_t : Float[jax.Array, "lat lon"], optional
        Latitudes of the T grid.
        If ``lat_u``, ``lon_u``, ``lat_v``, and ``lon_v`` are not provided, ``lat_t`` and ``lon_t`` must be provided to compute them.
        Defaults to `None`
    lon_t : Float[jax.Array, "lat lon"], optional
        Longitudes of the T grid.
        If ``lat_u``, ``lon_u``, ``lat_v``, and ``lon_v`` are not provided, ``lat_t`` and ``lon_t`` must be provided to compute them.
        Defaults to `None`
    lat_u : Float[jax.Array, "lat lon"], optional
        Latitudes of the U grid.
        Defaults to `None`
    lon_u : Float[jax.Array, "lat lon"], optional
        Longitudes of the U grid.
        Defaults to `None`
    lat_v : Float[jax.Array, "lat lon"], optional
        Latitudes of the V grid.
        Defaults to `None`
    lon_v : Float[jax.Array, "lat lon"], optional
        Longitudes of the V grid.
        Defaults to `None`
    mask : Float[jax.Array, "lat lon"], optional
        Mask defining the marine area of the spatial domain; `1` or `True` stands for masked (i.e. land).
        Defaults to `None`
    vel_on_uv : bool, optional
        If `True`, ``u`` and ``v`` are on the U and V grids.
        If `False`, they are on the T grid.
        Defaults to `True`

    Returns
    -------
    r : Float[jax.Array, "lat lon"]
        The radius of curvature of the velocity field
    """
    if lat_u is None or lon_u is None or lat_v is None or lon_v is None:
        if lat_t is None or lon_t is None:
            raise ValueError("Either lat_t and lon_t, or lat_u, lon_u, lat_v, and lon_v must be provided")
        lat_u, lon_u, lat_v, lon_v = geometry.compute_uv_grids(lat_t, lon_t)

    dx_u, dy_u = geometry.compute_spatial_step(lat_u, lon_u)
    dx_v, dy_v = geometry.compute_spatial_step(lat_v, lon_v)

    return _radius_of_curvature(u, v, dx_u, dx_v, dy_u, dy_v, mask, vel_on_uv)


# =============================================================================
# Internal Functions
# =============================================================================

def _cyclogeostrophic_loss(
    ug_u: Float[jax.Array, "lat lon"],
    vg_v: Float[jax.Array, "lat lon"],
    ucg_u: Float[jax.Array, "lat lon"],
    vcg_v: Float[jax.Array, "lat lon"],
    dx_u: Float[jax.Array, "lat lon"],
    dx_v: Float[jax.Array, "lat lon"],
    dy_u: Float[jax.Array, "lat lon"],
    dy_v: Float[jax.Array, "lat lon"],
    coriolis_factor_u: Float[jax.Array, "lat lon"],
    coriolis_factor_v: Float[jax.Array, "lat lon"],
    mask: Float[jax.Array, "lat lon"]
) -> Float[jax.Array, ""]:
    u_imbalance, v_imbalance = _cyclogeostrophic_imbalance(
        ug_u, vg_v, ucg_u, vcg_v,
        dx_u, dx_v, dy_u, dy_v, coriolis_factor_u, coriolis_factor_v,
        mask
    )

    return jnp.nansum(u_imbalance ** 2 + v_imbalance ** 2)


def _cyclogeostrophic_imbalance(
    ug_u: Float[jax.Array, "lat lon"],
    vg_v: Float[jax.Array, "lat lon"],
    ucg_u: Float[jax.Array, "lat lon"],
    vcg_v: Float[jax.Array, "lat lon"],
    dx_u: Float[jax.Array, "lat lon"],
    dx_v: Float[jax.Array, "lat lon"],
    dy_u: Float[jax.Array, "lat lon"],
    dy_v: Float[jax.Array, "lat lon"],
    coriolis_factor_u: Float[jax.Array, "lat lon"],
    coriolis_factor_v: Float[jax.Array, "lat lon"],
    mask: Float[jax.Array, "lat lon"]
) -> tuple[Float[jax.Array, "lat lon"], Float[jax.Array, "lat lon"]]:
    u_adv_v, v_adv_u = _advection(ucg_u, vcg_v, dx_u, dx_v, dy_u, dy_v, mask)

    u_imbalance_u = ucg_u + v_adv_u / coriolis_factor_u - ug_u
    v_imbalance_v = vcg_v - u_adv_v / coriolis_factor_v - vg_v

    return u_imbalance_u, v_imbalance_v


def _advection(
    u_u: Float[jax.Array, "lat lon"],
    v_v: Float[jax.Array, "lat lon"],
    dx_u: Float[jax.Array, "lat lon"],
    dx_v: Float[jax.Array, "lat lon"],
    dy_u: Float[jax.Array, "lat lon"],
    dy_v: Float[jax.Array, "lat lon"],
    mask: Float[jax.Array, "lat lon"]
) -> tuple[Float[jax.Array, "lat lon"], Float[jax.Array, "lat lon"]]:
    """
    Computes the advection terms of a 2d velocity field, on a C-grid, following NEMO convention.

    Parameters
    ----------
    u_u : Float[jax.Array, "lat lon"]
        U component of the velocity field (on the U grid)
    v_v : Float[jax.Array, "lat lon"]
        V component of the SSC velocity field (on the V grid)
    dx_u : Float[jax.Array, "lat lon"]
        Spatial steps on the U grid along `x`, in meters
    dy_u : Float[jax.Array, "lat lon"]
        Spatial steps on the U grid along `y`, in meters
    dx_v : Float[jax.Array, "lat lon"]
        Spatial steps on the V grid along `x`, in meters
    dy_v : Float[jax.Array, "lat lon"]
        Spatial steps on the V grid along `y`, in meters
    mask : Float[jax.Array, "lat lon"]
        Mask defining the marine area of the spatial domain; `1` or `True` stands for masked (i.e. land)

    Returns
    -------
    u_adv_v : Float[jax.Array, "lat lon"]
        U component of the advection term, on the V grid
    v_adv_u : Float[jax.Array, "lat lon"]
        V component of the advection term, on the U grid
    """
    u_adv_v = _u_advection_v(u_u, v_v, dx_v, dy_v, mask)
    v_adv_u = _v_advection_u(u_u, v_v, dx_u, dy_u, mask)

    return u_adv_v, v_adv_u


def _u_advection_v(
    u_u: Float[jax.Array, "lat lon"],
    v_v: Float[jax.Array, "lat lon"],
    dx_u: Float[jax.Array, "lat lon"],
    dy_u: Float[jax.Array, "lat lon"],
    mask: Float[jax.Array, "lat lon"]
) -> Float[jax.Array, "lat lon"]:
    dudx_t = operators.derivative(u_u, dx_u, mask, axis=1, padding="left")   # (U(i), U(i+1)) -> T(i+1)
    dudx_v = operators.interpolation(dudx_t, mask, axis=0, padding="right")  # (T(j), T(j+1)) -> V(j)

    dudy_f = operators.derivative(u_u, dy_u, mask, axis=0, padding="right")  # (U(j), U(j+1)) -> F(j)
    dudy_v = operators.interpolation(dudy_f, mask, axis=1, padding="left")   # (F(i), F(i+1)) -> V(i+1)

    u_t = operators.interpolation(u_u, mask, axis=1, padding="left")   # (U(i), U(i+1)) -> T(i+1)
    u_v = operators.interpolation(u_t, mask, axis=0, padding="right")  # (T(j), T(j+1)) -> V(j)

    u_adv_v = u_v * dudx_v + v_v * dudy_v  # V(j)

    return u_adv_v


def _v_advection_u(
    u_u: Float[jax.Array, "lat lon"],
    v_v: Float[jax.Array, "lat lon"],
    dx_v: Float[jax.Array, "lat lon"],
    dy_v: Float[jax.Array, "lat lon"],
    mask: Float[jax.Array, "lat lon"]
) -> Float[jax.Array, "lat lon"]:
    dvdx_f = operators.derivative(v_v, dx_v, mask, axis=1, padding="right")  # (V(i), V(i+1)) -> F(i)
    dvdx_u = operators.interpolation(dvdx_f, mask, axis=0, padding="left")   # (F(j), F(j+1)) -> U(j+1)

    dvdy_t = operators.derivative(v_v, dy_v, mask, axis=0, padding="left")   # (V(j), V(j+1)) -> T(j+1)
    dvdy_u = operators.interpolation(dvdy_t, mask, axis=1, padding="right")  # (T(i), T(i+1)) -> U(i)

    v_t = operators.interpolation(v_v, mask, axis=0, padding="left")   # (V(j), V(j+1)) -> T(j+1)
    v_u = operators.interpolation(v_t, mask, axis=1, padding="right")  # (T(i), T(i+1)) -> U(i)

    v_adv_u = u_u * dvdx_u + v_u * dvdy_u  # U(i)

    return v_adv_u


def _radius_of_curvature(
    u: Float[jax.Array, "lat lon"],
    v: Float[jax.Array, "lat lon"],
    dx_u: Float[jax.Array, "lat lon"],
    dx_v: Float[jax.Array, "lat lon"],
    dy_u: Float[jax.Array, "lat lon"],
    dy_v: Float[jax.Array, "lat lon"],
    mask: Float[jax.Array, "lat lon"],
    vel_on_uv: bool
) -> Float[jax.Array, "lat lon"]:
    if not vel_on_uv:
        u_t = u
        v_t = v
        u_u = operators.interpolation(u, mask, axis=1, padding="right")
        v_v = operators.interpolation(v, mask, axis=0, padding="right")
    else:
        u_t = operators.interpolation(u, mask, axis=1, padding="left")
        v_t = operators.interpolation(v, mask, axis=0, padding="left")
        u_u = u
        v_v = v

    V_t = kinematics.magnitude(u_t, v_t, vel_on_uv=False)

    du_dx_t = operators.derivative(u_u, dx_u, mask, axis=1, padding="left")  # (U(i), U(i+1)) -> T(i+1)
    du_dy_f = operators.derivative(u_u, dy_u, mask, axis=0, padding="right")  # (U(j), U(j+1)) -> F(j)

    dv_dx_f = operators.derivative(v_v, dx_v, mask, axis=1, padding="right")  # (V(i), V(i+1)) -> F(i)
    dv_dy_t = operators.derivative(v_v, dy_v, mask, axis=0, padding="left")  # (V(j), V(j+1)) -> T(j+1)

    du_dy_v = operators.interpolation(du_dy_f, mask, axis=1, padding="left")  # (F(i), F(i+1)) -> V(i+1)
    du_dy_t = operators.interpolation(du_dy_v, mask, axis=0, padding="left")  # (V(j), V(j+1)) -> T(j+1)
    dv_dx_u = operators.interpolation(dv_dx_f, mask, axis=0, padding="left")  # (F(j), F(j+1)) -> U(j+1)
    dv_dx_t = operators.interpolation(dv_dx_u, mask, axis=1, padding="left")  # (U(i), U(i+1)) -> T(i+1)

    numerator = V_t ** 3
    denominator = u_t ** 2 * dv_dx_t - v_t ** 2 * du_dy_t - u_t * v_t * (du_dx_t - dv_dy_t)
    r = numerator / denominator

    return r
