from datasets import *
from qiskit_aqua.utils import split_dataset_to_data_and_labels
from qiskit_aqua.input import SVMInput
from qiskit_aqua import run_algorithm
import numpy as np

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

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

aqua_dict = {
    'problem': {'name': 'svm_classification'},
    'algorithm': {
        'name': 'SVM'
    },
    'multiclass_extension': {'name': 'OneAgainstRest'}
}

algo_input = SVMInput(training_input, test_input, total_array)

extensions = [
   {'name': 'OneAgainstRest'},
   {'name': 'AllPairs'}, 
   {'name': 'ErrorCorrectingCode', 'code_size': 5}
]

for extension in extensions:
    aqua_dict['multiclass_extension'] = extension
    result = run_algorithm(aqua_dict, algo_input)
    print("\n----- Using multiclass extension: '{}' -----\n".format(extension['name']))
    for k,v in result.items():
        print("'{}' : {}".format(k, v))
