from math import inf

import os

import logging

import numpy as np

import scanpy as sc

import anndata as ad

import pandas as pd

import matplotlib

from matplotlib.figure import Figure

import matplotlib.pyplot as plt

from scipy.sparse.csr import spmatrix

from scipy.stats import chi2

from typing import Mapping, Sequence, Tuple, Iterable, Union

from scipy.sparse import issparse

from sklearn.metrics import adjusted_rand_score, normalized_mutual_info_score, silhouette_samples

from sklearn.neighbors import NearestNeighbors

import psutil

import scib

_cpu_count: Union[None, int] = psutil.cpu_count(logical=False)

if _cpu_count is None:

    _cpu_count: int = psutil.cpu_count(logical=True)

_logger = logging.getLogger(__name__)

def evaluate(adata: ad.AnnData,

             n_epoch: int,

             embedding_key: str = 'delta',

             n_neighbors: int = 15,

             resolutions: Iterable[float] = [0.01, 0.02, 0.04, 0.08, 0.16, 0.32, 0.64],

             clustering_method: str = "leiden",

             cell_type_col: str = "cell_types",

             batch_col: Union[str, None] = "batch_indices",

             color_by: Iterable[str] = None,

             return_fig: bool = False,

             plot_fname: str = "umap",

             plot_ftype: str = "jpg",

             plot_dir: Union[str, None] = None,

             plot_dpi: int = 300,

             min_dist: float = 0.3,

             spread: float = 1,

             n_jobs: int = 1,

             random_state: Union[None, int, np.random.RandomState, np.random.Generator] = 0,

             umap_kwargs: dict = dict()

             ) -> Mapping[str, Union[float, None, Figure]]:

    """Evaluates the clustering and batch correction performance of the given

    embeddings, and optionally plots the embeddings.

    Embeddings will be plotted if return_fig is True or plot_dir is provided.

    When tensorboard_dir is provided, will also save the embeddings using a

    tensorboard SummaryWriter.

    NOTE: Set n_jobs to 1 if you encounter pickling error.

    Args:

        adata: the dataset with the embedding to be evaluated.

        embedding_key: the key to the embedding. Must be in adata.obsm.

        n_neighbors: #neighbors used when computing neithborhood graph and

            calculating entropy of batch mixing / kBET.

        resolutions: a sequence of resolutions used for clustering.

        clustering_method: clustering method used. Should be one of 'leiden' or

            'louvain'.

        cell_type_col: a key in adata.obs to the cell type column.

        batch_col: a key in adata.obs to the batch column.

        return_fig: whether to return the Figure object. Useful for visualizing

            the plot.

        color_by: a list of adata.obs column keys to color the embeddings by.

            If None, will look up adata.uns['color_by']. Only used if is

            drawing.

        plot_fname: file name of the generated plot. Only used if is drawing.

        plot_ftype: file type of the generated plot. Only used if is drawing.

        plot_dir: directory to save the generated plot. If None, do not save

            the plot.

        plot_dpi: dpi to save the plot.

        writer: an initialized SummaryWriter to save the UMAP plot to. Only

            used if is drawing.

        min_dist: the min_dist argument in sc.tl.umap. Only used is drawing.

        spread: the spread argument in sc.tl.umap. Only used if is drawing.

        n_jobs: # jobs to generate. If <= 0, this is set to the number of

            physical cores.

        random_state: random state for knn calculation.

        umap_kwargs: other kwargs to pass to sc.pl.umap.

    Returns:

        A dict storing the ari, nmi, asw, ebm and k_bet of the cell embeddings

        with key "ari", "nmi", "asw", "ebm", "k_bet", respectively. If draw is

        True and return_fig is True, will also store the plotted figure with

        key "fig".

    """

    if cell_type_col and not pd.api.types.is_categorical_dtype(adata.obs[cell_type_col]):

        #_logger.warning("scETM.evaluate assumes discrete cell types. Converting cell_type_col to categorical.")

        adata.obs[cell_type_col] = adata.obs[cell_type_col].astype(str).astype('category')

    if batch_col and not pd.api.types.is_categorical_dtype(adata.obs[batch_col]):

        #_logger.warning("scETM.evaluate assumes discrete batches. Converting batch_col to categorical.")

        adata.obs[batch_col] = adata.obs[batch_col].astype(str).astype('category')

    # calculate neighbors

    _get_knn_indices(adata, use_rep=embedding_key, n_neighbors=n_neighbors, random_state=random_state, calc_knn=True)

    # calculate clustering metrics

    if cell_type_col in adata.obs and len(resolutions) > 0:

        cluster_key, best_ari, best_nmi = clustering(adata, resolutions=resolutions, cell_type_col=cell_type_col, batch_col=batch_col, clustering_method=clustering_method)

    else:

        cluster_key = best_ari = best_nmi = None

    if adata.obs[cell_type_col].nunique() > 1:

        sw = silhouette_samples(adata.X if embedding_key == 'X' else adata.obsm[embedding_key],

                                adata.obs[cell_type_col])

        adata.obs['silhouette_width'] = sw

        asw = np.mean(sw)

        #print(f'{embedding_key}_ASW: {asw:7.4f}')

        asw_2 = scib.me.silhouette(adata, group_key=cell_type_col, embed=embedding_key)

        if batch_col and cell_type_col:

            sw_table = adata.obs.pivot_table(index=cell_type_col, columns=batch_col, values="silhouette_width",

                                             aggfunc="mean")

            #print(f'SW: {sw_table}')

            if plot_dir is not None:

                sw_table.to_csv(os.path.join(plot_dir, f'{plot_fname}.csv'))

    else:

        asw = 0.

        asw_2 = 0.

    # calculate batch correction metrics

    need_batch = batch_col and adata.obs[batch_col].nunique() > 1

    if need_batch:

        ebm = calculate_entropy_batch_mixing(adata,

                                             use_rep=embedding_key,

                                             batch_col=batch_col,

                                             n_neighbors=n_neighbors,

                                             calc_knn=False,

                                             n_jobs=n_jobs,

                                             )

        #print(f'{embedding_key}_BE: {ebm:7.4f}')

        k_bet = calculate_kbet(adata,

                               use_rep=embedding_key,

                               batch_col=batch_col,

                               n_neighbors=n_neighbors,

                               calc_knn=False,

                               n_jobs=n_jobs,

                               )[2]

        #print(f'{embedding_key}_kBET: {k_bet:7.4f}')

        batch_asw = scib.me.silhouette_batch(adata, batch_key=batch_col, group_key='cell_type', embed=embedding_key, verbose=False)

        batch_graph_score = get_graph_connectivity(adata, use_rep=embedding_key,)

    else:

        ebm = k_bet = batch_asw = batch_graph_score = None

    # plot UMAP embeddings

    if return_fig or plot_dir is not None:

        if color_by is None:

            color_by = [batch_col, cell_type_col] if need_batch else [cell_type_col]

        color_by = list(color_by)

        if 'color_by' in adata.uns:

            for col in adata.uns['color_by']:

                if col not in color_by:

                    color_by.insert(0, col)

        if cluster_key is not None:

            color_by = [cluster_key] + color_by

        fig = draw_embeddings(adata=adata, color_by=color_by,

                              min_dist=min_dist, spread=spread,

                              ckpt_dir=plot_dir, fname=f'{plot_fname+str(n_epoch)}.{plot_ftype}', return_fig=return_fig,

                              dpi=plot_dpi,

                              umap_kwargs=umap_kwargs)

    else:

        fig = None

    return dict(

        ari=best_ari,

        nmi=best_nmi,

        asw=asw,

        asw_2=asw_2,

        ebm=ebm,

        k_bet=k_bet,

        batch_asw=batch_asw,

        batch_graph_score=batch_graph_score,

        fig=fig

    )

def evaluate_ari(adata: ad.AnnData,

             n_epoch: int,

             embedding_key: str = 'delta',

             n_neighbors: int = 15,

             resolutions: Iterable[float] = [0.01, 0.02, 0.04, 0.08, 0.16, 0.32, 0.64],

             clustering_method: str = "leiden",

             cell_type_col: str = "cell_types",

             batch_col: Union[str, None] = "batch_indices",

             color_by: Iterable[str] = None,

             return_fig: bool = False,

             plot_fname: str = "umap",

             plot_ftype: str = "jpg",

             plot_dir: Union[str, None] = None,

             plot_dpi: int = 300,

             min_dist: float = 0.3,

             spread: float = 1,

             n_jobs: int = 1,

             random_state: Union[None, int, np.random.RandomState, np.random.Generator] = 0,

             umap_kwargs: dict = dict()

             ) -> Mapping[str, Union[float, None, Figure]]:

    """Evaluates the clustering and batch correction performance of the given

    embeddings, and optionally plots the embeddings.

    Embeddings will be plotted if return_fig is True or plot_dir is provided.

    When tensorboard_dir is provided, will also save the embeddings using a

    tensorboard SummaryWriter.

    NOTE: Set n_jobs to 1 if you encounter pickling error.

    Args:

        adata: the dataset with the embedding to be evaluated.

        embedding_key: the key to the embedding. Must be in adata.obsm.

        n_neighbors: #neighbors used when computing neithborhood graph and

            calculating entropy of batch mixing / kBET.

        resolutions: a sequence of resolutions used for clustering.

        clustering_method: clustering method used. Should be one of 'leiden' or

            'louvain'.

        cell_type_col: a key in adata.obs to the cell type column.

        batch_col: a key in adata.obs to the batch column.

        return_fig: whether to return the Figure object. Useful for visualizing

            the plot.

        color_by: a list of adata.obs column keys to color the embeddings by.

            If None, will look up adata.uns['color_by']. Only used if is

            drawing.

        plot_fname: file name of the generated plot. Only used if is drawing.

        plot_ftype: file type of the generated plot. Only used if is drawing.

        plot_dir: directory to save the generated plot. If None, do not save

            the plot.

        plot_dpi: dpi to save the plot.

        writer: an initialized SummaryWriter to save the UMAP plot to. Only

            used if is drawing.

        min_dist: the min_dist argument in sc.tl.umap. Only used is drawing.

        spread: the spread argument in sc.tl.umap. Only used if is drawing.

        n_jobs: # jobs to generate. If <= 0, this is set to the number of

            physical cores.

        random_state: random state for knn calculation.

        umap_kwargs: other kwargs to pass to sc.pl.umap.

    Returns:

        A dict storing the ari, nmi, asw, ebm and k_bet of the cell embeddings

        with key "ari", "nmi", "asw", "ebm", "k_bet", respectively. If draw is

        True and return_fig is True, will also store the plotted figure with

        key "fig".

    """

    if cell_type_col and not pd.api.types.is_categorical_dtype(adata.obs[cell_type_col]):

        #_logger.warning("scETM.evaluate assumes discrete cell types. Converting cell_type_col to categorical.")

        adata.obs[cell_type_col] = adata.obs[cell_type_col].astype(str).astype('category')

    if batch_col and not pd.api.types.is_categorical_dtype(adata.obs[batch_col]):

        #_logger.warning("scETM.evaluate assumes discrete batches. Converting batch_col to categorical.")

        adata.obs[batch_col] = adata.obs[batch_col].astype(str).astype('category')

    # calculate neighbors

    _get_knn_indices(adata, use_rep=embedding_key, n_neighbors=n_neighbors, random_state=random_state, calc_knn=True)

    # calculate clustering metrics

    if cell_type_col in adata.obs and len(resolutions) > 0:

        cluster_key, best_ari, best_nmi = clustering(adata, resolutions=resolutions, cell_type_col=cell_type_col, batch_col=batch_col, clustering_method=clustering_method)

    else:

        cluster_key = best_ari = best_nmi = None

    return best_ari

def _eff_n_jobs(n_jobs: Union[None, int]) -> int:

    """If n_jobs <= 0, set it as the number of physical cores _cpu_count"""

    if n_jobs is None:

        return 1

    return int(n_jobs) if n_jobs > 0 else _cpu_count

def _calculate_kbet_for_one_chunk(knn_indices, attr_values, ideal_dist, n_neighbors):

    dof = ideal_dist.size - 1

    ns = knn_indices.shape[0]

    results = np.zeros((ns, 2))

    for i in range(ns):

        # NOTE: Do not use np.unique. Some of the batches may not be present in

        # the neighborhood.

        observed_counts = pd.Series(attr_values[knn_indices[i, :]]).value_counts(sort=False).values

        expected_counts = ideal_dist * n_neighbors

        stat = np.sum((observed_counts - expected_counts) ** 2 / expected_counts)

        p_value = 1 - chi2.cdf(stat, dof)

        results[i, 0] = stat

        results[i, 1] = p_value

    return results

def _get_knn_indices(adata: ad.AnnData,

                     use_rep: str = "delta",

                     n_neighbors: int = 25,

                     random_state: int = 0,

                     calc_knn: bool = True

                     ) -> np.ndarray:

    if calc_knn:

        assert use_rep == 'X' or use_rep in adata.obsm, f'{use_rep} not in adata.obsm and is not "X"'

        neighbors = sc.Neighbors(adata)

        neighbors.compute_neighbors(n_neighbors=n_neighbors, knn=True, use_rep=use_rep, random_state=random_state,

                                    write_knn_indices=True)

        adata.obsp['distances'] = neighbors.distances

        adata.obsp['connectivities'] = neighbors.connectivities

        adata.obsm['knn_indices'] = neighbors.knn_indices

        adata.uns['neighbors'] = {

            'connectivities_key': 'connectivities',

            'distances_key': 'distances',

            'knn_indices_key': 'knn_indices',

            'params': {

                'n_neighbors': n_neighbors,

                'use_rep': use_rep,

                'metric': 'euclidean',

                'method': 'umap'

            }

        }

    else:

        assert 'neighbors' in adata.uns, 'No precomputed knn exists.'

        assert adata.uns['neighbors']['params'][

                   'n_neighbors'] >= n_neighbors, f"pre-computed n_neighbors is {adata.uns['neighbors']['params']['n_neighbors']}, which is smaller than {n_neighbors}"

    return adata.obsm['knn_indices']

def get_graph_connectivity(

        adata: ad.AnnData,

        use_rep: str = "delta",):

    sc.pp.neighbors(adata, use_rep=use_rep)

    score = scib.me.graph_connectivity(adata, label_key='cell_type')

    return score

def calculate_kbet(

        adata: ad.AnnData,

        use_rep: str = "delta",

        batch_col: str = "batch_indices",

        n_neighbors: int = 25,

        alpha: float = 0.05,

        random_state: int = 0,

        n_jobs: Union[None, int] = None,

        calc_knn: bool = True

) -> Tuple[float, float, float]:

    """Calculates the kBET metric of the data.

    kBET measures if cells from different batches mix well in their local

    neighborhood.

    Args:

        adata: annotated data matrix.

        use_rep: the embedding to be used. Must exist in adata.obsm.

        batch_col: a key in adata.obs to the batch column.

        n_neighbors: # nearest neighbors.

        alpha: acceptance rate threshold. A cell is accepted if its kBET

            p-value is greater than or equal to alpha.

        random_state: random seed. Used only if method is "hnsw".

        n_jobs: # jobs to generate. If <= 0, this is set to the number of

            physical cores.

        calc_knn: whether to re-calculate the kNN graph or reuse the one stored

            in adata.

    Returns:

        stat_mean: mean kBET chi-square statistic over all cells.

        pvalue_mean: mean kBET p-value over all cells.

        accept_rate: kBET Acceptance rate of the sample.

    """

    _logger.info('Calculating kbet...')

    assert batch_col in adata.obs

    if adata.obs[batch_col].dtype.name != "category":

        _logger.warning(f'Making the column {batch_col} of adata.obs categorical.')

        adata.obs[batch_col] = adata.obs[batch_col].astype('category')

    ideal_dist = (

        adata.obs[batch_col].value_counts(normalize=True, sort=False).values

    )  # ideal no batch effect distribution

    nsample = adata.shape[0]

    nbatch = ideal_dist.size

    attr_values = adata.obs[batch_col].values.copy()

    attr_values.categories = range(nbatch)

    knn_indices = _get_knn_indices(adata, use_rep, n_neighbors, random_state, calc_knn)

    # partition into chunks

    n_jobs = min(_eff_n_jobs(n_jobs), nsample)

    starts = np.zeros(n_jobs + 1, dtype=int)

    quotient = nsample // n_jobs

    remainder = nsample % n_jobs

    for i in range(n_jobs):

        starts[i + 1] = starts[i] + quotient + (1 if i < remainder else 0)

    from joblib import Parallel, delayed, parallel_backend

    with parallel_backend("loky", n_jobs=n_jobs):

        kBET_arr = np.concatenate(

            Parallel()(

                delayed(_calculate_kbet_for_one_chunk)(

                    knn_indices[starts[i]: starts[i + 1], :], attr_values, ideal_dist, n_neighbors

                )

                for i in range(n_jobs)

            )

        )

    res = kBET_arr.mean(axis=0)

    stat_mean = res[0]

    pvalue_mean = res[1]

    accept_rate = (kBET_arr[:, 1] >= alpha).sum() / nsample

    return (stat_mean, pvalue_mean, accept_rate)

def _entropy(hist_data):

    _, counts = np.unique(hist_data, return_counts=True)

    freqs = counts / counts.sum()

    return (-freqs * np.log(freqs + 1e-30)).sum()

def _entropy_batch_mixing_for_one_pool(batches, knn_indices, nsample, n_samples_per_pool):

    indices = np.random.choice(

        np.arange(nsample), size=n_samples_per_pool)

    return np.mean(

        [

            _entropy(batches[knn_indices[indices[i]]])

            for i in range(n_samples_per_pool)

        ]

    )

def calculate_entropy_batch_mixing(

        adata: ad.AnnData,

        use_rep: str = "delta",

        batch_col: str = "batch_indices",

        n_neighbors: int = 50,

        n_pools: int = 50,

        n_samples_per_pool: int = 100,

        random_state: int = 0,

        n_jobs: Union[None, int] = None,

        calc_knn: bool = True

) -> float:

    """Calculates the entropy of batch mixing of the data.

    kBET measures if cells from different batches mix well in their local

    neighborhood.

    Args:

        adata: annotated data matrix.

        use_rep: the embedding to be used. Must exist in adata.obsm.

        batch_col: a key in adata.obs to the batch column.

        n_neighbors: # nearest neighbors.

        n_pools: #pools of cells to calculate entropy of batch mixing.

        n_samples_per_pool: #cells per pool to calculate within-pool entropy.

        random_state: random seed. Used only if method is "hnsw".

        n_jobs: # jobs to generate. If <= 0, this is set to the number of

            physical cores.

        calc_knn: whether to re-calculate the kNN graph or reuse the one stored

            in adata.

    Returns:

        score: the mean entropy of batch mixing, averaged from n_pools samples.

    """

    _logger.info('Calculating batch mixing entropy...')

    nsample = adata.n_obs

    knn_indices = _get_knn_indices(adata, use_rep, n_neighbors, random_state, calc_knn)

    from joblib import Parallel, delayed, parallel_backend

    with parallel_backend("loky", n_jobs=n_jobs, inner_max_num_threads=1):

        score = np.mean(

            Parallel()(

                delayed(_entropy_batch_mixing_for_one_pool)(

                    adata.obs[batch_col], knn_indices, nsample, n_samples_per_pool

                )

                for _ in range(n_pools)

            )

        )

    return score

def clustering(

        adata: ad.AnnData,

        resolutions: Sequence[float],

        clustering_method: str = "leiden",

        cell_type_col: str = "cell_types",

        batch_col: str = "batch_indices"

) -> Tuple[str, float, float]:

    """Clusters the data and calculate agreement with cell type and batch

    variable.

    This method cluster the neighborhood graph (requires having run sc.pp.

    neighbors first) with "clustering_method" algorithm multiple times with the

    given resolutions, and return the best result in terms of ARI with cell

    type.

    Other metrics such as NMI with cell type, ARi with batch are logged but not

    returned. (TODO: also return these metrics)

    Args:

        adata: the dataset to be clustered. adata.obsp shouhld contain the keys

            'connectivities' and 'distances'.

        resolutions: a list of leiden/louvain resolution parameters. Will

            cluster with each resolution in the list and return the best result

            (in terms of ARI with cell type).

        clustering_method: Either "leiden" or "louvain".

        cell_type_col: a key in adata.obs to the cell type column.

        batch_col: a key in adata.obs to the batch column.

    Returns:

        best_cluster_key: a key in adata.obs to the best (in terms of ARI with

            cell type) cluster assignment column.

        best_ari: the best ARI with cell type.

        best_nmi: the best NMI with cell type.

    """

    assert len(resolutions) > 0, f'Must specify at least one resolution.'

    if clustering_method == 'leiden':

        clustering_func = sc.tl.leiden

    elif clustering_method == 'louvain':

        clustering_func = sc.tl.louvain

    else:

        raise ValueError("Please specify louvain or leiden for the clustering method argument.")

    #_logger.info(f'Performing {clustering_method} clustering')

    assert cell_type_col in adata.obs, f"{cell_type_col} not in adata.obs"

    best_res, best_ari, best_nmi = None, -inf, -inf

    for res in resolutions:

        col = f'{clustering_method}_{res}'

        clustering_func(adata, resolution=res, key_added=col)

        ari = adjusted_rand_score(adata.obs[cell_type_col], adata.obs[col])

        nmi = normalized_mutual_info_score(adata.obs[cell_type_col], adata.obs[col])

        n_unique = adata.obs[col].nunique()

        if ari > best_ari:

            best_res = res

            best_ari = ari

        if nmi > best_nmi:

            best_nmi = nmi

        if batch_col in adata.obs and adata.obs[batch_col].nunique() > 1:

            ari_batch = adjusted_rand_score(adata.obs[batch_col], adata.obs[col])

            #print(f'Resolution: {res:5.3g}\tARI: {ari:7.4f}\tNMI: {nmi:7.4f}\tbARI: {ari_batch:7.4f}\t# labels: {n_unique}')

        else:

            #print(f'Resolution: {res:5.3g}\tARI: {ari:7.4f}\tNMI: {nmi:7.4f}\t# labels: {n_unique}')

            a=None

    return f'{clustering_method}_{best_res}', best_ari, best_nmi

def draw_embeddings(adata: ad.AnnData,

                    color_by: Union[str, Sequence[str], None] = None,

                    min_dist: float = 0.3,

                    spread: float = 1,

                    ckpt_dir: str = '.',

                    fname: str = "umap.pdf",

                    return_fig: bool = False,

                    dpi: int = 300,

                    umap_kwargs: dict = dict()

                    ) -> Union[None, Figure]:

    """Embeds, plots and optionally saves the neighborhood graph with UMAP.

    Requires having run sc.pp.neighbors first.

    Args:

        adata: the dataset to draw. adata.obsp shouhld contain the keys

            'connectivities' and 'distances'.

        color_by: a str or a list of adata.obs keys to color the points in the

            scatterplot by. E.g. if both cell_type_col and batch_col is in

            color_by, then we would have two plots colored by cell type and

            batch variables, respectively.

        min_dist: The effective minimum distance between embedded points.

            Smaller values will result in a more clustered/clumped embedding

            where nearby points on the manifold are drawn closer together,

            while larger values will result on a more even dispersal of points.

        spread: The effective scale of embedded points. In combination with

            `min_dist` this determines how clustered/clumped the embedded

            points are.

        ckpt_dir: where to save the plot. If None, do not save the plot.

        fname: file name of the saved plot. Only used if ckpt_dir is not None.

        return_fig: whether to return the Figure object. Useful for visualizing

            the plot.

        dpi: the dpi of the saved plot. Only used if ckpt_dir is not None.

        umap_kwargs: other kwargs to pass to sc.pl.umap.

    Returns:

        If return_fig is True, return the figure containing the plot.

    """

    #_logger.info(f'Plotting UMAP embeddings...')

    sc.tl.umap(adata, min_dist=min_dist, spread=spread)

    fig = sc.pl.umap(adata, color=color_by, show=False, return_fig=True, **umap_kwargs)

    if ckpt_dir is not None:

        assert os.path.exists(ckpt_dir), f'ckpt_dir {ckpt_dir} does not exist.'

        fig.savefig(

            os.path.join(ckpt_dir, fname),

            dpi=dpi, bbox_inches='tight'

        )

    if return_fig:

        return fig

    fig.clf()

    plt.close(fig)

def set_figure_params(

        matplotlib_backend: str = 'agg',

        dpi: int = 120,

        frameon: bool = True,

        vector_friendly: bool = True,

        fontsize: int = 10,

        figsize: Sequence[int] = (10, 10)

):

    """Set figure parameters.

    Args

        backend: the backend to switch to.  This can either be one of th

            standard backend names, which are case-insensitive:

            - interactive backends:

                GTK3Agg, GTK3Cairo, MacOSX, nbAgg,

                Qt4Agg, Qt4Cairo, Qt5Agg, Qt5Cairo,

                TkAgg, TkCairo, WebAgg, WX, WXAgg, WXCairo

            - non-interactive backends:

                agg, cairo, pdf, pgf, ps, svg, template

            or a string of the form: ``module://my.module.name``.

        dpi: resolution of rendered figures – this influences the size of

            figures in notebooks.

        frameon: add frames and axes labels to scatter plots.

        vector_friendly: plot scatter plots using `png` backend even when

            exporting as `pdf` or `svg`.

        fontsize: the fontsize for several `rcParams` entries.

        figsize: plt.rcParams['figure.figsize'].

    """

    matplotlib.use(matplotlib_backend)

    sc.set_figure_params(dpi=dpi, figsize=figsize, fontsize=fontsize, frameon=frameon, vector_friendly=vector_friendly)