"""
The Haxby dataset: different multi-class strategies
===================================================

In this example,
we compare ``one vs all`` and ``one vs one`` multi-class strategies:
We compare their overall cross-validated accuracy and their confusion matrix.

See the scikit-learn documentation
about `multiclass classification <https://scikit-learn.org/stable/modules/multiclass.html>`_.

"""

import numpy as np
import pandas as pd

from nilearn import datasets

# %%
# Load and prepare the Haxby dataset
# ----------------------------------
# By default 2nd subject from haxby datasets will be fetched.
haxby_dataset = datasets.fetch_haxby()

func_filename = haxby_dataset.func[0]
mask_filename = haxby_dataset.mask

print(f"Mask nifti images are located at: {mask_filename}")
print(f"Functional nifti images are located at: {func_filename}")


# %%
# We load the behavioral data that we will predict and
# remove the ``"rest"`` condition, as it is of no interest to us.
#
# .. seealso::
#
#   For more information about the dataset
#   see its :ref:`description <haxby_dataset>`.
#

labels = pd.read_csv(haxby_dataset.session_target[0], sep=" ")

non_rest = labels["labels"] != "rest"

y = labels["labels"][non_rest]

run = labels["chunks"]
n_runs = len(np.unique(run))


# %%
# We extract the data with a NiftiMasker.
# For decoding, standardizing is often very important,
# so we set ``standardize="zscore_sample"``.

from nilearn.maskers import NiftiMasker

nifti_masker = NiftiMasker(
    mask_img=mask_filename,
    runs=run,
    smoothing_fwhm=4,
    standardize="zscore_sample",
    memory="nilearn_cache",
    memory_level=1,
    verbose=1,
)
X = nifti_masker.fit_transform(func_filename)

# %%
# Remove the ``"rest"`` condition
X = X[non_rest]
run = run[non_rest]

# %%
# Build the decoders, using scikit-learn
# --------------------------------------
# Nilearn does not have dedicated multiclass estimators,
# as we can directly use those from sklearn.
# Here we use a Support Vector Classification (:class:`sklearn.svm.SVC`),
# with a linear kernel, and a simple feature selection step.

from sklearn.feature_selection import SelectKBest, f_classif
from sklearn.multiclass import OneVsOneClassifier, OneVsRestClassifier
from sklearn.pipeline import Pipeline
from sklearn.svm import SVC

# %%
svc_ovo = OneVsOneClassifier(
    Pipeline(
        [
            ("anova", SelectKBest(f_classif, k=500)),
            ("svc", SVC(kernel="linear")),
        ]
    )
)
svc_ovo

# %%
svc_ova = OneVsRestClassifier(
    Pipeline(
        [
            ("anova", SelectKBest(f_classif, k=500)),
            ("svc", SVC(kernel="linear")),
        ]
    )
)
svc_ova

# %%
# Now we compute cross-validation scores
# --------------------------------------
# The :term:`fMRI` data is acquired by runs,
# and the noise is autocorrelated in a given run.
# Hence, it is better to predict across runs when doing cross-validation.
# Here we do a 5 fold cross-validation.

from sklearn.model_selection import cross_val_score

cv = 5

# %%
cv_scores_ovo = cross_val_score(svc_ovo, X, y, cv=cv, verbose=1)
cv_scores_ovo

# %%
cv_scores_ova = cross_val_score(svc_ova, X, y, cv=cv, verbose=1)
cv_scores_ova

# %%
print("OvO:", cv_scores_ovo.mean().round(decimals=3))
print("OvA:", cv_scores_ova.mean().round(decimals=3))

# %%
# Plot barplots of the prediction scores
# --------------------------------------

from matplotlib import pyplot as plt

from nilearn.plotting import show

plt.figure(figsize=(4, 3))
plt.boxplot([cv_scores_ova, cv_scores_ovo])
plt.xticks([1, 2], ["One vs All", "One vs One"])
plt.title("Prediction: accuracy score")

show()

# %%
# Plot the cross-validated confusion matrices
# --------------------------------------------
# Instead of fitting on a single train/test split, we use
# :func:`sklearn.model_selection.cross_val_predict` with the same
# cross-validation as above to get, for each sample,
# a prediction made by a model that never saw that sample during training.
# We can then build a single confusion matrix from those out-of-fold
# predictions for the whole dataset.

from sklearn.model_selection import cross_val_predict

y_pred_ovo = cross_val_predict(svc_ovo, X, y, cv=cv, verbose=1)

y_pred_ova = cross_val_predict(svc_ova, X, y, cv=cv, verbose=1)

# %%
# We get the labels of the numerical conditions represented by the vector y
# and we sort the conditions by the order of appearance.

from sklearn.metrics import confusion_matrix

from nilearn.plotting import plot_matrix

unique_conditions, order = np.unique(y, return_index=True)
unique_conditions = unique_conditions[np.argsort(order)]

im = plot_matrix(
    confusion_matrix(y_pred_ovo, y),
    labels=unique_conditions,
    title="Confusion matrix: One vs One",
    cmap="inferno",
    figure=(6, 5),
    auto_fit=False,
    vmax=108,
)
ax = im.axes
ax.set_ylabel("True label")
ax.set_xlabel("Predicted label")

im = plot_matrix(
    confusion_matrix(y_pred_ova, y),
    labels=unique_conditions,
    title="Confusion matrix: One vs All",
    cmap="inferno",
    figure=(6, 5),
    auto_fit=False,
    vmax=108,
)
ax = im.axes
ax.set_ylabel("True label")
ax.set_xlabel("Predicted label")

show()

# sphinx_gallery_dummy_images=3