# This file is part of the Open Data Cube, see https://opendatacube.org for more information
#
# Copyright (c) 2015-2025 ODC Contributors
# SPDX-License-Identifier: Apache-2.0
from collections.abc import Generator
from math import ceil
from typing import Any

import numpy
import odc.geo.math as geomath
import xarray as xr
from deprecat import deprecat
from odc.geo import SomeResolution
from odc.geo.xr import spatial_dims as xr_spatial_dims

from datacube.migration import ODC2DeprecationWarning


def unsqueeze_data_array(
    da: xr.DataArray, dim: str, pos: int, coord: Any = 0, attrs: dict | None = None
) -> xr.DataArray:
    """
    Add a 1-length dimension to a data array.

    :param da: array to add a 1-length dimension
    :param dim: name of new dimension
    :param pos: position of dim
    :param coord: label of the coordinate on the unsqueezed dimension
    :param attrs: attributes for the coordinate dimension
    :return: A new xarray with a dimension added
    """
    new_dims = list(da.dims)
    new_dims.insert(pos, dim)
    new_shape = da.data.shape[:pos] + (1,) + da.data.shape[pos:]
    new_data = da.data.reshape(new_shape)
    new_coords = dict(da.coords.items())
    new_coords[dim] = xr.DataArray([coord], dims=[dim], attrs=attrs)
    return xr.DataArray(new_data, dims=new_dims, coords=new_coords, attrs=da.attrs)


def unsqueeze_dataset(
    ds: xr.Dataset, dim: str, coord: int = 0, pos: int = 0
) -> xr.Dataset:
    ds = ds.map(unsqueeze_data_array, dim=dim, pos=pos, keep_attrs=True, coord=coord)
    return ds


@deprecat(
    reason="This method has been moved to odc-geo.",
    version="1.9.0",
    category=ODC2DeprecationWarning,
)
def spatial_dims(
    xx: xr.DataArray | xr.Dataset, relaxed: bool = False
) -> tuple[str, str] | None:
    return xr_spatial_dims(xx, relaxed)


@deprecat(
    reason="This method has been moved to odc-geo.",
    version="1.9.0",
    category=ODC2DeprecationWarning,
)
def maybe_zero(x: float, tol: float) -> float:
    return geomath.maybe_zero(x, tol)


@deprecat(
    reason="This method has been moved to odc-geo.",
    version="1.9.0",
    category=ODC2DeprecationWarning,
)
def maybe_int(x: float, tol: float) -> int | float:
    return geomath.maybe_int(x, tol)


@deprecat(
    reason="This method has been moved to odc-geo.",
    version="1.9.0",
    category=ODC2DeprecationWarning,
)
def snap_scale(s, tol: float = 1e-6):
    return geomath.snap_scale(s, tol)


@deprecat(
    reason="This method has been moved to odc-geo.",
    version="1.9.0",
    category=ODC2DeprecationWarning,
)
def clamp(x, lo, up):
    return geomath.clamp(x, lo, up)


@deprecat(
    reason="This method has been moved to odc-geo.",
    version="1.9.0",
    category=ODC2DeprecationWarning,
)
def is_almost_int(x: float, tol: float):
    return geomath.is_almost_int(x, tol)


def dtype_is_float(dtype) -> bool:
    """
    Check if `dtype` is floating-point.
    """
    return numpy.dtype(dtype).kind == "f"


def valid_mask(xx, nodata):
    """
    Compute mask such that xx[mask] contains only valid pixels.
    """
    if dtype_is_float(xx.dtype):
        if nodata is None or numpy.isnan(nodata):
            return ~numpy.isnan(xx)
        return ~numpy.isnan(xx) & (xx != nodata)

    if nodata is None:
        return numpy.full_like(xx, True, dtype=bool)
    return xx != nodata


def invalid_mask(xx, nodata):
    """
    Compute mask such that xx[mask] contains only invalid pixels.
    """
    if dtype_is_float(xx.dtype):
        if nodata is None or numpy.isnan(nodata):
            return numpy.isnan(xx)
        return numpy.isnan(xx) | (xx == nodata)

    if nodata is None:
        return numpy.full_like(xx, False, dtype=bool)
    return xx == nodata


def num2numpy(
    x: int | float | None,
    dtype: str | type[numpy.float32 | numpy.float64] | numpy.dtype,
    ignore_range: bool | None = None,
) -> numpy.float32 | numpy.float64 | numpy.dtype | None:
    """
    Cast python numeric value to numpy.

    :param x: Numerical value to convert to numpy.type
    :param dtype: Destination dtype
    :param ignore_range: If set to True skip range check and cast anyway (for example: -1 -> 255)
                         (Not supported in numpy 2.0+)

    :returns: None if x is None
    :returns: None if x is outside the valid range of dtype and ignore_range is not set
    :returns: dtype.type(x) if x is within range or ignore_range=True
    """
    if x is None:
        return None

    if isinstance(dtype, str | type):
        dtype = numpy.dtype(dtype)

    if ignore_range or dtype.kind == "f":
        return dtype.type(x)

    info = numpy.iinfo(dtype)
    if info.min <= x <= info.max:
        return dtype.type(x)

    return None


@deprecat(
    reason="This method has been moved to odc-geo.",
    version="1.9.0",
    category=ODC2DeprecationWarning,
)
def data_resolution_and_offset(data, fallback_resolution: float | None = None):
    return geomath.data_resolution_and_offset(data, fallback_resolution)


@deprecat(
    reason="This method has been moved to odc-geo.",
    version="1.9.0",
    category=ODC2DeprecationWarning,
)
def affine_from_axis(xx, yy, fallback_resolution: SomeResolution | None = None):
    return geomath.affine_from_axis(xx, yy, fallback_resolution)


def iter_slices(
    shape: tuple[int, ...], chunk_size: tuple[int, ...]
) -> Generator[tuple]:
    """
    Generate slices for a given shape.

    E.g. ``shape=(4000, 4000), chunk_size=(500, 500)``
    Would yield 64 tuples of slices, each indexing 500x500.

    If the shape is not divisible by the chunk_size, the last chunk in each dimension will be smaller.

    :param shape: Shape of an array
    :param chunk_size: length of each slice for each dimension
    :return: Yields slices that can be used on an array of the given shape

    >>> list(iter_slices((5,), (2,)))
    [(slice(0, 2, None),), (slice(2, 4, None),), (slice(4, 5, None),)]
    """
    assert len(shape) == len(chunk_size)
    num_grid_chunks = [ceil(s / float(c)) for s, c in zip(shape, chunk_size)]
    for grid_index in numpy.ndindex(*num_grid_chunks):
        yield tuple(
            slice(min(d * c, stop), min((d + 1) * c, stop))
            for d, c, stop in zip(grid_index, chunk_size, shape)
        )
