"""
Slight modification to numpy's histogramdd to suit Igor like binning. Based on numpy 2.4.0.
https://github.com/numpy/numpy/blob/v2.4.0/numpy/lib/_histograms_impl.py

Copyright (c) 2005-2025, NumPy Developers.
All rights reserved.
"""

import operator

import numpy as np
from numpy.lib._histograms_impl import _get_outer_edges

# range is a keyword argument to many functions, so save the builtin so they can use it.
_range = range


def igor_histogramdd(sample, bins=10, range=None, density=None, weights=None):
    """
    Compute the multidimensional histogram of some data.

    Parameters
    ----------
    sample : (N, D) array, or (N, D) array_like
        The data to be histogrammed.

        Note the unusual interpretation of sample when an array_like:

        * When an array, each row is a coordinate in a D-dimensional space -
          such as ``histogramdd(np.array([p1, p2, p3]))``.
        * When an array_like, each element is the list of values for single
          coordinate - such as ``histogramdd((X, Y, Z))``.

        The first form should be preferred.

    bins : sequence or int, optional
        The bin specification:

        * A sequence of arrays describing the monotonically increasing bin
          edges along each dimension.
        * The number of bins for each dimension (nx, ny, ... =bins)
        * The number of bins for all dimensions (nx=ny=...=bins).

    range : sequence, optional
        A sequence of length D, each an optional (lower, upper) tuple giving
        the outer bin edges to be used if the edges are not given explicitly in
        `bins`.
        An entry of None in the sequence results in the minimum and maximum
        values being used for the corresponding dimension.
        The default, None, is equivalent to passing a tuple of D None values.
    density : bool, optional
        If False, the default, returns the number of samples in each bin.
        If True, returns the probability *density* function at the bin,
        ``bin_count / sample_count / bin_volume``.
    weights : (N,) array_like, optional
        An array of values `w_i` weighing each sample `(x_i, y_i, z_i, ...)`.
        Weights are normalized to 1 if density is True. If density is False,
        the values of the returned histogram are equal to the sum of the
        weights belonging to the samples falling into each bin.

    Returns
    -------
    H : ndarray
        The multidimensional histogram of sample x. See density and weights
        for the different possible semantics.
    edges : tuple of ndarrays
        A tuple of D arrays describing the bin edges for each dimension.

    See Also
    --------
    histogram: 1-D histogram
    histogram2d: 2-D histogram

    Examples
    --------
    >>> import numpy as np
    >>> rng = np.random.default_rng()
    >>> r = rng.normal(size=(100,3))
    >>> H, edges = np.histogramdd(r, bins = (5, 8, 4))
    >>> H.shape, edges[0].size, edges[1].size, edges[2].size
    ((5, 8, 4), 6, 9, 5)

    """

    # shift bins by half a bin (Igor like binning)
    bins = [edges.copy() for edges in bins]
    for i in _range(3):
        bins[i] += (bins[i][1] - bins[i][0]) / 2

    try:
        # Sample is an ND-array.
        _, d = sample.shape
    except (AttributeError, ValueError):
        # Sample is a sequence of 1D arrays.
        sample = np.atleast_2d(sample).T
        _, d = sample.shape

    nbin = np.empty(d, np.intp)
    edges = d * [None]
    dedges = d * [None]
    if weights is not None:
        weights = np.asarray(weights)

    try:
        m = len(bins)
        if m != d:
            raise ValueError(
                "The dimension of bins must be equal to the dimension of the sample x."
            )
    except TypeError:
        # bins is an integer
        bins = d * [bins]

    # normalize the range argument
    if range is None:
        range = (None,) * d
    elif len(range) != d:
        raise ValueError("range argument must have one entry per dimension")

    # Create edge arrays
    for i in _range(d):
        if np.ndim(bins[i]) == 0:
            if bins[i] < 1:
                raise ValueError(f"`bins[{i}]` must be positive, when an integer")
            smin, smax = _get_outer_edges(sample[:, i], range[i])
            try:
                n = operator.index(bins[i])

            except TypeError as e:
                raise TypeError(f"`bins[{i}]` must be an integer, when a scalar") from e

            edges[i] = np.linspace(smin, smax, n + 1)
        elif np.ndim(bins[i]) == 1:
            edges[i] = np.asarray(bins[i])
            if np.any(edges[i][:-1] > edges[i][1:]):
                raise ValueError(f"`bins[{i}]` must be monotonically increasing, when an array")
        else:
            raise ValueError(f"`bins[{i}]` must be a scalar or 1d array")

        nbin[i] = len(edges[i]) + 1  # includes an outlier on each end
        dedges[i] = np.diff(edges[i])

    # Compute the bin number each sample falls into.
    n_count = tuple(
        # avoid np.digitize to work around gh-11022
        np.searchsorted(edges[i], sample[:, i], side="right")
        for i in _range(d)
    )

    # Removed code: Keep values that fall on an edge in the left bin (used to be right bin).

    # Compute the sample indices in the flattened histogram matrix.
    # This raises an error if the array is too large.
    xy = np.ravel_multi_index(n_count, nbin)

    # Compute the number of repetitions in xy and assign it to the
    # flattened histmat.
    hist = np.bincount(xy, weights, minlength=nbin.prod())

    # Shape into a proper matrix
    hist = hist.reshape(nbin)

    # This preserves the (bad) behavior observed in gh-7845, for now.
    hist = hist.astype(float, casting="safe")

    # Remove outliers (indices 0 and -1 for each dimension).
    core = d * (slice(1, -1),)
    hist = hist[core]

    if density:
        # calculate the probability density function
        s = hist.sum()
        for i in _range(d):
            shape = np.ones(d, int)
            shape[i] = nbin[i] - 2
            hist = hist / dedges[i].reshape(shape)
        hist /= s

    if (hist.shape != nbin - 2).any():
        raise RuntimeError("Internal Shape Error")
    return hist, edges
