"""Embed high-dimensional data using PHATE."""

from __future__ import annotations

from typing import TYPE_CHECKING

from ... import logging as logg
from ..._compat import old_positionals
from ..._settings import settings
from ..._utils._doctests import doctest_needs

if TYPE_CHECKING:
    from typing import Literal

    from anndata import AnnData

    from ..._utils.random import _LegacyRandom


@old_positionals(
    "k",
    "a",
    "n_landmark",
    "t",
    "gamma",
    "n_pca",
    "knn_dist",
    "mds_dist",
    "mds",
    "n_jobs",
    "random_state",
    "verbose",
    "copy",
)
@doctest_needs("phate")
def phate(  # noqa: PLR0913
    adata: AnnData,
    n_components: int = 2,
    *,
    k: int = 5,
    a: int = 15,
    n_landmark: int = 2000,
    t: int | str = "auto",
    gamma: float = 1.0,
    n_pca: int = 100,
    knn_dist: str = "euclidean",
    mds_dist: str = "euclidean",
    mds: Literal["classic", "metric", "nonmetric"] = "metric",
    n_jobs: int | None = None,
    random_state: _LegacyRandom = None,
    verbose: bool | int | None = None,
    copy: bool = False,
    **kwargs,
) -> AnnData | None:
    """PHATE :cite:p:`Moon2019`.

    Potential of Heat-diffusion for Affinity-based Trajectory Embedding (PHATE)
    embeds high dimensional single-cell data into two or three dimensions for
    visualization of biological progressions.

    For more information and access to the object-oriented interface, read the
    `PHATE documentation <https://phate.readthedocs.io/>`__.  For
    tutorials, bug reports, and R/MATLAB implementations, visit the `PHATE
    GitHub page <https://github.com/KrishnaswamyLab/PHATE/>`__. For help
    using PHATE, go `here <https://krishnaswamylab.org/get-help>`__.

    Parameters
    ----------
    adata
        Annotated data matrix.
    n_components
        number of dimensions in which the data will be embedded
    k
        number of nearest neighbors on which to build kernel
    a
        sets decay rate of kernel tails.
        If None, alpha decaying kernel is not used
    n_landmark
        number of landmarks to use in fast PHATE
    t
        power to which the diffusion operator is powered
        sets the level of diffusion. If 'auto', t is selected
        according to the knee point in the Von Neumann Entropy of
        the diffusion operator
    gamma
        Informational distance constant between -1 and 1.
        `gamma=1` gives the PHATE log potential, `gamma=0` gives
        a square root potential.
    n_pca
        Number of principal components to use for calculating
        neighborhoods. For extremely large datasets, using
        n_pca < 20 allows neighborhoods to be calculated in
        log(n_samples) time.
    knn_dist
        recommended values: 'euclidean' and 'cosine'
        Any metric from `scipy.spatial.distance` can be used
        distance metric for building kNN graph
    mds_dist
        recommended values: 'euclidean' and 'cosine'
        Any metric from `scipy.spatial.distance` can be used
        distance metric for MDS
    mds
        Selects which MDS algorithm is used for dimensionality reduction.
    n_jobs
        The number of jobs to use for the computation.
        If `None`, `sc.settings.n_jobs` is used.
        If -1 all CPUs are used. If 1 is given, no parallel computing code is
        used at all, which is useful for debugging.
        For n_jobs below -1, (n_cpus + 1 + n_jobs) are used. Thus for
        n_jobs = -2, all CPUs but one are used
    random_state
        Random seed. Defaults to the global `numpy` random number generator
    verbose
        If `True` or an :class:`int`/:class:`~scanpy.Verbosity` ≥ 2/:attr:`~scanpy.Verbosity.hint`, print status messages.
        If `None`, :attr:`scanpy.settings.verbosity` is used.
    copy
        Return a copy instead of writing to `adata`.
    kwargs
        Additional arguments to `phate.PHATE`

    Returns
    -------
    Depending on `copy`, returns or updates `adata` with the following fields.

    **X_phate** : `np.ndarray`, (`adata.obs`, shape=[n_samples, n_components], dtype `float`)
        PHATE coordinates of data.

    Examples
    --------
    >>> from anndata import AnnData
    >>> import scanpy.external as sce
    >>> import phate
    >>> tree_data, tree_clusters = phate.tree.gen_dla(
    ...     n_dim=100,
    ...     n_branch=20,
    ...     branch_length=100,
    ... )
    >>> tree_data.shape
    (2000, 100)
    >>> adata = AnnData(tree_data)
    >>> sce.tl.phate(adata, k=5, a=20, t=150)
    >>> adata.obsm["X_phate"].shape
    (2000, 2)
    >>> sce.pl.phate(adata)

    """
    start = logg.info("computing PHATE")
    adata = adata.copy() if copy else adata
    verbosity = settings.verbosity if verbose is None else verbose
    verbose = verbosity if isinstance(verbosity, bool) else verbosity >= 2
    n_jobs = settings.n_jobs if n_jobs is None else n_jobs
    try:
        import phate
    except ImportError as e:
        msg = (
            "You need to install the package `phate`: please run `pip install "
            "--user phate` in a terminal."
        )
        raise ImportError(msg) from e
    x_phate = phate.PHATE(
        n_components=n_components,
        k=k,
        a=a,
        n_landmark=n_landmark,
        t=t,
        gamma=gamma,
        n_pca=n_pca,
        knn_dist=knn_dist,
        mds_dist=mds_dist,
        mds=mds,
        n_jobs=n_jobs,
        random_state=random_state,
        verbose=verbose,
        **kwargs,
    ).fit_transform(adata)
    # update AnnData instance
    adata.obsm["X_phate"] = x_phate  # annotate samples with PHATE coordinates
    logg.info(
        "    finished",
        time=start,
        deep=("added\n    'X_phate', PHATE coordinates (adata.obsm)"),
    )
    return adata if copy else None
