from datasets import *
from qiskit import Aer
from qiskit_aqua.utils import split_dataset_to_data_and_labels, map_label_to_class_name
from qiskit_aqua import run_algorithm, QuantumInstance
from qiskit_aqua.input import SVMInput
from qiskit_aqua.components.feature_maps import SecondOrderExpansion
from qiskit_aqua.algorithms import QSVMKernel

feature_dim = 2 # dimension of each data point
training_dataset_size = 20
testing_dataset_size = 10
random_seed = 10598
shots = 1024

sample_Total, training_input, test_input, class_labels = ad_hoc_data(training_size=training_dataset_size, 
                                                                     test_size=testing_dataset_size, 
                                                                     n=feature_dim, gap=0.3, PLOT_DATA=False)

datapoints, class_to_label = split_dataset_to_data_and_labels(test_input)
print(class_to_label)

params = {
    'problem': {'name': 'svm_classification', 'random_seed': random_seed},
    'algorithm': {
        'name': 'QSVM.Kernel'
    },
    'backend': {'shots': shots},
    'feature_map': {'name': 'SecondOrderExpansion', 'depth': 2, 'entanglement': 'linear'}
}
backend = Aer.get_backend('qasm_simulator')
algo_input = SVMInput(training_input, test_input, datapoints[0])

result = run_algorithm(params, algo_input, backend=backend)

print("kernel matrix during the training:")
kernel_matrix = result['kernel_matrix_training']
img = plt.imshow(np.asmatrix(kernel_matrix),interpolation='nearest',origin='upper',cmap='bone_r')
plt.show()

print("testing success ratio: ", result['testing_accuracy'])
print("predicted classes:", result['predicted_classes'])

backend = Aer.get_backend('qasm_simulator')
feature_map = SecondOrderExpansion(num_qubits=feature_dim, depth=2, entangler_map={0: [1]})
svm = QSVMKernel(feature_map, training_input, test_input, None)# the data for prediction can be feeded later.
svm.random_seed = random_seed
quantum_instance = QuantumInstance(backend, shots=shots, seed=random_seed, seed_mapper=random_seed)
result = svm.run(quantum_instance)

print("kernel matrix during the training:")
kernel_matrix = result['kernel_matrix_training']
img = plt.imshow(np.asmatrix(kernel_matrix),interpolation='nearest',origin='upper',cmap='bone_r')
plt.show()

print("testing success ratio: ", result['testing_accuracy'])

predicted_labels = svm.predict(datapoints[0])

predicted_classes = map_label_to_class_name(predicted_labels, svm.label_to_class)
print("ground truth: {}".format(datapoints[1]))
print("preduction:   {}".format(predicted_labels))