import qiskit

import numpy as np

X_train = np.asarray([[1.0, 1.0], [-1.0, 1.0], [-1.0, -1.0], [1.0, -1.0]])
y_train = [0, 1, 0, 1]

X_test = np.asarray([[0.2, 0.4], [0.4, -0.8]])
y_test = [0, 1]

import matplotlib.pyplot as plt

plt.scatter(
    X_train[:,0], X_train[:,1], 
    color=['red' if yy == 0 else 'blue' for yy in y_train],
    marker='x', s=200)
plt.scatter(
    X_test[:,0], X_test[:,1], 
    color=['red' if yy == 0 else 'blue' for yy in y_test],
    marker='o', s=200)
plt.hlines(0.0, -1.0, 1.0)
plt.vlines(0.0, -1.0, 1.0)
plt.show()

from dc_qiskit_qml.encoding_maps import EncodingMap
import scipy.sparse as sparse


class MyEncodingMap(EncodingMap):
        def map(self, input_vector: list) -> sparse.dok_matrix:
            result = sparse.dok_matrix((4,1))
            index = 0
            if input_vector[0] > 0 and input_vector[1] > 0:
                index = 0
            if input_vector[0] < 0 < input_vector[1]:
                index = 1
            if input_vector[0] < 0 and input_vector[1] < 0:
                index = 2
            if input_vector[0] > 0 > input_vector[1]:
                index = 3
            result[index, 0] = 1.0
            return result


encoding_map = MyEncodingMap()

list(zip(X_train, [encoding_map.map(x).keys() for x in X_train]))

list(zip(X_test, [encoding_map.map(x).keys() for x in X_test]))

from qiskit_aer.backends.aerbackend import AerBackend
from dc_qiskit_qml.distance_based.hadamard import QmlHadamardNeighborClassifier
from dc_qiskit_qml.distance_based.hadamard.state import QmlBinaryDataStateCircuitBuilder
from dc_qiskit_qml.distance_based.hadamard.state.cnot import CCXToffoli

initial_state_builder = QmlBinaryDataStateCircuitBuilder(CCXToffoli(), do_optimizations=True)

execution_backend: AerBackend = qiskit.Aer.get_backend('qasm_simulator')
qml = QmlHadamardNeighborClassifier(backend=execution_backend,
                                    shots=8192,
                                    encoding_map=encoding_map,
                                    classifier_circuit_factory=initial_state_builder)

qml.fit(X_train, y_train)
qml.predict(X_test), y_test

qiskit.IBMQ.load_account()

qiskit.IBMQ.get_provider().backends()

qiskit.Aer.backends()

b16 = qiskit.IBMQ.get_provider().get_backend('ibmq_16_melbourne')
b16.status()

execution_backend = b16
qml = QmlHadamardNeighborClassifier(backend=execution_backend,
                                    shots=8192,
                                    encoding_map=encoding_map,
                                    classifier_circuit_factory=initial_state_builder)
qml.fit(X_train, y_train)
async_qml_job = qml.predict_async(X_test)

async_qml_job.job.status(), async_qml_job.job.job_id()

async_qml_job.predict_result(), y_test

qml.last_predict_probability, qml.last_predict_p_acc

qml._last_predict_circuits[0].draw(output='mpl')

from dc_qiskit_qml.distance_based.hadamard.state.cnot import CCXMottonen

initial_state_builder = QmlBinaryDataStateCircuitBuilder(CCXMottonen())

execution_backend: AerBackend = qiskit.Aer.get_backend('qasm_simulator')
qml = QmlHadamardNeighborClassifier(backend=execution_backend,
                                    shots=8192,
                                    encoding_map=encoding_map,
                                    classifier_circuit_factory=initial_state_builder)

qml.fit(X_train, y_train)
qml.predict(X_test), y_test

qml._last_predict_circuits[0].draw(output='mpl', fold=-1)

execution_backend = b16
qml = QmlHadamardNeighborClassifier(backend=execution_backend,
                                    shots=8192,
                                    encoding_map=encoding_map,
                                    classifier_circuit_factory=initial_state_builder)

qml.fit(X_train, y_train)
async_qml_job.job.status(), async_qml_job.job.job_id()

qml.last_predict_probability, qml.last_predict_p_acc