from math import exp, sqrt
from typing import Any, Callable, Dict, Generic, Iterator, Optional, Sequence, Tuple, TypeVar, Union

import numpy as np

# from .numba import opjit

T = TypeVar("T")

RGB_YELLOW = (255, 255, 0)
RGB_WHITE = (255, 255, 255)


def assert_square(name: str, value: np.ndarray) -> None:
    """Raises a ValueError if matrix `value` is not square.
    value: [A, A]
    """

    if not len(value.shape) == 2 or value.shape[0] != value.shape[1]:
        raise ValueError(f"{name} must be a square")


def normalized_choice(
    p_ind, p_prob
):  # is this neccessary? I think `np.random.choice` can handle unnormalized probabilities
    return np.random.choice(p_ind, p=p_prob / np.sum(p_prob))


def shannon_entropy(ps: np.ndarray, base: int = 2) -> float:
    """Calculates the Shannon entropy for probabilities `ps` with `base`."""

    return -np.sum(ps * np.log(ps) / np.log(base))


def is_rgb(img: np.ndarray) -> bool:
    """Simply tests if `img` has 3 channels."""

    return len(img.shape) >= 1 and img.shape[-1] == 3


# @opjit() # np.errstate doesn't work in numba...
def rgb_to_hsi(image: np.ndarray) -> np.ndarray:
    """Converts an array [..., channels] of RGB values to HSI color values (H in rad).
    RGB values are assumed to be normalized to (0, 1).
    """

    if not is_rgb(image):
        raise ValueError("Input needs to be an array of RGB values")

    r = image[..., 0]
    g = image[..., 1]
    b = image[..., 2]

    out = np.zeros_like(image)

    # allequal = (img == img[:, :, 0, np.newaxis]).all(axis=-1)

    with np.errstate(invalid="ignore"):
        tmp = (2.0 * r - g - b) / 2.0 / np.sqrt((r - g) ** 2 + (r - b) * (g - b))  # if r==g==b then 0/0

        theta = np.arccos(np.clip(tmp, -1.0, +1.0))

        out[..., 0] = np.where(b <= g, theta, 2 * np.pi - theta)  # H
        out[..., 2] = np.sum(image, axis=-1) / 3.0  # I
        out[..., 1] = 1 - np.amin(image, axis=-1) / out[..., 2]  # S if r==g==b==0 then 0/0

    np.nan_to_num(out[..., 0:2], copy=False)

    return out


# @opjit() # np.dot with more than 2 dimensions is not supported...
def rgb_to_ycbcr(image: np.ndarray) -> np.ndarray:
    """Converts an array [..., channels] of RGB values to Digital Y'CbCr (0-255).
    RGB values are assumed to be normalized to (0, 1).
    Don't forget to cast to uint8 for pillow.
    """

    """  from RGB (0-1).
    """

    if not is_rgb(image):
        raise ValueError("Input needs to be an array of RGB values")

    m = np.array(
        [
            [+065.481, +128.553, +024.966],
            [-037.797, -074.203, +112.000],
            [+112.000, -093.786, -018.214],
        ]
    )
    a = np.array([16, 128, 128])

    return np.dot(image, m.T) + a


def random_triangular_matrix(size: int, lower: bool = True) -> np.ndarray:
    """Returns a triangular matrix with random value between 0 and 1 uniformly."""

    a = np.random.uniform(0, 1, (size, size))
    if lower:
        ind = np.triu_indices(5, 1)
    else:
        ind = np.tril_indices(5, 1)
    a[ind] = 0

    return a


def is_square(A: np.ndarray) -> bool:
    if len(A.shape) != 2:
        return False

    if A.shape[0] != A.shape[1]:
        return False

    return True


def batch_vTAv(A: np.ndarray, v: np.ndarray) -> np.ndarray:
    """Performs batched calculation of `v^T A v` transform.
    Special case of bilinear form `x^T A y`

    A: float[B, X, X]
    v: float[B, X]

    returns: float[B]
    """

    """ Faster than
        Av = np.matmul(A, v[...,:,None]) # [B, X, 1]
        return np.matmul(v[...,None,:], Av).squeeze((-2, -1)) # [B]
    """

    return np.einsum("...k,...kl,...l->...", v, A, v)


def batch_inner(a: np.ndarray, b: np.ndarray, verify: bool = True) -> np.ndarray:
    """Performs a batched inner product over the last dimension.
    Replacement for deprecated `from numpy.core.umath_tests import inner1d`.
    Shapes: (B, X), (B, X) -> (B, )
    """

    if verify and a.shape != b.shape:
        raise ValueError("All dimensions have to be equal")

    if a.shape[-1] == 0:
        return np.empty_like(a)

    return np.einsum("...i,...i->...", a, b)  # faster than np.sum(a * b, axis=-1)


def batch_outer(a: np.ndarray, b: np.ndarray, verify: bool = True) -> np.ndarray:
    """Performs a batched outer product over the last dimension.
    Shapes: (B, X), (B, Y) -> (B, X, Y)
    """

    if verify and a.shape[:-1] != b.shape[:-1]:
        raise ValueError("All except the last dimension have to be equal")

    return np.einsum("...i,...j->...ij", a, b)  # slightly faster than np.multiply(a[...,:,None], b[...,None,:])


def batchtopk(
    probs: np.ndarray, k: Optional[int] = None, axis: int = -1, reverse: bool = False
) -> Tuple[np.ndarray, np.ndarray]:
    """`probs` values ndarray
    `k` take the smallest `k` elements, if `reverse` is False
            and the largest `k` if `reverse` is True
    `axis` sorting and selection axis.
    """

    if k is not None and k <= 0:
        raise ValueError("k must be larger than 0. Use None to chose all elements.")

    if axis != -1:
        raise ValueError("Only last axis supported atm")

    if len(probs.shape) <= 1:
        raise ValueError("probs must be at least 2-dimensional")

    if reverse:
        sign = -1
    else:
        sign = 1

    indices = np.argsort(sign * probs, axis=-1)  # use argpartition?
    probs = np.take_along_axis(probs, indices[..., :k], axis=-1)

    return indices, probs


# @opjit() # np.diagonal not supported by numba as of version 0.52.0
def logtrace(m: np.ndarray) -> np.ndarray:
    """Calcuates the sum of the logs of the diagonal elements (batchwise if necessary)
    m: [..., x, x]
    """

    """ note: performance cannot easily be improve by numba.
        `np.diagonal` not supported by numba 0.52.0
    """

    return np.sum(np.log(np.diagonal(m, axis1=-2, axis2=-1)), axis=-1)


def shiftedexp(pvals: np.ndarray) -> np.ndarray:
    """Shifts `pvals` by the largest value in the last dimension
    before the exp is calculated to prevent overflow (batchwise if necessary).
    Can be used if probabilities are normalized again later.
    """

    if pvals.shape[-1] == 0:
        return np.empty_like(pvals)

    return np.exp(pvals - np.amax(pvals, axis=-1)[..., None])


class Sampler:
    """Sample from discrete CDF."""

    def __init__(self, cdf: np.ndarray) -> None:
        self.cdf = cdf
        self.psum = cdf[-1]

    def __call__(self) -> int:
        """Sample one."""

        rand = np.random.uniform(0, self.psum)
        return np.searchsorted(self.cdf, rand, side="right")

    def sample(self, n: int) -> np.ndarray:
        """Sample `n`."""

        rands = np.random.uniform(0, self.psum, n)
        return np.searchsorted(self.cdf, rands, side="right")

    def pdf(self, n: int, minlength: Optional[int] = None) -> np.ndarray:
        out = np.bincount(self.sample(n), minlength=minlength)
        return out / n


def sample_probabilities(pvals: np.ndarray) -> Callable[[], int]:
    """Sample from list of probabilities `pvals` with replacement.
    The probabilities don't need to be normalized.
    """

    return Sampler(np.cumsum(pvals))


class UnboundedSparseMatrix(Generic[T]):
    def __init__(self, dtype: type = float) -> None:
        self.dtype = dtype
        self.zero = self.dtype(0)
        self.m: Dict[Tuple[int, int], T] = {}
        self.cols = 0
        self.rows = 0

    def __getitem__(self, slice: Tuple[int, int]) -> T:
        return self.m.get(slice, self.zero)

    def __setitem__(self, slice: Tuple[int, int], value: T) -> None:
        c, r = slice
        self.cols = max(self.cols, c + 1)
        self.rows = max(self.rows, r + 1)
        self.m[slice] = value

    def todense(self) -> np.ndarray:
        ret = np.zeros((self.cols, self.rows), dtype=self.dtype)

        for slice, value in self.m.items():
            ret[slice] = value

        return ret


def normalize(pvals: np.ndarray) -> np.ndarray:
    return pvals / np.sum(pvals)


def categorical(pvals: np.ndarray) -> int:
    """Sample from the categorical distribution using `pvals`.
    See: https://en.wikipedia.org/wiki/Categorical_distribution
    """

    return sample_probabilities(pvals)()  # faster than: np.argmax(np.random.multinomial(1, normalize(pvals)))


def population2cdf(population: np.ndarray) -> np.ndarray:
    """Convert a population (list of observations) to a CDF."""

    population = np.sort(population)
    return np.searchsorted(population, population, side="right") / len(population)


def pmf2cdf(pdf: np.ndarray) -> np.ndarray:
    """Convert a discrete PDF into a discrete CDF."""

    cdf = np.cumsum(pdf)
    return cdf / cdf[-1]


def _two_sample_kolmogorov_smirnov_same_length(
    cdf1: np.ndarray, cdf2: np.ndarray, n1: int, n2: int
) -> Tuple[float, float]:
    # note: yields different results as `scipy.stats.ks_2samp`

    if len(cdf1) != len(cdf2):
        raise ValueError("Both CDFs must have the same length")

    D = np.amax(np.abs(cdf1 - cdf2))  # K-S statistic
    level = exp(-2 * (D / sqrt((n1 + n2) / (n1 * n2))) ** 2)
    return D, level


def _two_sample_kolmogorov_smirnov_population(
    p1: np.ndarray, p2: np.ndarray, alpha: float = 0.05
) -> Tuple[float, float, bool]:
    # note: yields different results as `scipy.stats.ks_2samp`

    cdf1 = population2cdf(p1)
    cdf2 = population2cdf(p2)

    statistic, pvalue = _two_sample_kolmogorov_smirnov_same_length(cdf1, cdf2, len(cdf1), len(cdf2))
    reject = pvalue < alpha
    return statistic, pvalue, reject


def _two_sample_kolmogorov_smirnov_pmf(
    pmf1: np.ndarray, pmf2: np.ndarray, alpha: float = 0.05
) -> Tuple[float, float, bool]:
    # note: yields different results as `scipy.stats.ks_2samp`
    """Tests the null hypothesis that both samples belong to the same distribution.
    See: https://en.wikipedia.org/wiki/Kolmogorov%E2%80%93Smirnov_test#Two-sample_Kolmogorov%E2%80%93Smirnov_test
    """

    cdf1 = np.cumsum(pmf1)
    cdf2 = np.cumsum(pmf2)

    n1 = cdf1[-1]
    n2 = cdf2[-1]

    # cannot be inplace because of type conversion
    cdf1 = cdf1 / n1
    cdf2 = cdf2 / n2

    statistic, pvalue = _two_sample_kolmogorov_smirnov_same_length(cdf1, cdf2, n1, n2)
    reject = pvalue < alpha
    return statistic, pvalue, reject


def inf_matrix_power(pm: np.ndarray, dtype=np.float64) -> np.ndarray:
    """Calculate stochastic matrix `pm` to the power of infinity,
    by finding the eigenvector which corresponds to the eigenvalue 1.
    """

    w, v = np.linalg.eig(
        pm
    )  # scipy.linalg.eig would probably by faster as it can return the left and right eigen vectors

    if not np.isclose(w[0], 1.0):
        raise ValueError("The first eigenvalue is not none. Is this a right stochastic matrix?")

    vi = np.linalg.inv(v)
    d = np.zeros(pm.shape[0], dtype=dtype)
    d[0] = 1.0

    return np.matmul(v, np.matmul(np.diag(d), vi))


def decompress(selectors: np.ndarray, data: np.ndarray, default) -> np.ndarray:
    """Same result as:
    `np.array(list(genutility.iter.decompress(selectors, iter(data), default)))`
    but faster.
    `selectors` should be a bool array.
    """

    out = np.full(len(selectors), default)
    out[selectors] = data
    return out


def unblock(arr: np.ndarray, n1: int, n2: int, axis1: int = -1, axis2: int = -2, blocksize: bool = False) -> np.ndarray:
    """Inverse of np.block.
    Set axis to (-2, -1) to modify the order of the result.
    """

    """ test (stackoverflow): Ok, so considering I have N block matrices with bm x bn dimension and want to stack them in a m x n matrix, provided N = m x n, I would then have x.reshape(m,n,bm,bn).swapaxes(1,2).reshape(bm*m,-1)
    """

    s = np.array(arr.shape)
    if s[axis1] % n1 != 0 or s[axis2] % n2 != 0:
        raise ValueError(f"{s[axis1]}x{s[axis2]} does not divide by {n1}x{n2}")

    if blocksize:
        n1 = s[axis1] // n1
        n2 = s[axis2] // n2

    # this first .split adds a new dimensions on the outside, so if a absolute index
    # is given for the second axis it must be moved one to the right
    if axis2 >= 0:
        _axis2 = axis2 + 1
    else:
        _axis2 = axis2

    arr = np.array(np.split(arr, n1, axis1))
    arr = np.array(np.split(arr, n2, _axis2))

    inv_blocksize = n1 * n2
    total = s[axis1] * s[axis2]
    s[axis2] = inv_blocksize
    s[axis1] = total // inv_blocksize

    return np.reshape(arr, s)


def unblock2d(arr: np.ndarray, n1: int, n2: int, blocksize: bool = False) -> np.ndarray:
    """
    arr: [width, height, ...]

    returns: (block_num, block_size)
    """

    s = np.array(arr.shape)
    if blocksize:
        bs1 = n1
        bs2 = n2
        bn1 = s[0] // n1
        bn2 = s[1] // n2
    else:
        bn1 = n1
        bn2 = n2
        bs1 = s[0] // n1
        bs2 = s[1] // n2

    if len(arr.shape) > 2:
        return arr.reshape(bn1, bs1, bn2, bs2, -1).swapaxes(1, 2).reshape(bn1 * bn2, bs1 * bs2, -1)
    else:
        return arr.reshape(bn1, bs1, bn2, bs2).swapaxes(1, 2).reshape(bn1 * bn2, bs1 * bs2)


def block2d(arr: np.ndarray, bn1: int, bn2: int, bs1: int, bs2: int) -> np.ndarray:
    """
    arr: (block_num, block_size)

    returns: (width, height)
    """

    s = np.array(arr.shape)
    assert s[0] == bn1 * bn2
    assert s[1] == bs1 * bs2

    return arr.reshape(bn1, bn2, bs1, bs2).swapaxes(1, 2).reshape(bn1 * bs1, bn2 * bs2)


def remove_color(img: np.ndarray, ratio: float, neutral_color: Tuple[int, int, int] = RGB_WHITE) -> None:
    """Replace colored pixels with a `neutral_color`. The `ratio` defines the 'colorfulness' above
    which level the pixel should be replace.
    I.e. if the `ratio` is 1 nothing will be replaced,
    if `ratio` is 0 only strict greys are kept unmodified.
    """

    channels = img.shape[-1]
    assert channels == 3, "Not a 3 channel color image"

    norm = np.std(np.array(RGB_YELLOW))  # this is the same for all pure colors

    sd = np.std(img, axis=-1)
    img[sd > ratio * norm] = neutral_color


# @opjit(cache=False) doesn't help a lot, because it's a generator
def sliding_window_2d(
    image: np.ndarray, window_size: Tuple[int, int], step_size: Tuple[int, int] = (1, 1)
) -> Iterator[np.ndarray]:
    win_x, win_y = window_size
    step_x, step_y = step_size
    height, width = image.shape[0:2]

    for y in range(0, height - win_y + 1, step_y):
        for x in range(0, width - win_x + 1, step_x):
            yield image[y : y + win_y, x : x + win_x, ...]


def histogram_correlation(hist1: np.ndarray, hist2: np.ndarray) -> np.ndarray:
    """Input shape of `hist1` and `hist2`: [batch..., levels]. The correlation is calculated over `levels`
    and then batched over the remaining dimensions.
    Similar to `cv2.compareHist(hist1, hist2, cv2.HISTCMP_CORREL)`.
    """

    if hist1.ndim < 1 or hist1.shape[-1] < 1 or hist1.shape != hist2.shape:
        raise ValueError(f"Histogram shapes unsupported: {hist1.shape}, {hist2.shape}")

    h1norm = (hist1.T - np.mean(hist1, axis=-1)).T
    h2norm = (hist2.T - np.mean(hist2, axis=-1)).T

    if np.count_nonzero(h1norm) == 0 or np.count_nonzero(h2norm) == 0:
        raise ValueError("Both histograms are completely flat")

    num = np.sum(h1norm * h2norm, axis=-1)
    denom = np.sqrt(np.sum(h1norm**2, axis=-1) * np.sum(h2norm**2, axis=-1))
    return num / denom


def bincount_batch(x: np.ndarray, axis: int = -1, minlength: int = 0) -> np.ndarray:
    if x.shape[axis] == 0:
        raise ValueError("Specified axis of x cannot be 0")

    minlength = max(minlength, x.max() + 1)
    return np.apply_along_axis(np.bincount, axis, x, minlength=minlength)


def stochastic(x: np.ndarray) -> np.ndarray:
    """It normalizes the last dimension of an ndarray to sum to 1.
    It can be used to convert (batches of) vectors to stochastic vectors
    or (batches of) matrices to *right* stochastic matrices.
    Right stochastic matrices are also called transitions matrices.
    """

    n = np.linalg.norm(x, 1, axis=-1, keepdims=True)
    # n = np.sum(x, axis=-1, keepdims=True) # todo: same result (except dtype), which is faster?

    with np.errstate(invalid="raise"):  # see: `normalized`
        return x / n


def sequence_mask(lengths: np.ndarray, maxlen: Optional[int] = None, dtype: Any = None) -> np.ndarray:
    """cf. tf.sequence_mask
    lengths: [N]
    """

    if not lengths:
        return np.array([], dtype=dtype or np.bool_)

    if maxlen is None:
        maxlen = max(lengths)
    row_vector = np.arange(maxlen)
    matrix = np.expand_dims(lengths, -1)

    ret = row_vector < matrix
    if dtype is not None:
        ret = ret.astype(dtype)

    return ret


# @opjit() # doesn't work atm because numba doesn't support `None` indexing
def _viterbi_dense_masked(
    p_emit: np.ndarray, p_trans: np.ndarray, p_trans0: np.ndarray, mask: np.ndarray
) -> np.ndarray:
    batch_size, T, N = p_emit.shape

    N1, N2 = p_trans.shape
    assert N == N1 == N2

    assert (batch_size, T) == mask.shape
    assert (N,) == p_trans0.shape

    trellis = np.zeros((T, batch_size, N), dtype=p_emit.dtype)
    states = np.zeros((T, batch_size, N), dtype=np.intp)

    # even if sequences of length 0 should be allowed, this does not have to be masked, as there cannot be a wrong result
    trellis[0] = p_trans0 + p_emit[:, 0, :]  # broadcast p_trans0: [N] -> batch_size*[[N]]

    for t in range(1, T):
        masked_p_trans = (
            mask[:, t, None, None] * p_trans[None, :, :]
        )  # [batch_size, 1, 1] * [1, N, N] == [batch_size, N, N]
        weighted_scores = (
            trellis[t - 1, :, :, None] + masked_p_trans
        )  # [batch_size, N, N] # scores and p_trans broadcasted
        max_scores = np.amax(weighted_scores, axis=1)  # [batch_size, N]
        trellis[t] = max_scores + p_emit[:, t, :]  # [batch_size, N] remember highest score of each path
        states[t] = np.argmax(
            weighted_scores, axis=1
        )  # [batch_size, N] remember index of best path, should be repeated for padding vals

    tokens = np.zeros((T, batch_size), dtype=np.intp)
    tokens[T - 1] = np.argmax(trellis[T - 1], axis=1)  # [batch_size]

    for t in range(T - 1, 0, -1):
        tokens[t - 1] = states[t, np.arange(batch_size), tokens[t]]  # [batch_size]

    return tokens.T  # [batch_size, T]


def viterbi_dense(
    p_emit: np.ndarray, p_trans: np.ndarray, p_trans0: Optional[np.ndarray] = None, mask: Optional[np.ndarray] = None
) -> np.ndarray:
    """Viterbi algorithm for finding the optimal path.
    One square transition matrix can be specified.

    T: max sequence length
    N: vocab size
    batch_size: to vectorize

    p_emit: masked emission scores [batch_size, T, N] float
    p_trans: transition scores [N, N] float
    p_trans0: initial transition scores [N] float
    mask: mask out padding values in transitions [batch_size, T] float
    """

    batch_size, T, N = p_emit.shape

    if mask is None:
        mask = np.ones((batch_size, T), dtype=p_trans.dtype)

    if p_trans0 is None:
        p_trans0 = np.zeros(N, dtype=p_emit.dtype)

    return _viterbi_dense_masked(p_emit, p_trans, p_trans0, mask)


# @opjit() # doesn't work atm because numba cannot store arbitrary elements in lists I think
def viterbi_sparse(p_emit: Sequence[np.ndarray], p_trans: Sequence[np.ndarray]) -> np.ndarray:
    """Viterbi algorithm for finding the optimal path.
    The number of emission probabilities per index can vary
    and a separate matrix can be specified for each transition.

    p_emit: Sequence of [x] matrices, where x can vary for every entry
    p_trans: Sequence of [y, z] matrices, where y and z can vary for every entry

    todo: using `Iterable`s instead of Sequence`s should be possible.
    """

    T = len(p_emit)

    assert T - 1 == len(p_trans)

    trellis = [p_emit[0]]
    states = [None]

    for t in range(1, T):
        weighted_scores = trellis[-1][:, None] + p_trans[t - 1]  # [x, y] # scores and p_trans broadcasted
        max_scores = np.amax(weighted_scores, axis=0)  # [y]
        trellis.append(np.add(max_scores, p_emit[t]))  # [y] remember highest score of each path
        states.append(np.argmax(weighted_scores, axis=0))  # [y] remember index of best path

    assert len(trellis) == T and len(states) == T

    tokens = [None] * T  # [T]
    tokens[-1] = np.argmax(trellis[-1], axis=0)  # []

    for t in range(T - 1, 0, -1):
        tokens[t - 1] = states[t][tokens[t]]  # []

    return tokens


_bit_counts = np.array([int(bin(x).count("1")) for x in range(256)]).astype(np.uint8)


def hamming_dist_packed(a: np.ndarray, b: np.ndarray, axis: Optional[int] = -1) -> np.ndarray:
    return np.sum(_bit_counts[np.bitwise_xor(a, b)], axis=axis)


def get_num_chunks(shape: np.ndarray, chunksize: np.ndarray) -> int:
    return np.prod(np.ceil(shape / chunksize).astype(np.int_))


def broadcast_shapes(*shapes: Tuple[int, ...]) -> Tuple[int, ...]:
    """np.broadcast_shapes requires `numpy==1.20.0`,
    which is not available for `python < 3.7`.
    """

    arrays = [np.empty(shape) for shape in shapes]
    return np.broadcast(*arrays).shape


def hamming_dist_packed_chunked(
    a: np.ndarray, b: np.ndarray, chunksize: Tuple[int, ...], axis: Optional[int] = -1
) -> Iterator[Tuple[Tuple[int, ...], np.ndarray]]:
    outshape = broadcast_shapes(a.shape, b.shape)
    select_broadcasted_axis = slice(0, 1)

    if len(outshape) - len(chunksize) != 1:
        raise ValueError("Length of `chunksize` must be one less the number of input dimensions")

    if len(outshape) == 2:
        for x in range(0, outshape[0], chunksize[0]):
            if a.shape[0] != outshape[0]:
                aix = select_broadcasted_axis
            else:
                aix = slice(x, x + chunksize[0])

            if b.shape[0] != outshape[0]:
                bix = select_broadcasted_axis
            else:
                bix = slice(x, x + chunksize[0])

            yield (x,), hamming_dist_packed(a[aix, :], b[bix, :], axis=axis)

    elif len(outshape) == 3:
        for x in range(0, outshape[0], chunksize[0]):
            for y in range(0, outshape[1], chunksize[1]):
                if a.shape[0] != outshape[0]:
                    aix = select_broadcasted_axis
                else:
                    aix = slice(x, x + chunksize[0])

                if b.shape[0] != outshape[0]:
                    bix = select_broadcasted_axis
                else:
                    bix = slice(x, x + chunksize[0])

                if a.shape[1] != outshape[1]:
                    aiy = select_broadcasted_axis
                else:
                    aiy = slice(y, y + chunksize[1])

                if b.shape[1] != outshape[1]:
                    biy = select_broadcasted_axis
                else:
                    biy = slice(y, y + chunksize[1])

                yield (x, y), hamming_dist_packed(a[aix, aiy, :], b[bix, biy, :], axis=axis)

    else:
        raise ValueError("Input must either be 2 or 3 dimensional")


def hamming_dist(a: np.ndarray, b: np.ndarray, axis: Optional[int] = -1) -> np.ndarray:
    return np.count_nonzero(a != b, axis=axis)


def center_of_mass_2d(arr: np.ndarray, dtype=np.float32) -> np.ndarray:
    """Batched center of mass calculation of 2d arrays
    `arr`: [..., x, y]
    """

    total = np.sum(arr, axis=(-1, -2))
    grids = np.ogrid[[slice(0, i) for i in arr.shape[-2:]]]
    with np.errstate(invalid="ignore"):
        results = np.array([np.sum(arr * grid.astype(dtype), axis=(-1, -2)) / total for grid in grids], dtype=dtype)
    results = np.moveaxis(results, 0, -1)
    return results


def ceildiv(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    return -(a // -b)


def image_grid(images: np.ndarray, nrow: int = 8, fill_value=(0, 0, 0)) -> np.ndarray:
    """Creates a picture grid (left to right, top to bottom).
    `images`: Input array of shape [n, height, widght, channel]
    `nrow`: Number of images per row
    `fill_value`: Pad empty space with pixels of shape [channel]
    """

    if not images.ndim == 4:
        raise ValueError("Input shape must be [n, height, width, channel]")

    h, w = images.shape[1:3]
    nbatch = images.shape[0]
    channels = images.shape[3]
    rows = ceildiv(nbatch, nrow)
    cols = min(nbatch, nrow)

    if channels != len(fill_value):
        raise ValueError("fill_value must match number of channels")

    out = np.full((h * rows, w * cols, channels), fill_value, dtype=images.dtype)

    for i in range(images.shape[0]):
        row = i // cols
        col = i % cols
        out[h * row : h * (row + 1), w * col : w * (col + 1), :] = images[i]

    return out


def limit_balanced(arr: np.ndarray, total: int) -> np.ndarray:
    """Use a list of counts `arr`, return a new list where the total count is `total` (if possible),
    in such a way that the selection is as balanced as possible.
    """

    assert len(arr.shape) == 1

    indices = np.argsort(arr)
    values = arr[indices]
    remaining_total = total
    remaining_cats = len(arr)
    balanced: list[int] = []
    for value in values:
        n_chose = remaining_total // remaining_cats
        if value >= n_chose:
            balanced.append(n_chose)
            remaining_total -= n_chose
        else:
            balanced.append(value)
            remaining_total -= value
        remaining_cats -= 1

    out = np.array(balanced)
    inverse_indices = np.argsort(indices)
    restored_out = out[inverse_indices]

    assert arr.shape == out.shape
    assert np.all(restored_out <= arr)
    assert np.sum(restored_out) == min(total, np.sum(arr))
    return restored_out


def indind(absolute: Union[slice, np.ndarray], relative: np.ndarray) -> np.ndarray:
    if isinstance(absolute, slice):
        assert absolute == slice(None)
        return relative
    else:
        return absolute[relative]


def _arg_sample_balanced(
    arr: np.ndarray,
    indices: Union[slice, np.ndarray],
    total: int,
    dim: int,
    rng: np.random.Generator,
) -> np.ndarray:
    subarray = arr[indices, dim]
    values, counts = np.unique(subarray, return_counts=True)
    n_chose_arr = limit_balanced(counts, total)

    out = []

    if dim == arr.shape[1] - 1:
        # last column, actually select values
        for v, subtotal in zip(values, n_chose_arr):
            (ind,) = np.nonzero(subarray == v)
            if len(ind) != subtotal:
                ind = rng.choice(ind, subtotal, replace=False)
            out.append(indind(indices, ind))
    else:
        # recurse into other columns
        for v, subtotal in zip(values, n_chose_arr):
            (ind,) = np.nonzero(subarray == v)
            out.append(_arg_sample_balanced(arr, indind(indices, ind), subtotal, dim + 1, rng))

    return np.concatenate(out)


def arg_sample_balanced(arr: np.ndarray, total: int, seed: int = 0) -> np.ndarray:
    """Select a `total` number of samples from `arr` and return the indices.
    `arr` needs to be a 2-dimensional array where the first dimension are the samples
    and the second one categorical columns. These categories are used to sample in a way,
    so that the result is a balanced as possible, with using additional columns with decreasing priority.
    The output are the indices used to select the samples. They are in no specific order.
    To select without reordering, sort the output.
    """

    if len(arr.shape) != 2 or arr.shape[0] < 1 or arr.shape[1] < 1:
        raise ValueError("array must be 2-dimensional with at least one element in each dimension")

    dim = 0
    rng = np.random.default_rng(seed=seed)

    out = _arg_sample_balanced(arr, slice(None), total, dim, rng)
    assert len(out) == min(total, len(arr))
    assert len(np.unique(out)) == len(out)
    return out
