import enum
from typing import Any, Optional, Sequence, Union

import numba as nb
import numpy as np
import scipy.sparse as sp
from numba import prange
from sklearn.utils.sparsefuncs import csc_median_axis_0

from anndata import AnnData

import cellrank.logging as logg
from cellrank._utils._docs import d, inject_docs
from cellrank._utils._enum import ModeEnum
from cellrank._utils._parallelize import parallelize
from cellrank._utils._utils import valuedispatch

__all__ = ["_get_offset"]

_OFFSET_KEY = "cellrank_offset"


class NormMode(ModeEnum):  # noqa: D101
    TMM = enum.auto()
    RLE = enum.auto()
    UPPER_QUANT = enum.auto()
    NONE = enum.auto()


@d.dedent
def _extract_data(data: AnnData, layer: Optional[str] = None, use_raw: bool = True) -> Union[np.ndarray, sp.spmatrix]:
    """Extract expression data from an object.

    Parameters
    ----------
    data
        Annotated data object or an array.
    layer
        Key in :attr:`~anndata.AnnData.layers` accessed when ``use_raw = False``.
        Only used when ``data`` is :class:`~anndata.AnnData`.
    use_raw
        Whether to access :attr:`anndata.AnnData.raw`. Only used when ``data`` is :class:`~anndata.AnnData`.

    Returns
    -------
    The extracted expression data.
    """
    if isinstance(data, AnnData):
        if use_raw:
            if not hasattr(data, "raw"):
                raise AttributeError("No `.raw` attribute found.")
            if data.raw is None:
                raise ValueError("Attribute `.raw` is None.")
            x = data.raw.X
        elif layer is not None:
            if layer not in data.layers:
                raise KeyError(
                    f"Layer `{layer!s}` not found in `adata.layers`. "
                    f"Valid options are: `{list(data.layers.keys())}`."
                )
            x = data.layers[layer]
        else:
            x = data.X
    elif not isinstance(data, (np.ndarray, sp.spmatrix)):
        raise TypeError(
            f"Expected parameter `data` to be either `anndata.AnnData`, `numpy.ndarray` "
            f"or `scipy.sparse.spmatrix`, found `{type(data).__name__!r}`."
        )
    else:
        x = data

    return x


@nb.njit(fastmath=True, cache=True)
def _rankdata(a: np.ndarray, method: str = "average") -> np.ndarray:
    assert method in (
        "average",
        "min",
        "max",
        "dense",
        "ordinal",
    ), "Invalid ranking method."

    arr = np.ravel(np.asarray(a))
    # algo = 'mergesort' if method == 'ordinal' else 'quicksort'
    sorter = np.argsort(arr)  # , kind=algo)

    inv = np.empty(sorter.size, dtype=np.intp)
    inv[sorter] = np.arange(sorter.size, dtype=np.intp)

    if method == "ordinal":
        return inv + 1.0

    arr = arr[sorter]
    obs = np.ones(len(arr), dtype=np.bool_)
    obs[1:] = arr[1:] != arr[:-1]
    dense = obs.cumsum()[inv]

    if method == "dense":
        return dense * 1.0

    # cumulative counts of each unique value
    count = np.append(np.nonzero(obs)[0], len(obs))

    if method == "max":
        return count[dense] * 1.0

    if method == "min":
        return count[dense - 1] + 1.0

    # average method
    return 0.5 * (count[dense] + count[dense - 1] + 1.0)


@inject_docs(m=NormMode)
def _calculate_norm_factors(
    data: Union[AnnData, np.ndarray, sp.spmatrix],
    method: str = NormMode.TMM,
    layer: Optional[str] = None,
    use_raw: bool = True,
    library_size: Optional[np.ndarray] = None,
    ref_ix: Optional[int] = None,
    logratio_trim: float = 0.3,
    sum_trim: float = 0.05,
    weight: bool = True,
    a_cutoff: float = -1e10,
    perc: float = 0.75,
) -> np.ndarray:
    """Calculate normalization factors.

    This uses the `edgeR <https://www.rdocumentation.org/packages/edgeR/versions/3.14.0/topics/calcNormFactors>`_
    implementation.

    Parameters
    ----------
    data
        Data of shape ``(n_cells, n_genes)`` containing, e.g., the counts.
    method
        Which method to use. Valid options are:

        - ``{m.TMM!r}`` - weighted trimmed mean of M-values from :cite:`robinson:10`.
        - ``{m.RLE!r}`` - relative log expression from [Anders10]_.
        - ``{m.UPPER_QUANT!r}`` - upper-quartile normalization method from [Bullard10]_.
        - ``{m.NONE!r}`` - all the factors are set to :math:`1`.
    layer
        Layer in :attr:`anndata.AnnData.layers` or `None` for :attr:`anndata.AnnData.X`.
        Only used when ``use_raw = False``.
    use_raw
        Whether to access :attr:`~anndata.AnnData.raw` or not.
    library_size
        Library size. If :obj:`None`, it will be set as the sum of values for each cell.
    ref_ix
        Index of a reference cell. If :obj:`None`, it will be determined automatically.
    logratio_trim
        Amount of trim to use on log-ratios ('M' values) when ``method = {m.TMM!r}``.
    sum_trim
        Amount of trim to use on the combined absolute levels ('A' values) when ``method = {m.TMM!r}``.
    weight
        Whether to compute asymptotic binomial precision weights when ``method = {m.TMM!r}``.
    a_cutoff
        Cutoff on 'A' values to use before trimming when ``method = {m.TMM!r}``.
    perc
        Percentile of the counts that is aligned when ``method = {m.UPPER_QUANT!r}``.

    Returns
    -------
    Array of shape ``(data.shape[0],)`` containing the factors.

    References
    ----------
    .. [Anders10] Anders, S. *et al.* (2010),
        *Differential expression analysis for sequence count data*,
        `Genome Biology <https://doi.org/10.1186/gb-2010-11-10-r106>`__.
    .. [Bullard10] Bullard, J. H. *et al.* (2010),
        *Evaluation of statistical methods for normalization and differential expression in mRNA-Seq experiments*,
        `BMC Bioinformatics <https://doi.org/10.1186/1471-2105-11-94>`__.
    """
    method = NormMode(method)

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

    x = _extract_data(data, layer, use_raw)

    if library_size is None:
        library_size = np.array(x.sum(1)).squeeze()
    elif library_size.shape != (x.shape[0],):
        raise ValueError(f"Expected `library_size` to be of shape `({x.shape[0]},)`, found `{library_size.shape}`.")

    f = _dispatch_computation(
        method,
        x=x,
        library_size=library_size,
        ref_ix=ref_ix,
        logratio_trim=logratio_trim,
        sum_trim=sum_trim,
        weight=weight,
        a_cutoff=a_cutoff,
        p=perc,
    )

    return f / np.exp(np.mean(np.log(f)))


@valuedispatch
def _dispatch_computation(mode, *_args, **_kwargs):
    raise NotImplementedError(mode)


@_dispatch_computation.register(NormMode.TMM)
def _(
    x: Union[np.ndarray, sp.spmatrix],
    library_size: np.ndarray,
    ref_ix: Optional[int] = None,
    **kwargs,
) -> np.ndarray:
    if ref_ix is None:
        f75 = _calc_factor_quant(x, library_size=library_size, p=0.75)
        ref_ix = np.argmin(np.abs(f75 - np.mean(f75)))

    ref = x[ref_ix]
    if sp.issparse(ref):
        ref = ref.A.squeeze(0)  # (genes,)

    return parallelize(
        _calc_factor_weighted,
        collection=np.arange(x.shape[0]),
        show_progress_bar=False,
        as_array=False,
        extractor=np.concatenate,
        backend="threading",
        n_jobs=4,
    )(
        obs_=x,
        obs_lib_size_=library_size,
        ref=ref,
        ref_lib_size=library_size[ref_ix],
        **kwargs,
    )


@_dispatch_computation.register(NormMode.UPPER_QUANT)
def _calc_factor_quant(x: Union[np.ndarray, sp.spmatrix], library_size: np.ndarray, p: float, **_) -> np.ndarray:
    # unused, because we fix the reference cell to 0

    library_size = np.array(library_size).reshape((-1, 1))
    if not sp.issparse(x):
        # this is same as below, which is R's default
        # return mquantiles(x / library_size, prob=q, alphap=1, betap=1, axis=1).data.squeeze()
        return np.quantile(x / library_size, p, axis=1)

    y = x.multiply(1.0 / library_size).tocsr()
    return _calc_factor_quant_sparse_helper(y.data, y.indptr, p=p, n_cols=y.shape[1])


@nb.njit(parallel=True, fastmath=True, cache=True)
def _calc_factor_quant_sparse_helper(
    data: np.ndarray,
    indptr: np.ndarray,
    p: float,
    n_cols: int,
) -> np.ndarray:
    n = len(indptr) - 1
    out = np.empty(n)

    # maybe disable parallel
    for ix in prange(n):
        d = data[indptr[ix] : indptr[ix + 1]]
        tmp = np.zeros(shape=(n_cols,))
        tmp[: len(d)] = d
        out[ix] = np.quantile(tmp, p)

    return out


@_dispatch_computation.register(NormMode.RLE)
def _(x: Union[np.ndarray, sp.spmatrix], **_) -> np.ndarray:
    # unused

    # log 0 -> inf -> masked out anyway, so we select only genes in every cell
    # this is the exact replica an in edgeR
    mask = np.array((x > 0).sum(0)).squeeze() == x.shape[0]
    tmp = x[:, mask]
    if sp.issparse(tmp):
        tmp = tmp.A

    gm = np.array(np.exp(np.mean(np.log(tmp), axis=0))).squeeze()
    gm_mask = (gm > 0) & np.isfinite(gm)

    if not sp.issparse(x):
        return np.median(x[:, mask][:, gm_mask] / gm[gm_mask], axis=1)

    return csc_median_axis_0(x.tocsr()[:, mask][:, gm_mask].multiply(1.0 / gm[gm_mask]).tocsr().T)


@_dispatch_computation.register(NormMode.NONE)
def _(x: Union[np.ndarray, sp.spmatrix], **_) -> np.ndarray:
    # unused
    return np.ones((x.shape[0],), dtype=x.dtype)


def _calc_factor_weighted(
    ixs: np.ndarray,
    obs_: Union[np.ndarray, sp.spmatrix],
    obs_lib_size_: np.ndarray,
    ref: np.ndarray,
    ref_lib_size: float,
    logratio_trim: float = 0.3,
    sum_trim: float = 0.05,
    weight: bool = True,
    a_cutoff: float = -1e10,
    queue=None,
    **_: Any,
) -> np.ndarray:
    with np.errstate(divide="ignore", invalid="ignore"):
        log_ref = np.log2(ref / ref_lib_size)
    res = np.empty(shape=(len(ixs),))

    for i, ix in enumerate(ixs):
        if queue is not None:
            queue.put(1)

        obs = obs_[ix]
        if sp.issparse(obs):
            obs = obs.A.squeeze(0)  # (genes,)
        obs_lib_size = obs_lib_size_[ix]

        res[i] = _calc_factor_weighted_helper(
            obs=obs,
            obs_lib_size=obs_lib_size,
            ref=ref,
            ref_lib_size=ref_lib_size,
            log_ref=log_ref,
            a_cutoff=a_cutoff,
            logratio_trim=logratio_trim,
            sum_trim=sum_trim,
            weight=weight,
        )

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

    return res


@nb.njit(fastmath=True, cache=True)
def _calc_factor_weighted_helper(
    obs: np.ndarray,
    obs_lib_size: float,
    ref: np.ndarray,
    ref_lib_size: float,
    log_ref: np.ndarray,
    a_cutoff: float,
    logratio_trim: float,
    sum_trim: float,
    weight: bool,
) -> float:
    o_scaled = obs / obs_lib_size
    log_obs = np.log2(o_scaled)

    log_r = log_obs - log_ref
    abs_e = (log_obs + log_ref) / 2.0
    v = ((obs_lib_size - obs) / obs_lib_size / obs) + (ref_lib_size - ref) / ref_lib_size / ref

    mask = np.isfinite(log_r) & np.isfinite(abs_e) & (abs_e > a_cutoff)

    log_r = log_r[mask]
    abs_e = abs_e[mask]
    v = v[mask]

    if not len(log_r) or np.max(np.abs(log_r)) < 1e-6:
        return 1.0

    n = len(log_r)
    lol = np.floor(n * logratio_trim) + 1.0
    hil = n + 1 - lol

    los = np.floor(n * sum_trim) + 1.0
    his = n + 1 - los

    rank_log_r, rank_abs_e = _rankdata(log_r), _rankdata(abs_e)
    mask = (rank_log_r >= lol) & (rank_log_r <= hil) & (rank_abs_e >= los) & (rank_abs_e <= his)

    f = (np.nansum(log_r[mask] / v[mask]) / np.nansum(1.0 / v[mask])) if weight else np.nanmean(log_r[mask])

    return 2 ** f if np.isfinite(f) else 1.0


def _get_knotlocs(
    pseudotime: Union[np.ndarray, Sequence],
    n_knots: int,
    uniform: bool = False,
) -> np.ndarray:
    """Find knot locations.

    The first and last knots are always placed at the beginning and the of the ``pseudotime``.

    Parameters
    ----------
    pseudotime
        Pseudotemporal ordering of cells.
    n_knots
        Number of knots.
    uniform
        Whether to place the knots uniformly across the ``pseudotime`` or place them based on ``pseudotime``'s density.

    Returns
    -------
    Array of shape ``(n_knots,)`` containing the locations of knots along the pseudotime.
    """
    if n_knots <= 0:
        raise ValueError(f"Expected number of knots to be positive, found `{n_knots}`.")

    if np.any(~np.isfinite(pseudotime)):
        raise ValueError("Not all pseudotime values are finite.")

    pseudotime = np.asarray(pseudotime)

    unique = np.unique(pseudotime)
    if len(unique) in (0, 1):
        raise ValueError(f"All pseudotime values are the same: `{unique}`.")

    if pseudotime.ndim == 2 and pseudotime.shape[1] == 1:
        pseudotime = pseudotime.squeeze(1)
    if pseudotime.ndim != 1:
        raise ValueError(f"Expected `pseudotime` to have `1` dimension, found `{pseudotime.ndim}`.")

    if uniform:
        # replicate the result from not uniform
        return (
            np.linspace(np.min(pseudotime), np.max(pseudotime), n_knots, endpoint=True)
            if n_knots > 1
            else np.array([np.max(pseudotime)])
        )

    x = np.quantile(pseudotime, q=np.arange(n_knots, dtype=np.float64) / max(n_knots - 1, 1))
    x[0] = np.min(pseudotime)
    x[-1] = np.max(pseudotime)

    u, ix, c = np.unique(x, return_index=True, return_counts=True)

    if len(u) != len(x):
        knotlocs = []
        for start, end, size in zip(x[ix], x[ix[1:]], c):
            knotlocs.extend(np.linspace(start, end, size, endpoint=False))
        knotlocs.extend(np.linspace(knotlocs[-1], x[ix[-1]], c[-1] + 1, endpoint=True)[1:])
        knotlocs = np.array(knotlocs)
    else:
        knotlocs = x

    logg.debug(f"Setting knot locations to `{list(knotlocs)}`")

    return knotlocs


@inject_docs(key=_OFFSET_KEY)
@d.dedent
def _get_offset(
    adata: Union[AnnData, np.ndarray, sp.spmatrix],
    layer: Optional[str] = None,
    use_raw: bool = True,
    ref_ix: Optional[int] = None,
    recompute: bool = False,
    **kwargs,
) -> np.ndarray:
    """Return an offset for GAM.

    Parameters
    ----------
    %(adata)s
    layer
        Layer in :attr:`~anndata.AnnData.layers`` or :obj:`None` for :attr:`~anndata.AnnData.X`.
        Only used when ``use_raw=False``.
    use_raw
        Whether to access :attr:`~anndata.AnnData.raw`.
    ref_ix
        Index of a reference cell. If :obj:`None`, it will be determined automatically.
    recompute
        Whether to recompute the offset if already found in :attr:`adata.obs['{key}'] <anndata.AnnData.obs>`.
    kwargs
        Keyword arguments for :func:`_utils._calc_norm_factors`.

    Returns
    -------
    Array of shape ``(adata.n_obs,)`` containing the offset.
    """
    with np.errstate(divide="ignore", invalid="ignore"):
        if not recompute and isinstance(adata, AnnData) and _OFFSET_KEY in adata.obs:
            logg.debug(f"Fetching offset from `adata.obs[{_OFFSET_KEY!r}]`")
            return adata.obs[_OFFSET_KEY].values.copy()

        logg.debug(f"Calculating offset for `{adata.shape[0]}` cells")

        data = _extract_data(adata, layer=layer, use_raw=use_raw)
        try:
            nf = _calculate_norm_factors(adata, layer=layer, use_raw=use_raw, ref_ix=ref_ix, **kwargs)
        except Exception as e:  # noqa: BLE001
            logg.debug(f"Unable to calculate the normalization factors, setting them to `1`. Reason: `{e}`")
            nf = np.ones(len(adata), dtype=np.float64)

        offset = np.log(nf * np.array(data.sum(1)).squeeze())
        offset[offset == 0] = 1.0

        mask = ~np.isfinite(offset) | np.isnan(offset)
        if np.any(mask):
            logg.warning(f"`{np.sum(mask)}` elements are not finite. Setting them to `1`")
            offset[mask] = 1.0

        if isinstance(adata, AnnData):
            adata.obs[_OFFSET_KEY] = offset

        return offset
