import os

import pickle

import time

from sklearn.ensemble import RandomForestClassifier

from sklearn.linear_model import LogisticRegression

# import sklearn.linear_model as lm

from sklearn.metrics import (

    accuracy_score,

    auc,

    average_precision_score,

    balanced_accuracy_score,

    classification_report,

    confusion_matrix,

    matthews_corrcoef,

    recall_score,

    roc_auc_score,

)

from sklearn.model_selection import GridSearchCV

from sklearn.neighbors import KNeighborsClassifier

from sklearn.svm import SVC

from .GATclassifier import GATclassifier

# import pandas as pd

# import numpy as np

from .utils_func import *

import sklearn.ensemble._forest

import sklearn.linear_model._logistic

import sklearn.neighbors._classification

import sklearn.svm._classes

from anndata._core.anndata import AnnData

from matplotlib.axes._axes import Axes

from matplotlib.figure import Figure

from matplotlib.gridspec import GridSpec

from numpy import float64, ndarray

from pandas.core.frame import DataFrame

from pandas.core.series import Series

from scpanel.GATclassifier import GATclassifier

from torch import Tensor

from typing import Any, Dict, List, Optional, Tuple, Union

def transform_adata(adata_train: AnnData, adata_test_dict: Dict[str, AnnData], selected_gene: Optional[List[str]]=None) -> Tuple[AnnData, AnnData]:

    ## Transforming train set and test set from the same dataset (batch effect free)

    ## subset adata_train with selected genes

    ## subset adata_test_dict with selected cell types and genes

    ## WATCH OUT: X matrix in adata_test_dict is log-normalized, need to scale further

    if selected_gene == None:

        selected_gene = adata_train.uns["svm_rfe_genes"]

    adata_train_final = adata_train[:, selected_gene]

    mean = adata_train_final.var["mean"].values

    std = adata_train_final.var["std"].values

    ct_selected = adata_train_final.obs.ct.unique()[0]

    # transform test data with selected gene, celltype and scaling

    adata_test = adata_test_dict[ct_selected].copy()

    adata_test_final = adata_test[:, selected_gene].copy()

    if isinstance(adata_test_final.X, np.ndarray):

        test_X = adata_test_final.X

    else:

        test_X = adata_test_final.X.toarray()

    test_X -= mean

    test_X /= std

    max_value = 10

    test_X[test_X > max_value] = max_value

    adata_test_final.X = test_X

    return adata_train_final, adata_test_final

def models_train(adata_train_final: AnnData, search_grid: bool, out_dir: Optional[str]=None, param_grid: Optional[Dict[str, Dict[str, int]]]=None) -> List[Union[Tuple[str, sklearn.linear_model._logistic.LogisticRegression], Tuple[str, sklearn.ensemble._forest.RandomForestClassifier], Tuple[str, sklearn.svm._classes.SVC], Tuple[str, sklearn.neighbors._classification.KNeighborsClassifier], Tuple[str, GATclassifier]]]:

    X_tr, y_tr, adj_tr = get_X_y_from_ann(

        adata_train_final, return_adj=True, n_neigh=10

    )

    sample_weight = compute_cell_weight(adata_train_final)

    # Make sure no nan in matrix

    X_tr = np.nan_to_num(X_tr)

    grid_search = search_grid

    models = [

        ("LR", LogisticRegression(solver="saga", max_iter=500, random_state=42)),

        ("RF", RandomForestClassifier(random_state=42)),

        ("SVM", SVC(probability=True, random_state=42)),

        ("KNN", KNeighborsClassifier()),

        (

            "GAT",

            GATclassifier(

                nFeatures=adata_train_final.n_vars, NumParts=10, nEpochs=1000, verbose=1

            ),

        ),

    ]

    # Parameter tuning grids-------------------------

    LR_params = [{"C": [10, 1.0, 0.1, 0.01], "max_iter": [10, 50, 200, 500]}]

    RF_params = [

        {"max_depth": [2, 5, 10, 15, 20, 30, None], "n_estimators": [50, 100, 500]}

    ]

    SVM_params = [{"C": [100, 10, 1.0, 0.1, 0.001], "gamma": [1, 0.1, 0.01, 0.001]}]

    KNN_params = [{"n_neighbors": [3, 5, 10, 20, 50], "p": [1, 2]}]

    my_grid = {"LR": LR_params, "RF": RF_params, "SVM": SVM_params, "KNN": KNN_params}

    clfs = []

    names = []

    runtimes = []

    best_params = []

    for name, model in models:

        start_time = time.time()

        if grid_search:

            if name != "GAT":

                clf = GridSearchCV(

                    model, my_grid[name], cv=5, scoring="roc_auc", n_jobs=10

                )

            else:

                clf = model

        else:

            clf = model

            if param_grid is not None:

                if name in param_grid:

                    clf.set_params(**param_grid[name])

        if name == "GAT":

            clf.fit(X_tr, y_tr, adj_tr)

        elif name == "KNN":

            clf.fit(X_tr, y_tr)

        else:

            clf.fit(X_tr, y_tr, sample_weight=sample_weight)

        runtime = time.time() - start_time

        # save outputs

        clfs.append((name, clf))

        names.append(name)

        runtimes.append(runtime)

        print("---%s finished in %s seconds ---" % (name, runtime))

    # save models

    if out_dir is not None:

        if not os.path.exists(out_dir):

            os.makedirs(out_dir)

        with open(f"{out_dir}/clfs.pkl", "wb") as f:

            pickle.dump(clfs, f, protocol=pickle.HIGHEST_PROTOCOL)

            f.close()

        with open(f"{out_dir}/adata_train_final.pkl", "wb") as f:

            pickle.dump(adata_train_final, f, protocol=pickle.HIGHEST_PROTOCOL)

            f.close()

    return clfs

def models_predict(clfs: List[Union[Tuple[str, sklearn.linear_model._logistic.LogisticRegression], Tuple[str, sklearn.ensemble._forest.RandomForestClassifier], Tuple[str, sklearn.svm._classes.SVC], Tuple[str, sklearn.neighbors._classification.KNeighborsClassifier], Tuple[str, GATclassifier]]], adata_test_final: AnnData, out_dir: Optional[str]=None) -> Tuple[AnnData, List[Union[Tuple[str, ndarray], Tuple[str, Tensor]]], List[Tuple[str, ndarray]]]:

    X_test, y_test, adj_test = get_X_y_from_ann(

        adata_test_final, return_adj=True, n_neigh=10

    )

    X_test = np.nan_to_num(X_test)

    ## Predicting---------------

    y_pred_list = []

    y_pred_score_list = []

    for name, clf in clfs:

        if name == "GAT":

            y_pred = clf.predict(X_test, y_test, adj_test)

            y_pred_score = clf.predict_proba(X_test, y_test, adj_test)

        else:

            y_pred = clf.predict(X_test)

            y_pred_score = clf.predict_proba(X_test)

        y_pred_list.append((name, y_pred))

        y_pred_score_list.append((name, y_pred_score))

    # add prediction result to adata_test_final

    y_pred = pd.DataFrame(dict([(name + "_pred", pred) for name, pred in y_pred_list]))

    y_pred_score = pd.DataFrame(

        dict([(name + "_pred_score", pred[:, 1]) for name, pred in y_pred_score_list])

    )

    y_pred_df = pd.concat([y_pred, y_pred_score], axis=1)

    y_pred_df.index = adata_test_final.obs.index

    if set(y_pred_df.columns).issubset(set(adata_test_final.obs.columns)):

        print("Prediction result already exits in test adata, overwrite it...")

        adata_test_final.obs.update(y_pred_df)

    else:

        adata_test_final.obs = pd.concat([adata_test_final.obs, y_pred_df], axis=1)

    # calcuate median prediction score out of 5 classifiers

    pred_col = [

        col for col in adata_test_final.obs.columns if col.endswith("_pred_score")

    ]

    adata_test_final.obs["median_pred_score"] = adata_test_final.obs[pred_col].median(

        axis=1

    )

    return adata_test_final, y_pred_list, y_pred_score_list

def models_score(adata_test_final, y_pred_list, y_pred_score_list, out_dir=None):

    X_test, y_test = get_X_y_from_ann(adata_test_final)

    ## Scoring-------------------------------------

    ## define scoring metrics (from sklearn)

    scorers = {

        "accuracy": (accuracy_score, {}),

        "balanced_accuracy": (balanced_accuracy_score, {}),

        "MCC": (matthews_corrcoef, {}),

    }  # Passing Dictionary as Arguments to Function

    scorers_prob = {

        "AUROC": (roc_auc_score, {}),

        "AUPRC": (average_precision_score, {}),

    }

    ## calculate

    eval_res_1 = pd.DataFrame()

    for name, y_pred in y_pred_list:

        eval_res_dict = dict(

            [

                (score_name, score_func(y_test, y_pred, **score_para))

                for score_name, (score_func, score_para) in scorers.items()

            ]

        )

        eval_res_i = pd.DataFrame(eval_res_dict, index=[name])

        eval_res_1 = pd.concat(objs=[eval_res_1, eval_res_i], axis=0)

    eval_res_2 = pd.DataFrame()

    for name, y_pred_score in y_pred_score_list:

        eval_res_dict = dict(

            [

                (score_name, score_func(y_test, y_pred_score[:, 1], **score_para))

                for score_name, (score_func, score_para) in scorers_prob.items()

            ]

        )

        eval_res_i = pd.DataFrame(eval_res_dict, index=[name])

        eval_res_2 = pd.concat(objs=[eval_res_2, eval_res_i], axis=0)

    eval_res = pd.concat(objs=[eval_res_2, eval_res_1], axis=1)

    if out_dir is not None:

        if not os.path.exists(out_dir):

            os.makedirs(out_dir)

        eval_res.to_csv(f"{out_dir}/eval_res.csv")

    return eval_res

def cal_sample_auc(df: DataFrame, score_col: str) -> float64:

    cell_prob = df[score_col].sort_values()

    # rank the cell probability ascendingly and normalize

    cell_rank = cell_prob.rank(method="first") / cell_prob.rank(method="first").max()

    sample_auc = auc(cell_rank, cell_prob)

    return sample_auc

def auc_pvalue(row: Series) -> float:

    if row.name[1] == 1:

        p_value = np.mean(row < 0.5)

    elif row.name[1] == 0:

        p_value = np.mean(row > 0.5)

    if p_value == 0:

        p_value = 1 / row.size

    return p_value

def pt_pred(adata_test_final: AnnData, cell_pred_col: str="median_pred_score", num_bootstrap: Optional[int]=None) -> AnnData:

    sample_auc = adata_test_final.obs.groupby("patient_id").apply(

        lambda df: cal_sample_auc(df, cell_pred_col)

    )

    adata_test_final.obs[cell_pred_col + "_sample_auc"] = (

        adata_test_final.obs["patient_id"].map(sample_auc).astype(float)

    )

    adata_test_final.obs[cell_pred_col + "_sample_pred"] = (

        adata_test_final.obs[cell_pred_col + "_sample_auc"] >= 0.5

    ).astype(int)

    if num_bootstrap is not None:

        auc_df = pd.DataFrame()

        for i in range(num_bootstrap):

            df = adata_test_final.obs.groupby("patient_id").sample(

                frac=1, replace=True, random_state=i

            )

            auc = (

                df.groupby(["patient_id", cell_pred_col + "_sample_pred"])

                .apply(lambda df: cal_sample_auc(df, cell_pred_col))

                .to_frame(name=i)

            )

            auc_df = pd.concat([auc_df, auc], axis=1)

        auc_df[cell_pred_col + "_sample_auc_pvalue"] = auc_df.apply(

            lambda row: auc_pvalue(row), axis=1

        )

        # store auc from each bootstrap iteration in adata.uns

        adata_test_final.uns[cell_pred_col + "_auc_df"] = auc_df

        # store auc_pvalue for each sample in adata.obs

        auc_df = auc_df.droplevel(cell_pred_col + "_sample_pred")

        adata_test_final.obs[cell_pred_col + "_sample_auc_pvalue"] = (

            adata_test_final.obs["patient_id"].map(

                auc_df[cell_pred_col + "_sample_auc_pvalue"]

            )

        )

    return adata_test_final

def pt_score(adata_test_final: AnnData, cell_pred_col: str="median_pred_score") -> AnnData:

    ## Calculate precision, recall, f1score and accuracy at patient level

    from sklearn.metrics import precision_recall_fscore_support

    pred_col = cell_pred_col

    res_prefix = cell_pred_col

    pt_pred_res = (

        adata_test_final.obs[["label", "patient_id", f"{res_prefix}_sample_pred"]]

        .drop_duplicates()

        .set_index("patient_id")

    )

    # precision, recall, f1score

    pt_score_res = precision_recall_fscore_support(

        pt_pred_res["label"],

        pt_pred_res[f"{res_prefix}_sample_pred"],

        average="weighted",

    )

    # accuracy

    pt_acc_res = accuracy_score(

        pt_pred_res["label"], pt_pred_res[f"{res_prefix}_sample_pred"]

    )

    # specificity

    pt_spec_res = recall_score(

        pt_pred_res["label"], pt_pred_res[f"{res_prefix}_sample_pred"], pos_label=0

    )

    pt_score_res = pd.DataFrame(list(pt_score_res) + [pt_acc_res] + [pt_spec_res])

    pt_score_res = pt_score_res.iloc[[0, 1, 2, 4, 5], :]

    pt_score_res.index = [

        "precision",

        "sensitivity",

        "f1score",

        "accuracy",

        "specificity",

    ]

    pt_score_res.columns = [res_prefix]

    if "sample_score" not in adata_test_final.uns:

        adata_test_final.uns["sample_score"] = pt_score_res

    else:

        adata_test_final.uns["sample_score"] = adata_test_final.uns[

            "sample_score"

        ].merge(pt_score_res, left_index=True, right_index=True, suffixes=("_x", ""))

        adata_test_final.uns["sample_score"].drop(

            adata_test_final.uns["sample_score"].filter(regex="_x$").columns,

            axis=1,

            inplace=True,

        )

    return adata_test_final

from math import pi

# Plot functions

import matplotlib.pyplot as plt

import seaborn as sns

from matplotlib import rcParams

def _panel_grid(hspace: float, wspace: float, ncols: int, num_panels: int) -> Tuple[Figure, GridSpec]:

    from matplotlib import gridspec

    n_panels_x = min(ncols, num_panels)

    n_panels_y = np.ceil(num_panels / n_panels_x).astype(int)

    # each panel will have the size of rcParams['figure.figsize']

    fig = plt.figure(

        figsize=(

            n_panels_x * rcParams["figure.figsize"][0] * (1 + wspace),

            n_panels_y * rcParams["figure.figsize"][1],

        ),

    )

    left = 0.2 / n_panels_x

    bottom = 0.13 / n_panels_y

    gs = gridspec.GridSpec(

        nrows=n_panels_y,

        ncols=n_panels_x,

        left=left,

        right=1 - (n_panels_x - 1) * left - 0.01 / n_panels_x,

        bottom=bottom,

        top=1 - (n_panels_y - 1) * bottom - 0.1 / n_panels_y,

        hspace=hspace,

        wspace=wspace,

    )

    return fig, gs

def plot_roc_curve(

    adata_test_final: AnnData,

    sample_id: Series,

    cell_pred_col: str,

    ncols: int=4,

    hspace: float=0.25,

    wspace: None=None,

    ax: None=None,

    scatter_kws: Optional[Dict[str, int]]=None,

    legend_kws: Optional[Dict[str, Dict[str, int]]]=None,

) -> List[Axes]:

    """

    Parameters

    ----------

    - adata_test_final: AnnData,

    - sample_id: str | Sequence,

    - cell_pred_col: str = 'median_pred_score',

    - ncols: int = 4,

    - hspace: float =0.25,

    - wspace: float | None = None,

    - ax: Axes | None = None,

    - scatter_kws: dict | None = None, Arguments to pass to matplotlib.pyplot.scatter()

    Returns

    -------

    Axes

    Examples

    --------

    plot_roc_curve(adata_test_final,

               sample_id = ['C3','C6','H1'],

               cell_pred_col = 'median_pred_score',

               scatter_kws={'s':10})

    """

    # turn sample_id into a python list

    ## if sample_id is string or None, wrap it with []

    ## if sample_id is already sequential, turn it into a list

    sample_id = (

        [sample_id]

        if isinstance(sample_id, str) or sample_id is None

        else list(sample_id)

    )

    ##########

    # Layout #

    ##########

    if scatter_kws is None:

        scatter_kws = {}

    if legend_kws is None:

        legend_kws = {}

    if wspace is None:

        #  try to set a wspace that is not too large or too small given the

        #  current figure size

        wspace = 0.75 / rcParams["figure.figsize"][0] + 0.02

    # if plotting multiple panels for elements in sample_id

    if len(sample_id) > 1:

        if ax is not None:

            raise ValueError(

                "Cannot specify `ax` when plotting multiple panels "

                "(each for a given value of 'color')."

            )

        fig, grid = _panel_grid(hspace, wspace, ncols, len(sample_id))

    else:

        grid = None

        if ax is None:

            fig = plt.figure()

            ax = fig.add_subplot(111)

    ############

    # Plotting #

    ############

    axs = []

    for count, _sample_id in enumerate(sample_id):

        if grid:

            ax = plt.subplot(grid[count])

            axs.append(ax)

        # prediction probability of class 1 for sample_id

        cell_prob = adata_test_final.obs.loc[

            adata_test_final.obs["patient_id"] == sample_id[count]

        ][cell_pred_col]

        cell_prob = cell_prob.sort_values(ascending=True)

        # rank of cell_prob and normalize

        cell_rank = (

            cell_prob.rank(method="first") / cell_prob.rank(method="first").max()

        )

        # auc

        sample_auc = adata_test_final.obs.loc[

            adata_test_final.obs["patient_id"] == sample_id[count]

        ][cell_pred_col + "_sample_auc"].unique()[0]

        # auc-pvalue

        sample_auc_pvalue = adata_test_final.obs.loc[

            adata_test_final.obs["patient_id"] == sample_id[count]

        ][cell_pred_col + "_sample_auc_pvalue"].unique()[0]

        ax.scatter(x=cell_rank, y=cell_prob, c=".3", **scatter_kws)

        ax.plot(

            cell_rank,

            cell_prob,

            label=f"AUC = {sample_auc:.3f} \np-value = {sample_auc_pvalue:.1e}",

            zorder=0,

        )

        ax.plot(

            [0, 1], [0, 1], linestyle="--", color=".5", zorder=0, label="Random guess"

        )

        # ax.text(x = 0.99, y = 0.01, s = f'AUC: {sample_auc:.3f}',

        #         horizontalalignment='right',

        #         verticalalignment='bottom')

        ax.spines[["right", "top"]].set_visible(False)

        ax.set_xlabel("Rank")

        ax.set_ylabel("Prediction Probability (Cell)")

        ax.set_title(f"{_sample_id}")

        ax.set_aspect("equal")

        if not bool(legend_kws):

            ax.legend(prop=dict(size=8 * rcParams["figure.figsize"][0] / ncols))

        else:

            ax.legend(**legend_kws)

    axs = axs if grid else ax

    return axs

def convert_pvalue_to_asterisks(pvalue: float) -> str:

    if pvalue <= 0.0001:

        return "****"

    elif pvalue <= 0.001:

        return "***"

    elif pvalue <= 0.01:

        return "**"

    elif pvalue <= 0.05:

        return "*"

    return "ns"

# plot cell level probabilities for each patient

def plot_violin(

    adata: AnnData,

    cell_pred_col: str="median_pred_score",

    dot_size: int=2,

    ax: Optional[Axes]=None,

    palette: Optional[Dict[str, str]]=None,

    xticklabels_color: bool=False,

    text_kws: Dict[Any, Any]={},

) -> Axes:

    """

    Violin Plots for cell-level prediction probabilities in each sample.

    Parameters:

    - adata: AnnData Object

    - cell_pred_col: string, name of the column with cell-level prediction probabilities

    in adata.obs (default: 'median_pred_score')

    - pt_stat: string, a test for the null hypothesis that the distribution of probabilities

    in this sample is different from the population (default: 'perm')

        Options:

        - 'perm': permutation test

        - 't-test': one-sample t-test

    - fig_size: tuple, size of figure (default: (10, 3))

    - dot_size: float, Radius of the markers in stripplot.

    Returns:

        ax

    """

    # A. organize input data for plotting--------------

    res_prefix = cell_pred_col

    ## cell-level data

    pred_score_df = adata.obs[

        [

            cell_pred_col,

            "y",

            "label",

            "patient_id",

            f"{res_prefix}_sample_auc",

            f"{res_prefix}_sample_auc_pvalue",

        ]

    ].copy()

    ## sample-level data

    sample_pData = pred_score_df.groupby(

        [

            "y",

            "label",

            "patient_id",

            f"{res_prefix}_sample_auc",

            f"{res_prefix}_sample_auc_pvalue",

        ],

        observed=True,

        as_index=False,

    )[cell_pred_col].max()

    sample_pData.rename(columns={cell_pred_col: "y_pos"}, inplace=True)

    sample_pData = sample_pData.sort_values(by=f"{res_prefix}_sample_auc").reset_index(

        drop=True

    )

    sample_order = sample_pData.patient_id.tolist()

    sample_threshold_index = (

        sample_pData[f"{res_prefix}_sample_auc"]

        .where(sample_pData[f"{res_prefix}_sample_auc"] >= 0.5)

        .first_valid_index()

    )

    if sample_threshold_index is None:

        if (sample_pData[f"{res_prefix}_sample_auc"] >= 0.5).all():

            sample_threshold = -0.5

        else:

            sample_threshold = len(sample_pData[f"{res_prefix}_sample_auc"]) - 0.5

    else:

        sample_threshold = sample_threshold_index - 0.5

    # B. plot--------------------------------------------

    if ax is None:

        ax = plt.gca()

    # Hide the right and top spines

    ax.spines[["right", "top"]].set_visible(False)

    # Violin plot

    sns.violinplot(

        y=cell_pred_col,

        x="patient_id",

        data=pred_score_df,

        order=sample_order,

        color="0.8",

        scale="width",

        fontsize=15,

        ax=ax,

        cut=0,

    )

    # Strip plot

    sns.stripplot(

        y=cell_pred_col,

        x="patient_id",

        hue="y",

        data=pred_score_df,

        order=sample_order,

        dodge=False,

        jitter=True,

        size=dot_size,

        ax=ax,

        palette=palette,

    )

    ax.axhline(y=0.5, color="0.8", linestyle="--")

    ax.axvline(x=sample_threshold, color="0.8", linestyle="--")

    # Add statistical signifiance (asterisks (*)) on top of each violin

    ## get position x

    yposlist = (sample_pData["y_pos"] + 0.03).tolist()

    ## get position y

    xposlist = range(len(yposlist))

    ## get text

    pvalue_list = sample_pData[f"{res_prefix}_sample_auc_pvalue"].tolist()

    asterisks_list = [convert_pvalue_to_asterisks(pvalue) for pvalue in pvalue_list]

    perm_stat_list = [

        "%.3f" % perm_stat

        for perm_stat in sample_pData[f"{res_prefix}_sample_auc"].tolist()

    ]

    text_list = [

        perm_stat + "\n" + asterisk

        for perm_stat, asterisk in zip(perm_stat_list, asterisks_list)

    ]

    for k in range(len(asterisks_list)):

        ax.text(x=xposlist[k], y=yposlist[k], s=text_list[k], ha="center", **text_kws)

    ax.set_title(cell_pred_col, pad=30)

    ax.set_xlabel(None)

    ax.set_ylabel("Prediction Probablity (Cell)", fontsize=13)

    ax.plot()

    ax.set_xticks(ax.get_xticks(), ax.get_xticklabels(), rotation=45, ha="right")

    if xticklabels_color:

        for xtick in ax.get_xticklabels():

            x_label = xtick.get_text()

            x_label_cate = sample_pData["y"][

                sample_pData["patient_id"] == x_label

            ].values[0]

            xtick.set_color(palette[x_label_cate])

    ax.legend(loc="upper left", title="Patient Label", bbox_to_anchor=(1.04, 1))

    return ax

### Plot patient level prediction scores

def make_single_spider(adata_test_final: AnnData, metric_idx: int, color: str, nrow: int, ncol: int) -> None:

    # number of variable

    categories = adata_test_final.uns["sample_score"].index.tolist()

    N = len(adata_test_final.uns["sample_score"].index)

    # We are going to plot the first line of the data frame.

    # But we need to repeat the first value to close the circular graph:

    values = (

        adata_test_final.uns["sample_score"]

        .iloc[:, metric_idx]

        .values.flatten()

        .tolist()

    )

    values += values[:1]

    # What will be the angle of each axis in the plot? (we divide the plot / number of variable)

    angles = [n / float(N) * 2 * pi for n in range(N)]

    angles += angles[:1]

    # Initialise the spider plot

    ax = plt.subplot(nrow, ncol, metric_idx + 1, polar=True)

    # If you want the first axis to be on top:

    ax.set_theta_offset(pi / 2)

    ax.set_theta_direction(-1)

    # Draw one axe per variable + add labels labels yet

    plt.xticks(angles[:-1], categories, color="grey", size=15)

    for label, i in zip(ax.get_xticklabels(), range(0, len(angles))):

        if i < len(angles) / 2:

            angle_text = angles[i] * (-180 / pi) + 90

            label.set_horizontalalignment("left")

        else:

            angle_text = angles[i] * (-180 / pi) - 90

            label.set_horizontalalignment("right")

        label.set_rotation(angle_text)

    # Draw ylabels

    ax.set_rlabel_position(0)

    plt.yticks([0.1, 0.3, 0.6], ["0.1", "0.3", "0.6"], color="grey", size=8)

    plt.ylim(0, 1.05)

    # Plot data

    ax.plot(angles, values, color=color, linewidth=2, linestyle="solid")

    ax.fill(angles, values, color=color, alpha=0.4)

    ax.grid(color="white")

    for ti, di in zip(angles, values):

        ax.text(

            ti, di - 0.02, "{0:.2f}".format(di), color="black", ha="center", va="center"

        )

    # Add a title

    t = adata_test_final.uns["sample_score"].columns[metric_idx]

    t = t.replace("_pred_score", "")

    plt.title(t, color="black", y=1.2, size=22)