import functools
import multiprocessing
import threading
import types
import warnings
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Dict,
    List,
    Literal,
    Mapping,
    NamedTuple,
    Optional,
    Sequence,
    Tuple,
    Type,
    Union,
)

import joblib as jl
import wrapt

import numpy as np
import pandas as pd
import scipy.sparse as sp
import scipy.stats as st
from statsmodels.stats.multitest import multipletests

from anndata import AnnData

from aligndg.uopt._logging import logger
from aligndg.uopt._types import ArrayLike, Str_Dict_t

if TYPE_CHECKING:
    from uopt.base.problems.compound_problem import BaseCompoundProblem
    from uopt.base.problems.problem import BaseProblem


Callback = Callable[..., Any]


class TimeScalesHeatKernel(NamedTuple):  # noqa: D101
    x: Optional[float]
    y: Optional[float]
    xy: Optional[float]


def _validate_annotations(
    df_source: pd.DataFrame,
    df_target: pd.DataFrame,
    source_annotation_key: Optional[str] = None,
    target_annotation_key: Optional[str] = None,
    source_annotations: Optional[List[Any]] = None,
    target_annotations: Optional[List[Any]] = None,
    aggregation_mode: Literal["annotation", "cell"] = "annotation",
    forward: bool = False,
) -> Tuple[List[Any], List[Any]]:
    if forward:
        if TYPE_CHECKING:  # checked in _check_argument_compatibility_cell_transition(
            assert target_annotations is not None
        target_annotations_verified = list(
            set(df_target[target_annotation_key].cat.categories).intersection(target_annotations)
        )
        if not len(target_annotations_verified):
            raise ValueError(f"None of `{target_annotations}`, found in the target annotations.")
        source_annotations_verified = _validate_annotations_helper(
            df_source, source_annotation_key, source_annotations, aggregation_mode
        )
        return source_annotations_verified, target_annotations_verified

    if TYPE_CHECKING:  # checked in _check_argument_compatibility_cell_transition(
        assert source_annotations is not None
    source_annotations_verified = list(
        set(df_source[source_annotation_key].cat.categories).intersection(set(source_annotations))
    )
    if not len(source_annotations_verified):
        raise ValueError(f"None of `{source_annotations}`, found in the source annotations.")
    target_annotations_verified = _validate_annotations_helper(
        df_target, target_annotation_key, target_annotations, aggregation_mode
    )
    return source_annotations_verified, target_annotations_verified


def _validate_annotations_helper(
    df: pd.DataFrame,
    annotation_key: Optional[str] = None,
    annotations: Optional[List[Any]] = None,
    aggregation_mode: Literal["annotation", "cell"] = "annotation",
) -> List[Any]:
    if aggregation_mode == "annotation":
        if TYPE_CHECKING:  # checked in _check_argument_compatibility_cell_transition(
            assert annotations is not None
        annotations_verified = set(df[annotation_key].cat.categories).intersection(set(annotations))
        if not len(annotations_verified):
            raise ValueError(f"None of `{annotations}` found in the distribution corresponding to `{annotation_key}`.")
        return list(annotations_verified)
    return [None]


def _check_argument_compatibility_cell_transition(
    source_annotation: Str_Dict_t,
    target_annotation: Str_Dict_t,
    key: Optional[str] = None,
    # TODO(MUCDK): unused variable
    other_key: Optional[str] = None,
    other_adata: Optional[AnnData] = None,
    aggregation_mode: Literal["annotation", "cell"] = "annotation",
    forward: bool = False,
    **_: Any,
) -> None:
    if key is None and other_adata is None:
        raise ValueError("Unable to infer distributions, missing `adata` and `key`.")
    if forward and target_annotation is None:
        raise ValueError("No target annotation provided.")
    if aggregation_mode == "annotation" and (not forward and source_annotation is None):
        raise ValueError("No source annotation provided.")
    if (aggregation_mode == "annotation") and (source_annotation is None or target_annotation is None):
        raise ValueError(
            "If aggregation mode is `'annotation'`, " "source or target annotation in `adata.obs` must be provided."
        )


def _get_df_cell_transition(
    adata: AnnData,
    annotation_keys: List[Optional[str]],
    filter_key: Optional[str] = None,
    filter_value: Optional[Any] = None,
) -> pd.DataFrame:
    if filter_key is not None:
        adata = adata[adata.obs[filter_key] == filter_value]
    return adata.obs[list({ak for ak in annotation_keys if ak is not None})].copy()


def _validate_args_cell_transition(
    adata: AnnData,
    arg: Str_Dict_t,
) -> Tuple[Optional[str], List[Any], Optional[List[str]]]:
    if arg is None:
        return None, [], None
    if isinstance(arg, str):
        try:
            return arg, adata.obs[arg].cat.categories, None
        except KeyError:
            raise KeyError(f"Unable to fetch data from `adata.obs[{arg!r}]`.") from None
        except AttributeError:
            raise AttributeError(f"Data in `adata.obs[{arg!r}]` is not categorical.") from None

    if isinstance(arg, dict):
        if len(arg) != 1:
            raise ValueError(f"Expected dictionary of length `1`, found `{len(arg)}`.")
        key, val = next(iter(arg.items()))
        if not set(val).issubset(adata.obs[key].cat.categories):
            raise ValueError(f"Not all values `{val}` are present in `adata.obs[{key!r}]`.")
        return key, val, val

    raise TypeError(f"Expected argument to be either `str` or `dict`, found `{type(arg)}`.")


def _assert_series_match(a: pd.Series, b: pd.Series) -> None:
    """Assert that two series are equal ignoring the names."""
    pd.testing.assert_series_equal(a, b, check_names=False)


def _assert_columns_and_index_match(a: pd.Series, b: pd.DataFrame) -> None:
    """Assert that a series and a dataframe's index and columns are matching."""
    _assert_series_match(a, b.index.to_series())
    _assert_series_match(a, b.columns.to_series())


def _get_cell_indices(
    adata: AnnData,
    key: Optional[str] = None,
    key_value: Optional[Any] = None,
) -> pd.Index:
    if key is None:
        return adata.obs.index
    return adata[adata.obs[key] == key_value].obs.index


def _get_categories_from_adata(
    adata: AnnData,
    key: Optional[str] = None,
    key_value: Optional[Any] = None,
    annotation_key: Optional[str] = None,
) -> pd.Series:
    if key is None:
        return adata.obs[annotation_key]
    return adata[adata.obs[key] == key_value].obs[annotation_key]


def _order_transition_matrix_helper(
    tm: pd.DataFrame,
    rows_verified: List[str],
    cols_verified: List[str],
    row_order: Optional[List[str]],
    col_order: Optional[List[str]],
) -> pd.DataFrame:
    if col_order is not None:
        tm = tm[[col for col in col_order if col in cols_verified]]
    tm = tm.T
    if row_order is not None:
        return tm[[row for row in row_order if row in rows_verified]]
    return tm


def _order_transition_matrix(
    tm: pd.DataFrame,
    source_annotations_verified: List[str],
    target_annotations_verified: List[str],
    source_annotations_ordered: Optional[List[str]],
    target_annotations_ordered: Optional[List[str]],
    forward: bool,
) -> pd.DataFrame:
    # TODO(michalk8): simplify
    if target_annotations_ordered is not None or source_annotations_ordered is not None:
        if forward:
            return _order_transition_matrix_helper(
                tm=tm,
                rows_verified=source_annotations_verified,
                cols_verified=target_annotations_verified,
                row_order=source_annotations_ordered,
                col_order=target_annotations_ordered,
            ).T
        return _order_transition_matrix_helper(
            tm=tm,
            rows_verified=target_annotations_verified,
            cols_verified=source_annotations_verified,
            row_order=target_annotations_ordered,
            col_order=source_annotations_ordered,
        )

    if target_annotations_verified == source_annotations_verified:
        annotations_ordered = tm.columns.sort_values()
        if forward:
            return _order_transition_matrix_helper(
                tm=tm,
                rows_verified=source_annotations_verified,
                cols_verified=target_annotations_verified,
                row_order=annotations_ordered,
                col_order=annotations_ordered,
            ).T
        return _order_transition_matrix_helper(
            tm=tm,
            rows_verified=target_annotations_verified,
            cols_verified=source_annotations_verified,
            row_order=annotations_ordered,
            col_order=annotations_ordered,
        )

    return tm if forward else tm.T


def _correlation_test(
    X: Union[ArrayLike, sp.spmatrix],
    Y: pd.DataFrame,
    feature_names: Sequence[str],
    corr_method: Literal["pearson", "spearman"] = "pearson",
    significance_method: Literal["fisher", "perm_test"] = "fisher",
    confidence_level: float = 0.95,
    n_perms: Optional[int] = None,
    seed: Optional[int] = None,
    **kwargs: Any,
) -> pd.DataFrame:
    """
    Perform a statistical test for correlation between X and .

    Return NaN for genes which don't vary across cells.

    Parameters
    ----------
    X
        Array or sparse matrix of shape ``(n_cells, n_features)`` containing the counts.
    Y
        Data frame of shape ``(n_cells, 1)`` containing the pull/push distribution.
    feature_names
        Sequence of shape ``(n_features,)`` containing the feature names.
    corr_method
        Which type of correlation to compute, options are `pearson`, and `spearman`.
    significance_method
        Method for p-value calculation.
    confidence_level
        Confidence level for the confidence interval calculation. Must be in `[0, 1]`.
    n_perms
        Number of permutations if ``method = 'perm_test'``.
    seed
        Random seed if ``method = 'perm_test'``.
    kwargs
        Keyword arguments for parallelization, e.g., `n_jobs`.

    Returns
    -------
    Dataframe of shape ``(n_genes, 5)`` containing the following columns, one for each lineage:
        - ``corr`` - correlation between the count data and push/pull distributions.
        - ``pval`` - calculated p-values for double-sided test.
        - ``qval`` - corrected p-values using Benjamini-Hochberg method at level `0.05`.
        - ``ci_low`` - lower bound of the ``confidence_level`` correlation confidence interval.
        - ``ci_high`` - upper bound of the ``confidence_level`` correlation confidence interval.
    """
    corr, pvals, ci_low, ci_high = _correlation_test_helper(
        X.T,
        Y.values,
        corr_method=corr_method,
        significance_method=significance_method,
        n_perms=n_perms,
        seed=seed,
        confidence_level=confidence_level,
        **kwargs,
    )
    invalid = (corr < -1) | (corr > 1)
    if np.any(invalid):
        logger.warning(
            f"Found `{np.sum(invalid)}` correlation(s) that are not in `[0, 1]`. "
            f"This usually happens when gene expression is constant across all cells. "
            f"Setting to `NaN`"
        )
        corr[invalid] = np.nan
        pvals[invalid] = np.nan
        ci_low[invalid] = np.nan
        ci_high[invalid] = np.nan

    res = pd.DataFrame(corr, index=feature_names, columns=[f"{c}_corr" for c in Y.columns])
    for idx, c in enumerate(Y.columns):
        p = pvals[:, idx]
        valid_mask = ~np.isnan(p)

        res[f"{c}_pval"] = p
        res[f"{c}_qval"] = np.nan
        if np.any(valid_mask):
            res.loc[np.asarray(feature_names)[valid_mask], f"{c}_qval"] = multipletests(
                p[valid_mask], alpha=0.05, method="fdr_bh"
            )[1]
        res[f"{c}_ci_low"] = ci_low[:, idx]
        res[f"{c}_ci_high"] = ci_high[:, idx]

    # fmt: off
    res = res[[f"{c}_{stat}" for c in Y.columns for stat in ("corr", "pval", "qval", "ci_low", "ci_high")]]
    return res.sort_values(by=[f"{c}_corr" for c in Y.columns], ascending=False)
    # fmt: on


def _correlation_test_helper(
    X: ArrayLike,
    Y: ArrayLike,
    corr_method: Literal["pearson", "spearman"] = "spearman",
    significance_method: Literal["fisher", "perm_test"] = "fisher",
    n_perms: Optional[int] = None,
    seed: Optional[int] = None,
    confidence_level: float = 0.95,
    **kwargs: Any,
) -> Tuple[ArrayLike, ArrayLike, ArrayLike, ArrayLike]:
    """
    Compute the correlation between rows in matrix ``X`` columns of matrix ``Y``.

    Parameters
    ----------
    X
        Array or matrix of `(M, N)` elements.
    Y
        Array of `(N, K)` elements.
    corr_method
        Which type of correlation to compute, options are `pearson`, and `spearman`.
    significance_method
        Method for p-value calculation.
    n_perms
        Number of permutations if ``method='perm_test'``.
    seed
        Random seed if ``method = 'perm_test'``.
    confidence_level
        Confidence level for the confidence interval calculation. Must be in `[0, 1]`.
    kwargs
        Keyword arguments for :func:`moscot._utils.parallelize`, e.g. `n_jobs`.

    Returns
    -------
    Correlations, p-values, corrected p-values, lower and upper bound of 95% confidence interval.
    """

    def perm_test_extractor(res: Sequence[Tuple[ArrayLike, ArrayLike]]) -> Tuple[ArrayLike, ArrayLike, ArrayLike]:
        pvals, corr_bs = zip(*res)
        if TYPE_CHECKING:
            assert isinstance(n_perms, int)
        pvals = np.sum(pvals, axis=0) / float(n_perms)

        corr_bs = np.concatenate(corr_bs, axis=0)
        corr_ci_low, corr_ci_high = np.quantile(corr_bs, q=ql, axis=0), np.quantile(corr_bs, q=qh, axis=0)

        return pvals, corr_ci_low, corr_ci_high  # type:ignore[return-value]

    if not (0 <= confidence_level <= 1):
        raise ValueError(f"Expected `confidence_level` to be in interval `[0, 1]`, found `{confidence_level}`.")

    n = X.shape[1]  # genes x cells
    ql = 1 - confidence_level - (1 - confidence_level) / 2.0
    qh = confidence_level + (1 - confidence_level) / 2.0

    if sp.issparse(X):
        X = sp.csr_matrix(X)

    if corr_method == "spearman":
        X, Y = st.rankdata(X, method="average", axis=0), st.rankdata(Y, method="average", axis=0)
    corr = _pearson_mat_mat_corr_sparse(X, Y) if sp.issparse(X) else _pearson_mat_mat_corr_dense(X, Y)

    if significance_method == "fisher":
        # see: https://en.wikipedia.org/wiki/Pearson_correlation_coefficient#Using_the_Fisher_transformation
        # for spearman see: https://www.sciencedirect.com/topics/mathematics/spearman-correlation
        mean, se = np.arctanh(corr), 1 / np.sqrt(n - 3)
        z_score = (np.arctanh(corr) - np.arctanh(0)) * np.sqrt(n - 3)

        z = st.norm.ppf(qh)
        corr_ci_low = np.tanh(mean - z * se)
        corr_ci_high = np.tanh(mean + z * se)
        pvals = 2 * st.norm.cdf(-np.abs(z_score))

    elif significance_method == "perm_test":
        if not isinstance(n_perms, int):
            raise TypeError(f"Expected `n_perms` to be an integer, found `{type(n_perms).__name__}`.")
        if n_perms <= 0:
            raise ValueError(f"Expected `n_perms` to be positive, found `{n_perms}`.")

        pvals, corr_ci_low, corr_ci_high = parallelize(
            _perm_test,  # type: ignore[arg-type]
            np.arange(n_perms),
            as_array=False,
            unit="permutation",
            extractor=perm_test_extractor,
            **kwargs,
        )(corr, X, Y, seed=seed)

    else:
        raise NotImplementedError(significance_method)

    return corr, pvals, corr_ci_low, corr_ci_high


def _pearson_mat_mat_corr_sparse(
    X: sp.csr_matrix,
    Y: ArrayLike,
) -> ArrayLike:
    n = X.shape[1]
    X_bar = np.reshape(np.array(X.mean(axis=1)), (-1, 1))
    X_std = np.reshape(np.sqrt(np.array(X.power(2).mean(axis=1)) - (X_bar**2)), (-1, 1))

    y_bar = np.reshape(np.mean(Y, axis=0), (1, -1))
    y_std = np.reshape(np.std(Y, axis=0), (1, -1))

    with warnings.catch_warnings():
        warnings.simplefilter("ignore", RuntimeWarning)
        return (X @ Y - (n * X_bar * y_bar)) / ((n - 1) * X_std * y_std)


def _pearson_mat_mat_corr_dense(X: ArrayLike, Y: ArrayLike) -> ArrayLike:
    n = X.shape[1]

    X_bar = np.reshape(np.mean(X, axis=1), (-1, 1))
    X_std = np.reshape(np.std(X, axis=1), (-1, 1))

    y_bar = np.reshape(np.mean(Y, axis=0), (1, -1))
    y_std = np.reshape(np.std(Y, axis=0), (1, -1))

    with warnings.catch_warnings():
        warnings.simplefilter("ignore", RuntimeWarning)
        return (X @ Y - (n * X_bar * y_bar)) / ((n - 1) * X_std * y_std)


def _perm_test(
    ixs: ArrayLike,
    corr: ArrayLike,
    X: Union[ArrayLike, sp.spmatrix],
    Y: ArrayLike,
    seed: Optional[int] = None,
    queue=None,
) -> Tuple[ArrayLike, ArrayLike]:
    rs = np.random.RandomState(None if seed is None else seed + ixs[0])
    cell_ixs = np.arange(X.shape[1])
    pvals = np.zeros_like(corr, dtype=np.float64)
    corr_bs = np.zeros((len(ixs), X.shape[0], Y.shape[1]))  # perms x genes x lineages

    mmc = _pearson_mat_mat_corr_sparse if sp.issparse(X) else _pearson_mat_mat_corr_dense

    for i, _ in enumerate(ixs):
        rs.shuffle(cell_ixs)
        corr_i = mmc(X, Y[cell_ixs, :])
        pvals += np.abs(corr_i) >= np.abs(corr)

        bootstrap_ixs = rs.choice(cell_ixs, replace=True, size=len(cell_ixs))
        corr_bs[i, :, :] = mmc(X[:, bootstrap_ixs], Y[bootstrap_ixs, :])

        if queue is not None:
            queue.put(1)

    if queue is not None:
        queue.put(None)

    return pvals, corr_bs


@wrapt.decorator
def require_solution(
    wrapped: Callable[[Any], Any], instance: "BaseProblem", args: Tuple[Any, ...], kwargs: Mapping[str, Any]
) -> Any:
    """Check whether problem has been solved."""
    from uopt.base.problems.compound_problem import BaseCompoundProblem
    from uopt.base.problems.problem import OTProblem

    if isinstance(instance, OTProblem) and instance.solution is None:
        raise RuntimeError("Run `.solve()` first.")
    if isinstance(instance, BaseCompoundProblem) and instance.solutions is None:
        raise RuntimeError("Run `.solve()` first.")
    return wrapped(*args, **kwargs)


@wrapt.decorator
def require_prepare(
    wrapped: Callable[[Any], Any],
    instance: "BaseCompoundProblem",  # type: ignore[type-arg]
    args: Tuple[Any, ...],
    kwargs: Mapping[str, Any],
) -> Any:
    """Check whether problem has been prepared."""
    if instance.problems is None:
        raise RuntimeError("Run `.prepare()` first.")
    return wrapped(*args, **kwargs)


@wrapt.decorator
def wrap_prepare(
    wrapped: Callable[[Any], Any], instance: "BaseProblem", args: Tuple[Any, ...], kwargs: Mapping[str, Any]
) -> Any:
    """Check and update the state when preparing :class:`moscot.problems.base.OTProblem`."""
    instance = wrapped(*args, **kwargs)
    if instance.problem_kind == "unknown":
        raise RuntimeError("Problem kind was not set after running `.prepare()`.")
    instance._stage = "prepared"
    return instance


@wrapt.decorator
def wrap_solve(
    wrapped: Callable[[Any], Any], instance: "BaseProblem", args: Tuple[Any, ...], kwargs: Mapping[str, Any]
) -> Any:
    """Check and update the state when solving :class:`moscot.problems.base.OTProblem`."""
    if instance.stage not in ("prepared", "solved"):
        raise RuntimeError(
            f"Expected problem's stage to be either `'prepared'` or `'solved'`, found `{instance.stage!r}`."
        )
    instance = wrapped(*args, **kwargs)
    instance._stage = "solved"
    return instance


def attributedispatch(func: Optional[Callback] = None, attr: Optional[str] = None) -> Callback:
    """Dispatch a function based on the first value."""

    def dispatch(value: Type[Any]) -> Callback:
        for typ in value.mro():
            if typ in registry:
                return registry[typ]
        return func  # type: ignore[return-value]

    def register(value: Type[Any], func: Optional[Callback] = None) -> Callback:
        if func is None:
            return lambda f: register(value, f)
        registry[value] = func
        return func

    def wrapper(instance: Any, *args: Any, **kwargs: Any) -> Any:
        typ = type(getattr(instance, str(attr)))
        return dispatch(typ)(instance, *args, **kwargs)

    if func is None:
        return functools.partial(attributedispatch, attr=attr)

    registry: Dict[Type[Any], Callback] = {}
    wrapper.register = register  # type: ignore[attr-defined]
    wrapper.dispatch = dispatch  # type: ignore[attr-defined]
    wrapper.registry = types.MappingProxyType(registry)  # type: ignore[attr-defined]
    functools.update_wrapper(wrapper, func)

    return wrapper


def parallelize(
    callback: Callable[[Any], Any],
    collection: Union[sp.spmatrix, Sequence[Any]],
    n_jobs: Optional[int] = None,
    n_split: Optional[int] = None,
    unit: str = "",
    as_array: bool = True,
    use_ixs: bool = False,
    backend: str = "loky",
    extractor: Optional[Callable[[Any], Any]] = None,
    show_progress_bar: bool = True,
) -> Any:
    """
    Parallelize function call over a collection of elements.

    Parameters
    ----------
    callback
        Function to parallelize.
    collection
        Sequence of items which to chunkify or an already.
    n_jobs
        Number of parallel jobs.
    n_split
        Split ``collection`` into ``n_split`` chunks. If `None`, split into ``n_jobs`` chunks.
    unit
        Unit of the progress bar.
    as_array
        Whether to convert the results not :class:`numpy.ndarray`.
    use_ixs
        Whether to pass indices to the callback.
    backend
        Which backend to use for multiprocessing. See :class:`joblib.Parallel` for valid options.
    extractor
        Function to apply to the result after all jobs have finished.
    show_progress_bar
        Whether to show a progress bar.

    Returns
    -------
    The result depending on ``callable``, ``extractor`` and ``as_array``.
    """
    if show_progress_bar:
        try:
            from tqdm.auto import tqdm
        except ImportError:
            try:
                from tqdm.std import tqdm
            except ImportError:
                tqdm = None
    else:
        tqdm = None

    def update(pbar, queue, n_total):
        n_finished = 0
        while n_finished < n_total:
            try:
                res = queue.get()
            except EOFError as e:
                if not n_finished != n_total:
                    raise RuntimeError(f"Finished only `{n_finished}` out of `{n_total}` tasks.`") from e
                break
            assert res in (None, (1, None), 1)  # (None, 1) means only 1 job
            if res == (1, None):
                n_finished += 1
                if pbar is not None:
                    pbar.update()
            elif res is None:
                n_finished += 1
            elif pbar is not None:
                pbar.update()

        if pbar is not None:
            pbar.close()

    def wrapper(*args, **kwargs):
        if pass_queue and show_progress_bar:
            pbar = None if tqdm is None else tqdm(total=col_len, unit=unit, mininterval=0.125)
            queue = multiprocessing.Manager().Queue()
            thread = threading.Thread(target=update, args=(pbar, queue, len(collections)))
            thread.start()
        else:
            pbar, queue, thread = None, None, None

        res = jl.Parallel(n_jobs=n_jobs, backend=backend)(
            jl.delayed(callback)(
                *((i, cs) if use_ixs else (cs,)),
                *args,
                **kwargs,
                queue=queue,
            )
            for i, cs in enumerate(collections)
        )

        res = np.array(res) if as_array else res
        if thread is not None:
            thread.join()

        return res if extractor is None else extractor(res)

    col_len = collection.shape[0] if sp.issparse(collection) else len(collection)  # type: ignore[union-attr]
    n_jobs = _get_n_cores(n_jobs, col_len)
    if n_split is None:
        n_split = n_jobs

    if sp.issparse(collection):
        n_split = max(1, min(n_split, collection.shape[0]))  # type: ignore
        if n_split == collection.shape[0]:  # type: ignore[union-attr]
            collections = [collection[[ix], :] for ix in range(collection.shape[0])]  # type: ignore
        else:
            step = collection.shape[0] // n_split  # type: ignore[union-attr]
            ixs = [np.arange(i * step, min((i + 1) * step, collection.shape[0])) for i in range(n_split)]  # type: ignore  # noqa: E501
            ixs[-1] = np.append(ixs[-1], np.arange(ixs[-1][-1] + 1, collection.shape[0]))  # type: ignore

            collections = [collection[ix, :] for ix in filter(len, ixs)]  # type:ignore[call-overload]
    else:
        collections = list(filter(len, np.array_split(collection, n_split)))

    n_split = len(collections)
    n_jobs = min(n_jobs, n_split)
    pass_queue = not hasattr(callback, "py_func")  # we'd be inside a numba function

    return wrapper


def _get_n_cores(n_cores: Optional[int], n_jobs: Optional[int]) -> int:
    """
    Make number of cores a positive integer.

    Parameters
    ----------
    n_cores
        Number of cores to use.
    n_jobs
        Number of jobs. This is just used to determine if the collection is a singleton.
        If `1`, always returns `1`.

    Returns
    -------
    Positive integer corresponding to how many cores to use.
    """
    if n_cores == 0:
        raise ValueError("Number of cores cannot be `0`.")
    if n_jobs == 1 or n_cores is None:
        return 1
    if n_cores < 0:
        return multiprocessing.cpu_count() + 1 + n_cores

    return n_cores
