from datasets import *
from qiskit import Aer
from qiskit_aqua.input import SVMInput
from qiskit_aqua import run_algorithm
import numpy as np

n = 2  # dimension of each data point
sample_Total, training_input, test_input, class_labels = Wine(training_size=40,
                                                              test_size=10, n=n, PLOT_DATA=True)

temp = [test_input[k] for k in test_input]
total_array = np.concatenate(temp)

aqua_dict = {
    'problem': {'name': 'svm_classification', 'random_seed': 10598},
    'algorithm': {
        'name': 'QSVM.Kernel'
    },
    'feature_map': {'name': 'SecondOrderExpansion', 'depth': 2, 'entangler_map': {0: [1]}},
    'multiclass_extension': {'name': 'AllPairs'},
    'backend': {'shots': 1024}
}

backend = Aer.get_backend('qasm_simulator')
algo_input = SVMInput(training_input, test_input, total_array)
result = run_algorithm(aqua_dict, algo_input, backend=backend)
for k,v in result.items():
    print("'{}' : {}".format(k, v))
