"""NumPy-related utility types and functions."""

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser General Public License for more details.

# You should have received a copy of the GNU Lesser General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.

from __future__ import annotations
from collections.abc import Callable
from typing import Any, ParamSpec, TypeAlias, TypeVar
import numpy as np
import numba  # type: ignore

# tuple[int, ...] used for all shapes to fix regression in Numpy 2.2.6 typing.
# Seems to be fine in Numpy 2.3, but that's currently not supported by Numba.

BoolArray1D: TypeAlias = np.ndarray[tuple[int, ...], np.dtype[np.bool_]]
"""Type alias for 1D bool NumPy arrays."""

UIntArray1D: TypeAlias = np.ndarray[tuple[int, ...], np.dtype[np.unsignedinteger[Any]]]
"""Type alias for 1D uint NumPy arrays."""

UInt8Array1D: TypeAlias = np.ndarray[tuple[int, ...], np.dtype[np.uint8]]
"""Type alias for 1D uint8 NumPy arrays."""

UInt16Array1D: TypeAlias = np.ndarray[tuple[int, ...], np.dtype[np.uint16]]
"""Type alias for 1D uint16 NumPy arrays."""

UInt8Array2D: TypeAlias = np.ndarray[tuple[int, ...], np.dtype[np.uint8]]
"""Type alias for 2D uint8 NumPy arrays."""

FloatArray1D: TypeAlias = np.ndarray[tuple[int, ...], np.dtype[np.floating[Any]]]
"""Type alias for 1D float NumPy arrays."""

ComplexArray1D: TypeAlias = np.ndarray[
    tuple[int, ...], np.dtype[np.complexfloating[Any]]
]
"""Type alias for 1D complex128 NumPy arrays."""

Complex128Array1D: TypeAlias = np.ndarray[tuple[int, ...], np.dtype[np.complex128]]
"""Type alias for 1D complex128 NumPy arrays."""

Complex128Array2D: TypeAlias = np.ndarray[tuple[int, ...], np.dtype[np.complex128]]
"""Type alias for 2D complex128 NumPy arrays."""

RNG: TypeAlias = np.random.Generator
"""Typa alias for a NumPy random number generator."""

ShapeT = TypeVar("ShapeT", bound=tuple[int, ...])
"""Type variable for the shape of NumPy arrays."""

_numba_jit = numba.jit(nopython=True, cache=True)
"""Decorator to apply :func:`numba.jit` with desired settings."""

_P = ParamSpec("_P")
"""Type alias for a generic parameter list."""

_R = TypeVar("_R")
"""Type alias for a generic return type."""


def numba_jit(func: Callable[_P, _R]) -> Callable[_P, _R]:
    """Decorator to apply :func:`numba.jit` with desired settings."""
    return _numba_jit(func)  # type: ignore


def canonicalize_phase(
    array: np.ndarray[ShapeT, np.dtype[np.complex128]],
    *,
    tol: float = 1e-8,
) -> None:
    """Normalises the phase of a complex array (in place)."""
    assert tol > 0
    idx = np.argmax(abs(array) >= tol)
    val = array.flatten()[idx]
    if (aval := abs(val)) >= tol:
        array *= aval / val
