# pylint: disable=import-error, wrong-import-position

# Python imports
import os
import sys

# External imports
from pylab import cm
from sklearn import metrics
import numpy as np
import matplotlib.pyplot as plt

%matplotlib inline

# Qiskit imports
from qiskit.visualization import circuit_drawer
from qiskit_algorithms.optimizers import SPSA
from qiskit_machine_learning.kernels import TrainableFidelityQuantumKernel
from qiskit_machine_learning.kernels.algorithms import QuantumKernelTrainer
from qiskit_machine_learning.algorithms import QSVC

# Put this repository on the Python path and import qkt pkgs
module_path = os.path.abspath(os.path.join("../.."))
sys.path.append(module_path)
from qkt.feature_maps import CovariantFeatureMap
from qkt.utils import train_test_split, QKTCallback

# Load the dataset and split into train and test sets
DATA_FILEPATH = "../../data/dataset_graph7.csv"
X_train, y_train, X_test, y_test = train_test_split(DATA_FILEPATH)

num_features = np.shape(X_train)[1]
entangler_map = [[0, 2], [3, 4], [2, 5], [1, 4], [2, 3], [4, 6]]

# Note that [[0,1],[2,3],[4,5],[6,7],[8,9],[1,2],[3,4],[5,6],[7,8]]
# is a suitable input for the 10-qubit dataset

fm = CovariantFeatureMap(
    feature_dimension=num_features,
    entanglement=entangler_map,
    single_training_parameter=True,
)
circuit_drawer(fm)

print(fm.training_parameters)

# Instantiate quantum kernel
quant_kernel = TrainableFidelityQuantumKernel(
    feature_map=fm, training_parameters=fm.training_parameters
)

# Set up the optimizer
cb_qkt = QKTCallback()
spsa_opt = SPSA(
    maxiter=5, callback=cb_qkt.callback, learning_rate=0.1, perturbation=0.1
)

# Instantiate a quantum kernel trainer.
qkt = QuantumKernelTrainer(
    quantum_kernel=quant_kernel,
    loss="svc_loss",
    optimizer=spsa_opt,
    initial_point=[0.1] * len(fm.training_parameters),
)

# Train the kernel using QKT directly
qka_results = qkt.fit(X_train, y_train)
optimized_kernel = qka_results.quantum_kernel

# Use QSVC for classification
qsvc = QSVC(quantum_kernel=optimized_kernel)

# Fit the QSVC
qsvc.fit(X_train, y_train)

# Predict the labels
labels_test = qsvc.predict(X_test)

# Evaluate the test accuracy
accuracy_test = metrics.balanced_accuracy_score(y_true=y_test, y_pred=labels_test)
print(f"accuracy test: {accuracy_test}")

plot_data = cb_qkt.get_callback_data()  # callback data
K = optimized_kernel.evaluate(
    X_train
)  # kernel matrix evaluated on the training samples

plt.rcParams["font.size"] = 20
fig, ax = plt.subplots(1, 2, figsize=(14, 5))
ax[0].plot(
    [i + 1 for i in range(len(plot_data[0]))], np.array(plot_data[2]), c="k", marker="o"
)
ax[0].set_xlabel("Iterations")
ax[0].set_ylabel("Loss")
ax[1].imshow(K, cmap=cm.get_cmap("bwr", 20))
ax[1].axis("off")
fig.tight_layout()
plt.show()

#import qiskit.tools.jupyter
#
#%qiskit_version_table
#%qiskit_copyright