!pip install pyRiemann-qiskit
!pip install moabb

from pyriemann_qiskit.pipelines import (
    QuantumMDMWithRiemannianPipeline,
)
from moabb.datasets import BI2012
from moabb.paradigms import P300
from sklearn.model_selection import train_test_split
from sklearn.metrics import balanced_accuracy_score
from sklearn.preprocessing import LabelEncoder
from time import time

# Instantiate dataset
# This has to be adapted for the data you used
dataset = BI2012() #selected dataset
# Instantiate paradigm:
# It defines the pre-processing steps for data extraction
# This is specific to the MOABB library
paradigm = P300()
# Get first subject's data
X, y, _ = paradigm.get_data(dataset, subjects=[1])

# Encode class label into numeric value
label_encoder = LabelEncoder()
y = label_encoder.fit_transform(y)

# Let's keep it simple without cross-validation
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=100, stratify=y)

# Instantiate our pipeline (inherits from sklearn mixins)
clf = QuantumMDMWithRiemannianPipeline(
    convex_metric="distance", quantum=True
)
display(clf._pipe)

start_time = time()
clf.fit(X_train, y_train)
fit_time = time() - start_time

# Perform classification (prediction)
# This use QAOA and a Aer/statevector simulator.
# An account token can also be provided to the estimator
# to run on real quantum backend
start_time = time()
y_pred = clf.predict(X_test)
pred_time = time() - start_time

# Compute dataset balance
n_class_0 = y[y == 1].shape[0]
n_class_1 = y[y == 0].shape[0]

# Our dataset is not balanced
balance = n_class_0 / n_class_1
print(f'balance = {balance}')

# Then we will use the balanced-accuracy
# which is a better metric than accuracy in this case:
score = balanced_accuracy_score(y_test, y_pred)
# Do not pay to much attention to the score.
# There is no cross-validation here.
# It likely just depends on the random_state.
print(f'score = {score}')

# Print execution time
# The fitting is classical, the prediction (here on simulator) is quantum.
# Other versions of the algorithm exist (with fitting quantum and prediction classical)
print("Fit time: {:.3f}s".format(fit_time))
print("Prediction time: {:.3f}s".format(pred_time))

# Time required for the classfication of a single trial:
n_train = y_train.shape[0]
n_test = y_test.shape[0]
print("Fit time (1 example): {:.3f}s".format(fit_time / n_train))
print("Prediction time (1 example): {:.3f}s".format(pred_time / n_test))