# TODO: uncomment the next line after release qlearnkit 0.2.0
#!pip install qlearnkit
!pip install matplotlib

import numpy as np
from qlearnkit.algorithms.qsvm import QSVClassifier
from sklearn.preprocessing import MinMaxScaler
from sklearn.datasets import load_iris
from sklearn.svm import SVC
from matplotlib import pyplot as plt
from qiskit import BasicAer
from qiskit.circuit.library import PauliFeatureMap
from qiskit.utils import QuantumInstance

# import some data to play with
iris = load_iris()
mms = MinMaxScaler()
X = iris.data[:, :2]  # we only take the first two features. We could
                      # avoid this ugly slicing by using a two-dim dataset
X = mms.fit_transform(X)
y = iris.target

seed = 42
encoding_map = PauliFeatureMap(2)
quantum_instance = QuantumInstance(BasicAer.get_backend('statevector_simulator'),
                                   shots=1024,
                                   optimization_level=1,
                                   seed_simulator=seed,
                                   seed_transpiler=seed)

svc = SVC(kernel='linear')
qsvc = QSVClassifier(encoding_map=encoding_map, quantum_instance=quantum_instance)

svc.fit(X,y)
qsvc.fit(X,y)

h = 0.1 # step size in the mesh
# create a mesh to plot in
x_min, x_max = X[:, 0].min() - 0.2, X[:, 0].max() + 0.2
y_min, y_max = X[:, 1].min() - 0.2, X[:, 1].max() + 0.2
xx, yy = np.meshgrid(np.arange(x_min, x_max, h),
                     np.arange(y_min, y_max, h))

# title for the plots
titles = ['SVC with linear kernel',
          'QSVC with Pauli feature map']

for i, clf in enumerate((svc, qsvc)):
    # Plot the decision boundary. For that, we will assign a color to each
    # point in the mesh [x_min, x_max]x[y_min, y_max].
    plt.subplot(2, 1, i + 1)
    plt.subplots_adjust(wspace=0.4, hspace=0.4)
    Z = clf.predict(np.c_[xx.ravel(), yy.ravel()])

    # Put the result into a color plot
    Z = Z.reshape(xx.shape)
    plt.contourf(xx, yy, Z, cmap="RdBu", alpha=0.8)

    # Plot also the training points
    plt.scatter(X[:, 0], X[:, 1], c=y, cmap="RdBu",edgecolor="grey" )
    plt.xlabel('Sepal length')
    plt.ylabel('Sepal width')
    plt.xlim(xx.min(), xx.max())
    plt.ylim(yy.min(), yy.max())
    plt.xticks(())
    plt.yticks(())
    plt.title(titles[i])

plt.show()