from textwrap import dedent

import numpy as np


# Conditional numpy imports for numpy 1.26 and 2.x compatibility
try:
    from numpy.lib.array_utils import normalize_axis_index, normalize_axis_tuple
except ModuleNotFoundError:
    # numpy < 2.0
    from numpy.core.multiarray import normalize_axis_index  # type: ignore[no-redef]
    from numpy.core.numeric import normalize_axis_tuple  # type: ignore[no-redef]


try:
    from numpy._core.einsumfunc import (  # type: ignore[attr-defined]
        _find_contraction,
        _parse_einsum_input,
    )
except ModuleNotFoundError:
    from numpy.core.einsumfunc import (  # type: ignore[no-redef]
        _find_contraction,
        _parse_einsum_input,
    )


# suppress linting warning by "using" the imports here:
__all__ = [
    "_find_contraction",
    "_parse_einsum_input",
    "normalize_axis_index",
    "normalize_axis_tuple",
]


numpy_version_tuple = tuple(int(n) for n in np.__version__.split(".")[:2])
numpy_version = np.lib.NumpyVersion(
    np.__version__
)  # used to compare with version strings, e.g. numpy_version < "1.16.0"
using_numpy_2 = numpy_version >= "2.0.0rc1"


if using_numpy_2:
    ndarray_c_version = np._core._multiarray_umath._get_ndarray_c_version()  # type: ignore[attr-defined]
else:
    ndarray_c_version = np.core._multiarray_umath._get_ndarray_c_version()  # type: ignore[attr-defined]


# used in tests: the type of error thrown if a value is too large for the specified
# numpy data type is different in numpy 2.x
UintOverflowError = OverflowError if using_numpy_2 else TypeError


# to patch up some of the C code, we need to use these special values...
if using_numpy_2:
    numpy_axis_is_none_flag = np.iinfo(np.int32).min  # the value of "NPY_RAVEL_AXIS"
else:
    # 32 is the value used to mark axis = None in Numpy C-API prior to version 2.0
    numpy_axis_is_none_flag = 32


# max number of dims is 64 in numpy 2.x; 32 in older versions
numpy_maxdims = 64 if using_numpy_2 else 32


# function that replicates np.unique from numpy < 2.0
def old_np_unique(
    arr, return_index=False, return_inverse=False, return_counts=False, axis=None
):
    """Replicate np.unique from numpy versions < 2.0"""
    if not return_inverse or not using_numpy_2:
        return np.unique(arr, return_index, return_inverse, return_counts, axis)

    outs = list(np.unique(arr, return_index, return_inverse, return_counts, axis))

    inv_idx = 2 if return_index else 1

    if axis is None:
        outs[inv_idx] = np.ravel(outs[inv_idx])
    else:
        inv_shape = (arr.shape[axis],)
        outs[inv_idx] = outs[inv_idx].reshape(inv_shape)

    return tuple(outs)


# compatibility header for C code
def npy_2_compat_header() -> str:
    """Compatibility header that Numpy suggests is vendored with code that uses Numpy < 2.0 and Numpy 2.x"""
    return dedent("""
    #ifndef NUMPY_CORE_INCLUDE_NUMPY_NPY_2_COMPAT_H_
    #define NUMPY_CORE_INCLUDE_NUMPY_NPY_2_COMPAT_H_


    /*
     * This header is meant to be included by downstream directly for 1.x compat.
     * In that case we need to ensure that users first included the full headers
     * and not just `ndarraytypes.h`.
     */

    #ifndef NPY_FEATURE_VERSION
      #error "The NumPy 2 compat header requires `import_array()` for which "  \\
             "the `ndarraytypes.h` header include is not sufficient.  Please "  \\
             "include it after `numpy/ndarrayobject.h` or similar."  \\
             "" \\
             "To simplify inclusion, you may use `PyArray_ImportNumPy()` " \\
             "which is defined in the compat header and is lightweight (can be)."
    #endif

    #if NPY_ABI_VERSION < 0x02000000
      /*
       * Define 2.0 feature version as it is needed below to decide whether we
       * compile for both 1.x and 2.x (defining it gaurantees 1.x only).
       */
      #define NPY_2_0_API_VERSION 0x00000012
      /*
       * If we are compiling with NumPy 1.x, PyArray_RUNTIME_VERSION so we
       * pretend the `PyArray_RUNTIME_VERSION` is `NPY_FEATURE_VERSION`.
       * This allows downstream to use `PyArray_RUNTIME_VERSION` if they need to.
       */
      #define PyArray_RUNTIME_VERSION NPY_FEATURE_VERSION
      /* Compiling on NumPy 1.x where these are the same: */
      #define PyArray_DescrProto PyArray_Descr
    #endif


    /*
     * Define a better way to call `_import_array()` to simplify backporting as
     * we now require imports more often (necessary to make ABI flexible).
     */
    #ifdef import_array1

    static inline int
    PyArray_ImportNumPyAPI()
    {
        if (NPY_UNLIKELY(PyArray_API == NULL)) {
            import_array1(-1);
        }
        return 0;
    }

    #endif  /* import_array1 */


    /*
     * NPY_DEFAULT_INT
     *
     * The default integer has changed, `NPY_DEFAULT_INT` is available at runtime
     * for use as type number, e.g. `PyArray_DescrFromType(NPY_DEFAULT_INT)`.
     *
     * NPY_RAVEL_AXIS
     *
     * This was introduced in NumPy 2.0 to allow indicating that an axis should be
     * raveled in an operation. Before NumPy 2.0, NPY_MAXDIMS was used for this purpose.
     *
     * NPY_MAXDIMS
     *
     * A constant indicating the maximum number dimensions allowed when creating
     * an ndarray.
     *
     * NPY_NTYPES_LEGACY
     *
     * The number of built-in NumPy dtypes.
     */
    #if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION
        #define NPY_DEFAULT_INT NPY_INTP
        #define NPY_RAVEL_AXIS NPY_MIN_INT
        #define NPY_MAXARGS 64

    #elif NPY_ABI_VERSION < 0x02000000
        #define NPY_DEFAULT_INT NPY_LONG
        #define NPY_RAVEL_AXIS 32
        #define NPY_MAXARGS 32

        /* Aliases of 2.x names to 1.x only equivalent names */
        #define NPY_NTYPES NPY_NTYPES_LEGACY
        #define PyArray_DescrProto PyArray_Descr
        #define _PyArray_LegacyDescr PyArray_Descr
        /* NumPy 2 definition always works, but add it for 1.x only */
        #define PyDataType_ISLEGACY(dtype) (1)
    #else
        #define NPY_DEFAULT_INT  \\
            (PyArray_RUNTIME_VERSION >= NPY_2_0_API_VERSION ? NPY_INTP : NPY_LONG)
        #define NPY_RAVEL_AXIS  \\
            (PyArray_RUNTIME_VERSION >= NPY_2_0_API_VERSION ? -1 : 32)
        #define NPY_MAXARGS  \\
            (PyArray_RUNTIME_VERSION >= NPY_2_0_API_VERSION ? 64 : 32)
    #endif


    /*
     * Access inline functions for descriptor fields.  Except for the first
     * few fields, these needed to be moved (elsize, alignment) for
     * additional space.  Or they are descriptor specific and are not generally
     * available anymore (metadata, c_metadata, subarray, names, fields).
     *
     * Most of these are defined via the `DESCR_ACCESSOR` macro helper.
     */
    #if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION || NPY_ABI_VERSION < 0x02000000
        /* Compiling for 1.x or 2.x only, direct field access is OK: */

        static inline void
        PyDataType_SET_ELSIZE(PyArray_Descr *dtype, npy_intp size)
        {
            dtype->elsize = size;
        }

        static inline npy_uint64
        PyDataType_FLAGS(const PyArray_Descr *dtype)
        {
        #if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION
            return dtype->flags;
        #else
            return (unsigned char)dtype->flags;  /* Need unsigned cast on 1.x */
        #endif
        }

        #define DESCR_ACCESSOR(FIELD, field, type, legacy_only)    \\
            static inline type                                     \\
            PyDataType_##FIELD(const PyArray_Descr *dtype) {       \\
                if (legacy_only && !PyDataType_ISLEGACY(dtype)) {  \\
                    return (type)0;                                \\
                }                                                  \\
                return ((_PyArray_LegacyDescr *)dtype)->field;     \\
            }
    #else  /* compiling for both 1.x and 2.x */

        static inline void
        PyDataType_SET_ELSIZE(PyArray_Descr *dtype, npy_intp size)
        {
            if (PyArray_RUNTIME_VERSION >= NPY_2_0_API_VERSION) {
                ((_PyArray_DescrNumPy2 *)dtype)->elsize = size;
            }
            else {
                ((PyArray_DescrProto *)dtype)->elsize = (int)size;
            }
        }

        static inline npy_uint64
        PyDataType_FLAGS(const PyArray_Descr *dtype)
        {
            if (PyArray_RUNTIME_VERSION >= NPY_2_0_API_VERSION) {
                return ((_PyArray_DescrNumPy2 *)dtype)->flags;
            }
            else {
                return (unsigned char)((PyArray_DescrProto *)dtype)->flags;
            }
        }

        /* Cast to LegacyDescr always fine but needed when `legacy_only` */
        #define DESCR_ACCESSOR(FIELD, field, type, legacy_only)        \\
            static inline type                                         \\
            PyDataType_##FIELD(const PyArray_Descr *dtype) {           \\
                if (legacy_only && !PyDataType_ISLEGACY(dtype)) {      \\
                    return (type)0;                                    \\
                }                                                      \\
                if (PyArray_RUNTIME_VERSION >= NPY_2_0_API_VERSION) {  \\
                    return ((_PyArray_LegacyDescr *)dtype)->field;     \\
                }                                                      \\
                else {                                                 \\
                    return ((PyArray_DescrProto *)dtype)->field;       \\
                }                                                      \\
            }
    #endif

    DESCR_ACCESSOR(ELSIZE, elsize, npy_intp, 0)
    DESCR_ACCESSOR(ALIGNMENT, alignment, npy_intp, 0)
    DESCR_ACCESSOR(METADATA, metadata, PyObject *, 1)
    DESCR_ACCESSOR(SUBARRAY, subarray, PyArray_ArrayDescr *, 1)
    DESCR_ACCESSOR(NAMES, names, PyObject *, 1)
    DESCR_ACCESSOR(FIELDS, fields, PyObject *, 1)
    DESCR_ACCESSOR(C_METADATA, c_metadata, NpyAuxData *, 1)

    #undef DESCR_ACCESSOR


    #if !(defined(NPY_INTERNAL_BUILD) && NPY_INTERNAL_BUILD)
    #if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION
        static inline PyArray_ArrFuncs *
        PyDataType_GetArrFuncs(const PyArray_Descr *descr)
        {
            return _PyDataType_GetArrFuncs(descr);
        }
    #elif NPY_ABI_VERSION < 0x02000000
        static inline PyArray_ArrFuncs *
        PyDataType_GetArrFuncs(const PyArray_Descr *descr)
        {
            return descr->f;
        }
    #else
        static inline PyArray_ArrFuncs *
        PyDataType_GetArrFuncs(const PyArray_Descr *descr)
        {
            if (PyArray_RUNTIME_VERSION >= NPY_2_0_API_VERSION) {
                return _PyDataType_GetArrFuncs(descr);
            }
            else {
                return ((PyArray_DescrProto *)descr)->f;
            }
        }
    #endif


    #endif  /* not internal build */

    #endif  /* NUMPY_CORE_INCLUDE_NUMPY_NPY_2_COMPAT_H_ */

    """)
