from matplotlib import pyplot as plt
import numpy as np
from qiskit import Aer, execute
from qiskit.tools import job_monitor
from qiskit.aqua import QuantumInstance
from qiskit.circuit.library import ZZFeatureMap, PauliFeatureMap, ZFeatureMap
from qiskit.aqua.algorithms import QSVM
from qiskit.aqua.components.multiclass_extensions import OneAgainstRest
from qiskit.aqua.utils import split_dataset_to_data_and_labels, map_label_to_class_name
from qiskit.aqua import MissingOptionalLibraryError
from sklearn import datasets
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, MinMaxScaler
from sklearn.decomposition import PCA
import pandas as pd

data = datasets.load_wine()
target = data.target
data = data.data

training_dataset_size = 10
testing_dataset_size = 10
class_labels = ["A", "B", "C"]

pd.DataFrame(data)

dim = 3
training_dataset_size = 20
testing_dataset_size = 10

sample_train, sample_test, label_train, label_test = train_test_split(data, target, random_state=10)

# Standardizing the dataset -- gaussian with unit variance and 0 mean
std = StandardScaler()
std.fit(sample_train)
sample_train = std.transform(sample_train)
sample_test = std.transform(sample_test)

# Using PCA to reduce no. of dimensions to 2 to match number of qubits
pca = PCA(dim)
pca.fit(sample_train)
sample_train = pca.transform(sample_train)
sample_test = pca.transform(sample_test)

# Scaling data to range from -1 to 1
minmax = MinMaxScaler()
minmax = minmax.fit(sample_train)
sample_train = minmax.transform(sample_train)
sample_test = minmax.transform(sample_test)

# Setting dataset size to number of training_dataset_size
training_input = {key: (sample_train[label_train == k, :])[:training_dataset_size] for k, key in enumerate(class_labels)}
test_input = {key: (sample_test[label_test == k, :])[:testing_dataset_size] for k, key in enumerate(class_labels)}

# Plotting data
for k in range(0, 3):
    plt.scatter(sample_train[label_train == k, 0][:training_dataset_size],
                sample_train[label_train == k, 1][:training_dataset_size])

plt.title("Wine dataset with reduced dimensions")
plt.show()

datapoints, class_to_label = split_dataset_to_data_and_labels(test_input)

backend = Aer.get_backend("qasm_simulator")
feature_map = ZZFeatureMap(dim, reps=2)
print(feature_map)

svm = QSVM(feature_map, training_input, test_input, None, multiclass_extension=OneAgainstRest())
quantum_instance = QuantumInstance(backend, shots=8000)
result = svm.run(quantum_instance)

predicted_labels = svm.predict(datapoints[0])
predicted_classes = map_label_to_class_name(predicted_labels, svm.label_to_class)
print(f"Ground truth: {datapoints[1]}")
print(f"Prediction: {predicted_labels}")
print(" RESULTS Testing success ratio: ", result['testing_accuracy'])