Utilities¶

def create_global_grid(
    res_lat: float,
    res_lon: float,
    add_bounds: bool = True,
    chunks: Optional[Union[int, Dict[str, int]]] = None,
) -> xr.Dataset:
    """
    Create a global rectilinear grid dataset.

    Parameters
    ----------
    res_lat : float
        Latitude resolution in degrees.
    res_lon : float
        Longitude resolution in degrees.
    add_bounds : bool, default True
        Whether to add cell boundary coordinates.
    chunks : int or dict, optional
        Chunk sizes for the resulting dask-backed dataset.
        If None (default), returns an eager NumPy-backed dataset.

    Returns
    -------
    xr.Dataset
        The global grid dataset containing 'lat' and 'lon'.
    """
    return _create_rectilinear_grid(
        lat_range=(-90, 90),
        lon_range=(0, 360),
        res_lat=res_lat,
        res_lon=res_lon,
        add_bounds=add_bounds,
        chunks=chunks,
        history_msg=f"Created global grid ({res_lat}x{res_lon}) using xregrid.",
    )

def create_regional_grid(
    lat_range: Tuple[float, float],
    lon_range: Tuple[float, float],
    res_lat: float,
    res_lon: float,
    add_bounds: bool = True,
    chunks: Optional[Union[int, Dict[str, int]]] = None,
) -> xr.Dataset:
    """
    Create a regional rectilinear grid dataset.

    Parameters
    ----------
    lat_range : tuple of float
        (min_lat, max_lat).
    lon_range : tuple of float
        (min_lon, max_lon).
    res_lat : float
        Latitude resolution in degrees.
    res_lon : float
        Longitude resolution in degrees.
    add_bounds : bool, default True
        Whether to add cell boundary coordinates.
    chunks : int or dict, optional
        Chunk sizes for the resulting dask-backed dataset.
        If None (default), returns an eager NumPy-backed dataset.

    Returns
    -------
    xr.Dataset
        The regional grid dataset containing 'lat' and 'lon'.
    """
    return _create_rectilinear_grid(
        lat_range=lat_range,
        lon_range=lon_range,
        res_lat=res_lat,
        res_lon=res_lon,
        add_bounds=add_bounds,
        chunks=chunks,
        history_msg=f"Created regional grid ({res_lat}x{res_lon}) using xregrid.",
    )

def create_grid_like(
    obj: Union[xr.DataArray, xr.Dataset],
    res: Union[float, Tuple[float, float]],
    add_bounds: bool = True,
    chunks: Optional[Union[int, Dict[str, int]]] = None,
    extent: Optional[Tuple[float, float, float, float]] = None,
    crs: Optional[Union[str, int, Any]] = None,
) -> xr.Dataset:
    """
    Create a new grid dataset with the same extent and CRS as an existing object.

    Automatically detects the CRS and spatial extent of the input object.
    Supports both geographic (lat-lon) and projected coordinate systems.

    Parameters
    ----------
    obj : xr.DataArray or xr.Dataset
        The input object to use as a template.
    res : float or tuple of float
        New grid resolution in the coordinate system units.
        If tuple, (res_x, res_y) or (res_lon, res_lat).
    add_bounds : bool, default True
        Whether to add cell boundary coordinates.
    chunks : int or dict, optional
        Chunk sizes for the resulting dask-backed dataset.
    extent : tuple of float, optional
        Override the detected extent (min_x, max_x, min_y, max_y).
        Use this to avoid hidden dask.compute() if you already know the extent.
    crs : str, int, or pyproj.CRS, optional
        Override the detected CRS.

    Returns
    -------
    xr.Dataset
        The new grid dataset.
    """
    if crs is not None:
        if pyproj is not None:
            crs_obj = pyproj.CRS(crs)
        else:
            crs_obj = crs
    else:
        crs_obj = get_crs_info(obj)

    if isinstance(res, (int, float)):
        res_x = res_y = float(res)
    else:
        res_x, res_y = map(float, res)

    obj_name = getattr(obj, "name", "input")
    history_msg_base = f"Created grid like {obj_name} using xregrid."
    if hasattr(obj, "attrs") and "history" in obj.attrs:
        history_msg_base += f"\nTemplate history:\n{obj.attrs['history']}"

    if extent is not None:
        if crs_obj is None or (
            hasattr(crs_obj, "is_geographic") and crs_obj.is_geographic
        ):
            # Lat-Lon
            return _create_rectilinear_grid(
                (extent[2], extent[3]),  # lat_range
                (extent[0], extent[1]),  # lon_range
                res_y,  # res_lat
                res_x,  # res_lon
                add_bounds=add_bounds,
                chunks=chunks,
                crs=crs_obj.to_wkt() if hasattr(crs_obj, "to_wkt") else "EPSG:4326",
                history_msg=history_msg_base + " (Override Extent).",
            )
        else:
            # Projected
            return create_grid_from_crs(
                crs_obj, extent, (res_x, res_y), add_bounds=add_bounds, chunks=chunks
            )

    # 1. Try to find projected coordinates
    try:
        x_da = obj.cf["projection_x_coordinate"]
        y_da = obj.cf["projection_y_coordinate"]

        try:
            # Use bounds for exact extent if available
            x_b = obj.cf.get_bounds("projection_x_coordinate")
            y_b = obj.cf.get_bounds("projection_y_coordinate")

            # Batch compute if lazy to minimize roundtrips
            if dask is not None and (
                hasattr(x_b.data, "dask") or hasattr(y_b.data, "dask")
            ):
                vals = dask.compute(x_b.min(), x_b.max(), y_b.min(), y_b.max())
                extent = tuple(map(float, vals))
            elif hasattr(x_b.data, "dask") or hasattr(y_b.data, "dask"):
                extent = (
                    float(x_b.min()),
                    float(x_b.max()),
                    float(y_b.min()),
                    float(y_b.max()),
                )
            else:
                extent = (
                    float(x_b.min()),
                    float(x_b.max()),
                    float(y_b.min()),
                    float(y_b.max()),
                )
        except Exception:
            # Fallback to centers
            # Discovery logic: we need min/max and average diff for heuristic
            if dask is not None and (
                hasattr(x_da.data, "dask") or hasattr(y_da.data, "dask")
            ):
                # Batch everything!
                tasks_dict = {
                    "x_min": x_da.min(),
                    "x_max": x_da.max(),
                    "y_min": y_da.min(),
                    "y_max": y_da.max(),
                }
                if x_da.size > 1:
                    tasks_dict["res_x"] = abs(x_da.diff(x_da.dims[0]).mean())
                if y_da.size > 1:
                    tasks_dict["res_y"] = abs(y_da.diff(y_da.dims[0]).mean())

                results = dask.compute(tasks_dict)[0]
                x_min, x_max, y_min, y_max = (
                    float(results["x_min"]),
                    float(results["x_max"]),
                    float(results["y_min"]),
                    float(results["y_max"]),
                )

                res_x_orig = float(results.get("res_x", 0))
                res_y_orig = float(
                    results.get("res_y", res_x_orig if res_x_orig else 0)
                )

                extent = (
                    x_min - res_x_orig / 2,
                    x_max + res_x_orig / 2,
                    y_min - res_y_orig / 2,
                    y_max + res_y_orig / 2,
                )
            elif hasattr(x_da.data, "dask") or hasattr(y_da.data, "dask"):
                # Non-batched fallback
                res_x_orig = (
                    abs(float(x_da.diff(x_da.dims[0]).mean())) if x_da.size > 1 else 0
                )
                res_y_orig = (
                    abs(float(y_da.diff(y_da.dims[0]).mean()))
                    if y_da.size > 1
                    else res_x_orig
                )
                extent = (
                    float(x_da.min()) - res_x_orig / 2,
                    float(x_da.max()) + res_x_orig / 2,
                    float(y_da.min()) - res_y_orig / 2,
                    float(y_da.max()) + res_y_orig / 2,
                )
            else:
                res_x_orig = (
                    abs(float(x_da.diff(x_da.dims[0]).mean())) if x_da.size > 1 else 0
                )
                res_y_orig = (
                    abs(float(y_da.diff(y_da.dims[0]).mean()))
                    if y_da.size > 1
                    else res_x_orig
                )
                extent = (
                    float(x_da.min()) - res_x_orig / 2,
                    float(x_da.max()) + res_x_orig / 2,
                    float(y_da.min()) - res_y_orig / 2,
                    float(y_da.max()) + res_y_orig / 2,
                )

        if crs_obj is None:
            # Fallback to generic geographic if no CRS found
            crs_obj = "EPSG:4326"

        return create_grid_from_crs(
            crs_obj, extent, (res_x, res_y), add_bounds=add_bounds, chunks=chunks
        )

    except (KeyError, AttributeError, ValueError):
        pass

    # 2. Fallback to Geographic (Lat-Lon)
    try:
        lat_da = _find_coord(obj, "latitude")
        lon_da = _find_coord(obj, "longitude")
        if lat_da is None or lon_da is None:
            raise KeyError("Coordinates not found")

        try:
            lat_b = obj.cf.get_bounds("latitude")
            lon_b = obj.cf.get_bounds("longitude")

            if dask is not None and (
                hasattr(lat_b.data, "dask") or hasattr(lon_b.data, "dask")
            ):
                vals = dask.compute(lat_b.min(), lat_b.max(), lon_b.min(), lon_b.max())
                lat_range = (float(vals[0]), float(vals[1]))
                lon_range = (float(vals[2]), float(vals[3]))
            elif hasattr(lat_b.data, "dask") or hasattr(lon_b.data, "dask"):
                lat_range = (float(lat_b.min()), float(lat_b.max()))
                lon_range = (float(lon_b.min()), float(lon_b.max()))
            else:
                lat_range = (float(lat_b.min()), float(lat_b.max()))
                lon_range = (float(lon_b.min()), float(lon_b.max()))
        except Exception:
            # Heuristic for resolution to calculate extent from centers
            if dask is not None and (
                hasattr(lat_da.data, "dask") or hasattr(lon_da.data, "dask")
            ):
                tasks_dict = {
                    "lat_min": lat_da.min(),
                    "lat_max": lat_da.max(),
                    "lon_min": lon_da.min(),
                    "lon_max": lon_da.max(),
                }
                if lat_da.size > 1:
                    tasks_dict["res_lat"] = abs(lat_da.diff(lat_da.dims[0]).mean())
                if lon_da.size > 1:
                    tasks_dict["res_lon"] = abs(lon_da.diff(lon_da.dims[-1]).mean())

                results = dask.compute(tasks_dict)[0]
                lat_min, lat_max, lon_min, lon_max = (
                    float(results["lat_min"]),
                    float(results["lat_max"]),
                    float(results["lon_min"]),
                    float(results["lon_max"]),
                )

                res_lat_orig = float(results.get("res_lat", 0))
                res_lon_orig = float(
                    results.get("res_lon", res_lat_orig if res_lat_orig else 0)
                )

                lat_range = (
                    lat_min - res_lat_orig / 2,
                    lat_max + res_lat_orig / 2,
                )
                lon_range = (
                    lon_min - res_lon_orig / 2,
                    lon_max + res_lon_orig / 2,
                )
            elif hasattr(lat_da.data, "dask") or hasattr(lon_da.data, "dask"):
                res_lat_orig = (
                    abs(float(lat_da.diff(lat_da.dims[0]).mean()))
                    if lat_da.size > 1
                    else 0
                )
                res_lon_orig = (
                    abs(float(lon_da.diff(lon_da.dims[-1]).mean()))
                    if lon_da.size > 1
                    else res_lat_orig
                )
                lat_range = (
                    float(lat_da.min()) - res_lat_orig / 2,
                    float(lat_da.max()) + res_lat_orig / 2,
                )
                lon_range = (
                    float(lon_da.min()) - res_lon_orig / 2,
                    float(lon_da.max()) + res_lon_orig / 2,
                )
            else:
                res_lat_orig = (
                    abs(float(lat_da.diff(lat_da.dims[0]).mean()))
                    if lat_da.size > 1
                    else 0
                )
                res_lon_orig = (
                    abs(float(lon_da.diff(lon_da.dims[-1]).mean()))
                    if lon_da.size > 1
                    else res_lat_orig
                )
                lat_range = (
                    float(lat_da.min()) - res_lat_orig / 2,
                    float(lat_da.max()) + res_lat_orig / 2,
                )
                lon_range = (
                    float(lon_da.min()) - res_lon_orig / 2,
                    float(lon_da.max()) + res_lon_orig / 2,
                )

        return _create_rectilinear_grid(
            lat_range,
            lon_range,
            res_y,  # res_lat
            res_x,  # res_lon
            add_bounds=add_bounds,
            chunks=chunks,
            crs=crs_obj.to_wkt() if crs_obj else "EPSG:4326",
            history_msg=history_msg_base,
        )
    except (KeyError, AttributeError, ValueError):
        raise ValueError(
            "Could not detect spatial coordinates (latitude/longitude or "
            "projection_x/y) in input object."
        )

def create_grid_from_crs(
    crs: Union[str, int, Any],
    extent: Tuple[float, float, float, float],
    res: Union[float, Tuple[float, float]],
    add_bounds: bool = True,
    chunks: Optional[Union[int, Dict[str, int]]] = None,
) -> xr.Dataset:
    """
    Create a structured grid dataset from a CRS and extent.

    Parameters
    ----------
    crs : str, int, or pyproj.CRS
        The CRS of the grid (Proj4 string, EPSG code, WKT, or CRS object).
    extent : Tuple[float, float, float, float]
        Grid extent in CRS units: (min_x, max_x, min_y, max_y).
    res : float or Tuple[float, float]
        Grid resolution in CRS units. If float, same resolution in x and y.
        If tuple, (res_x, res_y).
    add_bounds : bool, default True
        Whether to add cell boundary coordinates.
    chunks : int or Dict[str, int], optional
        Chunk sizes for the resulting dask-backed dataset.

    Returns
    -------
    xr.Dataset
        The grid dataset containing 'lat', 'lon' and projected coordinates 'x', 'y'.
    """
    if isinstance(res, (int, float)):
        res_x = res_y = float(res)
    else:
        res_x, res_y = map(float, res)

    x_chunks = chunks.get("x", -1) if isinstance(chunks, dict) else chunks
    y_chunks = chunks.get("y", -1) if isinstance(chunks, dict) else chunks

    # Generate 1D coordinates in projected space
    x = _lazy_arange(extent[0] + res_x / 2, extent[1], res_x, chunks=x_chunks)
    y = _lazy_arange(extent[2] + res_y / 2, extent[3], res_y, chunks=y_chunks)

    x_da = xr.DataArray(x, dims=["x"], name="x")
    y_da = xr.DataArray(y, dims=["y"], name="y")

    # Use xr.broadcast for lazy 2D arrays
    yy_da, xx_da = xr.broadcast(y_da, x_da)

    # Ensure (y, x) order
    yy_da = yy_da.transpose("y", "x")
    xx_da = xx_da.transpose("y", "x")

    if pyproj is None:
        raise ImportError("pyproj is required for create_grid_from_crs.")
    crs_obj = pyproj.CRS(crs)

    lon, lat = xr.apply_ufunc(
        _transform_coords,
        xx_da,
        yy_da,
        kwargs={"crs_in": crs_obj},
        dask="parallelized",
        output_dtypes=[float, float],
        input_core_dims=[[], []],
        output_core_dims=[[], []],
    )

    try:
        units = crs_obj.axis_info[0].unit_name or "m"
    except (IndexError, AttributeError):
        units = "m"

    ds = xr.Dataset(
        coords={
            "y": (
                ["y"],
                y,
                {"units": units, "standard_name": "projection_y_coordinate"},
            ),
            "x": (
                ["x"],
                x,
                {"units": units, "standard_name": "projection_x_coordinate"},
            ),
            "lat": (
                ["y", "x"],
                lat.data,
                {"units": "degrees_north", "standard_name": "latitude"},
            ),
            "lon": (
                ["y", "x"],
                lon.data,
                {"units": "degrees_east", "standard_name": "longitude"},
            ),
        }
    )

    ds.attrs["crs"] = crs_obj.to_wkt()

    if add_bounds:
        if chunks is not None and da is not None:
            x_b_raw = da.stack(
                [x - res_x / 2, x + res_x / 2, x + res_x / 2, x - res_x / 2]
            )
            y_b_raw = da.stack(
                [y - res_y / 2, y - res_y / 2, y + res_y / 2, y + res_y / 2]
            )
        else:
            x_b_raw = np.stack(
                [x - res_x / 2, x + res_x / 2, x + res_x / 2, x - res_x / 2]
            )
            y_b_raw = np.stack(
                [y - res_y / 2, y - res_y / 2, y + res_y / 2, y + res_y / 2]
            )

        x_b_da = xr.DataArray(x_b_raw, dims=["nv", "x"])
        y_b_da = xr.DataArray(y_b_raw, dims=["nv", "y"])

        yy_b_da, xx_b_da = xr.broadcast(y_b_da, x_b_da)

        lon_b, lat_b = xr.apply_ufunc(
            _transform_coords,
            xx_b_da,
            yy_b_da,
            kwargs={"crs_in": crs_obj},
            dask="parallelized",
            output_dtypes=[float, float],
            input_core_dims=[[], []],
            output_core_dims=[[], []],
        )

        ds.coords["lat_b"] = (
            ["y", "x", "nv"],
            lat_b.data.transpose(1, 2, 0),
            {"units": "degrees_north"},
        )
        ds.coords["lon_b"] = (
            ["y", "x", "nv"],
            lon_b.data.transpose(1, 2, 0),
            {"units": "degrees_east"},
        )
        ds["lat"].attrs["bounds"] = "lat_b"
        ds["lon"].attrs["bounds"] = "lon_b"

        # Add 1D projected bounds using backend-agnostic xarray operations
        x_da_1d = xr.DataArray(x, dims=["x"])
        y_da_1d = xr.DataArray(y, dims=["y"])

        # Create (N, 2) bounds
        x_b_1d = xr.concat(
            [x_da_1d - res_x / 2, x_da_1d + res_x / 2], dim="nbounds"
        ).transpose("x", "nbounds")
        y_b_1d = xr.concat(
            [y_da_1d - res_y / 2, y_da_1d + res_y / 2], dim="nbounds"
        ).transpose("y", "nbounds")

        ds.coords["x_b"] = (["x", "nbounds"], x_b_1d.data, {"units": units})
        ds.coords["y_b"] = (["y", "nbounds"], y_b_1d.data, {"units": units})
        ds["x"].attrs["bounds"] = "x_b"
        ds["y"].attrs["bounds"] = "y_b"

    # Backend detection for provenance
    is_lazy = chunks is not None
    backend = "Lazy" if is_lazy else "Eager"

    update_history(ds, f"Created grid from CRS {crs} using xregrid ({backend}).")
    if chunks is not None:
        ds = ds.chunk(chunks)
    return ds

def create_grid_from_ioapi(
    metadata: Dict[str, Any],
    add_bounds: bool = True,
    chunks: Optional[Union[int, Dict[str, int]]] = None,
) -> xr.Dataset:
    """
    Create a structured grid dataset from IOAPI-compliant metadata.

    Supports GDTYP:
    - 1: Lat-Lon
    - 2: Lambert Conformal
    - 3: Mercator
    - 4: Stereographic
    - 5: UTM
    - 6: Polar Stereographic
    - 7: Equatorial Mercator
    - 8: Transverse Mercator
    - 9: Albers Equal Area
    - 10: Lambert Azimuthal Equal Area
    - 13: Sinusoidal

    Parameters
    ----------
    metadata : dict
        IOAPI metadata containing GDTYP, P_ALP, P_BET, P_GAM, XCENT, YCENT,
        XORIG, YORIG, XCELL, YCELL, NCOLS, NROWS.
    add_bounds : bool, default True
        Whether to add cell boundary coordinates.
    chunks : int or dict, optional
        Chunk sizes for the resulting dask-backed dataset.

    Returns
    -------
    xr.Dataset
        The grid dataset.
    """
    gdtyp = metadata["GDTYP"]
    p_alp = metadata["P_ALP"]
    p_bet = metadata["P_BET"]
    xcent = metadata["XCENT"]
    ycent = metadata["YCENT"]
    xorig = metadata["XORIG"]
    yorig = metadata["YORIG"]
    xcell = metadata["XCELL"]
    ycell = metadata["YCELL"]
    ncols = metadata["NCOLS"]
    nrows = metadata["NROWS"]

    if gdtyp == 1:  # Lat-Lon
        crs = "EPSG:4326"
    elif gdtyp == 2:  # Lambert Conformal
        crs = (
            f"+proj=lcc +lat_1={p_alp} +lat_2={p_bet} +lat_0={ycent} "
            f"+lon_0={xcent} +x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs"
        )
    elif gdtyp == 3:  # Mercator
        crs = (
            f"+proj=merc +lat_ts={p_alp} +lon_0={xcent} +lat_0={ycent} "
            f"+x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs"
        )
    elif gdtyp == 4:  # Stereographic
        crs = (
            f"+proj=stere +lat_ts={p_alp} +lat_0={ycent} +lon_0={xcent} "
            f"+x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs"
        )
    elif gdtyp == 5:  # UTM
        crs = f"+proj=utm +zone={int(p_alp)} +datum=WGS84 +units=m +no_defs"
    elif gdtyp == 6:  # Polar Stereographic
        # lat_0 determined by p_alp (1.0 for North, -1.0 for South)
        lat_0 = 90.0 if p_alp > 0 else -90.0
        crs = (
            f"+proj=stere +lat_0={lat_0} +lat_ts={p_bet} +lon_0={xcent} "
            f"+x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs"
        )
    elif gdtyp == 7:  # Equatorial Mercator
        crs = (
            f"+proj=merc +lat_ts={p_alp} +lon_0={xcent} +lat_0=0 "
            f"+x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs"
        )
    elif gdtyp == 8:  # Transverse Mercator
        crs = (
            f"+proj=tmerc +lat_0={ycent} +k={p_bet} +lon_0={xcent} "
            f"+x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs"
        )
    elif gdtyp == 9:  # Albers Equal Area
        crs = (
            f"+proj=aea +lat_1={p_alp} +lat_2={p_bet} +lat_0={ycent} "
            f"+lon_0={xcent} +x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs"
        )
    elif gdtyp == 10:  # Lambert Azimuthal Equal Area
        crs = (
            f"+proj=laea +lat_0={ycent} +lon_0={xcent} "
            f"+x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs"
        )
    elif gdtyp == 13:  # Sinusoidal
        crs = f"+proj=sinu +lon_0={xcent} +x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs"
    else:
        raise ValueError(f"Unsupported IOAPI GDTYP: {gdtyp}")

    extent = (xorig, xorig + ncols * xcell, yorig, yorig + nrows * ycell)
    res = (xcell, ycell)

    ds = create_grid_from_crs(crs, extent, res, add_bounds=add_bounds, chunks=chunks)

    # Attach IOAPI metadata for provenance
    for k, v in metadata.items():
        ds.attrs[f"ioapi_{k}"] = v

    update_history(ds, f"Created grid from IOAPI metadata (GDTYP={gdtyp})")

    return ds

def create_mesh_from_coords(
    x: Union[np.ndarray, xr.DataArray],
    y: Union[np.ndarray, xr.DataArray],
    crs: Union[str, int, Any],
    chunks: Optional[Union[int, Dict[str, int]]] = None,
) -> xr.Dataset:
    """
    Create an unstructured mesh dataset from coordinates and a CRS.

    Parameters
    ----------
    x : np.ndarray or xr.DataArray
        1D array of x coordinates in CRS units.
    y : np.ndarray or xr.DataArray
        1D array of y coordinates in CRS units.
    crs : str, int, or pyproj.CRS
        The CRS of the coordinates.
    chunks : int or dict, optional
        Chunk sizes for the resulting dask-backed dataset.
        If None (default), returns an eager NumPy-backed dataset.

    Returns
    -------
    xr.Dataset
        The mesh dataset containing 'lat', 'lon' and 'x', 'y' as 1D arrays.
    """
    if pyproj is None:
        raise ImportError(
            "pyproj is required for create_mesh_from_coords. "
            "Install it with `pip install pyproj`."
        )
    crs_obj = pyproj.CRS(crs)

    # Force n_pts dimension to avoid alignment/broadcasting issues in apply_ufunc
    # Preserve backend (data) while dropping stale coordinates
    if isinstance(x, xr.DataArray):
        x_da = xr.DataArray(x.data, dims=["n_pts"], name="x", attrs=x.attrs)
    else:
        x_da = xr.DataArray(x, dims=["n_pts"], name="x")

    if isinstance(y, xr.DataArray):
        y_da = xr.DataArray(y.data, dims=["n_pts"], name="y", attrs=y.attrs)
    else:
        y_da = xr.DataArray(y, dims=["n_pts"], name="y")

    if chunks is not None:
        # If chunks is a dict, ensure we use the canonical dimension name
        if isinstance(chunks, dict):
            # Map any single-dimension chunking to n_pts
            n_pts_chunks = next(iter(chunks.values()))
            chunks = {"n_pts": n_pts_chunks}
        x_da = x_da.chunk(chunks)
        y_da = y_da.chunk(chunks)

    # Backend detection for provenance
    is_lazy = chunks is not None or hasattr(x_da.data, "dask")
    backend = "Lazy" if is_lazy else "Eager"

    # Use apply_ufunc with dask='parallelized'
    lon, lat = xr.apply_ufunc(
        _transform_coords,
        x_da,
        y_da,
        kwargs={"crs_in": crs_obj},
        dask="parallelized",
        output_dtypes=[float, float],
        input_core_dims=[[], []],
        output_core_dims=[[], []],
    )

    # Conditional metadata based on CRS type
    if crs_obj.is_geographic:
        x_std, y_std = "longitude", "latitude"
        x_units, y_units = "degrees_east", "degrees_north"
    else:
        x_std, y_std = "projection_x_coordinate", "projection_y_coordinate"
        try:
            x_units = y_units = crs_obj.axis_info[0].unit_name or "m"
        except (IndexError, AttributeError):
            x_units = y_units = "m"

    ds = xr.Dataset(
        coords={
            "n_pts": (["n_pts"], np.arange(x_da.size)),
            "x": (
                ["n_pts"],
                x_da.data,
                {"units": x_units, "standard_name": x_std},
            ),
            "y": (
                ["n_pts"],
                y_da.data,
                {"units": y_units, "standard_name": y_std},
            ),
            "lat": (
                ["n_pts"],
                lat.data,
                {"units": "degrees_north", "standard_name": "latitude"},
            ),
            "lon": (
                ["n_pts"],
                lon.data,
                {"units": "degrees_east", "standard_name": "longitude"},
            ),
        }
    )

    # Add grid_mapping variable for CF-compliance
    if crs_obj.is_projected:
        gm_name = "spatial_ref"
        ds[gm_name] = ([], 0, crs_obj.to_cf())
        ds.attrs["grid_mapping"] = gm_name
        for var in ["x", "y", "lat", "lon"]:
            ds[var].attrs["grid_mapping"] = gm_name

    ds.attrs["crs"] = crs_obj.to_wkt()

    # Capture extents for provenance (careful with lazy arrays)
    if not is_lazy:
        extent = (
            float(x_da.min()),
            float(x_da.max()),
            float(y_da.min()),
            float(y_da.max()),
        )
        extent_msg = f" Extent: {extent}."
    else:
        extent_msg = ""

    update_history(
        ds,
        f"Created mesh from coordinates and CRS {crs} using xregrid ({backend}).{extent_msg}",
    )

    return ds