"""
scanpy paga
"""

import numpy as np
import scanpy as sc
from ..obj_utils import (
    _backup_default_key,
    _delete_backup_key,
    _rename_default_key,
    _set_default_key,
    _restore_default_key,
)


def paga(
    adata,
    key_added=None,
    **kwargs,
):
    """
    Wrapper function for sc.tl.paga, for supporting named slot
    """
    sc.tl.paga(adata, **kwargs)

    if key_added:
        paga_key = f"paga_{key_added}"
        _rename_default_key(adata.uns, "paga", paga_key)
    else:
        _delete_backup_key(adata.uns, "paga")

    return adata


def plot_paga(
    adata,
    use_key="paga",
    basis=None,
    layout=None,
    init_pos=None,
    legend_loc="on data",
    color=None,
    size=None,
    title=None,
    show=None,
    save=None,
    **kwargs,
):
    """Make PAGA plot"""
    if basis is not None and f"X_{basis}" in adata.obsm.keys():
        ax = sc.pl.embedding(
            adata,
            basis=basis,
            color=color,
            legend_loc=legend_loc,
            size=size,
            title=None,
            save=False,
            show=False,
        )

        grouping = adata.uns[use_key]["groups"]
        categories = list(adata.obs[grouping].cat.categories)
        obsm = adata.obsm[f"X_{basis}"]
        group_pos = np.zeros((len(categories), 2))
        for i, label in enumerate(categories):
            offset = 1 if basis.startswith("diffmap") else 0
            _scatter = obsm[adata.obs[grouping] == label, (0 + offset) : (2 + offset)]
            x_pos, y_pos = np.median(_scatter, axis=0)
            group_pos[i] = [x_pos, y_pos]

        _set_default_key(adata.uns, "paga", use_key)
        kwargs["node_size_scale"] = 0
        kwargs["fontsize"] = 1
        kwargs["pos"] = group_pos
        kwargs["color"] = None
        try:
            sc.pl.paga(
                adata,
                ax=ax,
                title=title,
                show=show,
                save=save,
                **kwargs,
            )
        finally:
            _restore_default_key(adata.uns, "paga", use_key)
    else:
        _set_default_key(adata.uns, "paga", use_key)
        try:
            sc.pl.paga(
                adata,
                layout=layout,
                init_pos=init_pos,
                color=color,
                title=title,
                show=show,
                save=save,
                **kwargs,
            )
        finally:
            _restore_default_key(adata.uns, "paga", use_key)

    return adata
