import qiskit
from qiskit_aer import Aer, QasmSimulator
from qiskit.transpiler import PassManager
from qiskit.transpiler.passes import Decompose

from dc_qiskit_qml import QiskitOptions
from dc_qiskit_qml.encoding_maps import NormedAmplitudeEncoding
from dc_qiskit_qml.distance_based.hadamard import QmlHadamardNeighborClassifier
from dc_qiskit_qml.distance_based.hadamard.state import QmlGenericStateCircuitBuilder
from dc_qiskit_qml.distance_based.hadamard.state.sparsevector import MottonenStatePreparation
from dc_qiskit_algorithms.MottonenStatePreparation import MottonenStatePreparationGate

initial_state_builder = QmlGenericStateCircuitBuilder(MottonenStatePreparation())

execution_backend: QasmSimulator = Aer.get_backend('qasm_simulator')
options = QiskitOptions(shots=8192,optimization_level=3)
qml = QmlHadamardNeighborClassifier(
    backend=execution_backend, 
    classifier_circuit_factory=initial_state_builder,
    encoding_map=NormedAmplitudeEncoding(),
    options=options
)

import numpy as np

from sklearn.preprocessing import StandardScaler, Normalizer
from sklearn.pipeline import Pipeline
from sklearn.datasets import load_iris

X, y = load_iris(return_X_y=True)
X = np.asarray([x[0:2] for x, y in zip(X, y) if y != 2])
y = np.asarray([y for x, y in zip(X, y) if y != 2])

X_train = X[[33, 85]]
y_train = y[[33, 85]]

X_test = X[[28, 36]]
y_test = y[[28, 36]]

pipeline = Pipeline([
    ('scaler',  StandardScaler()), 
    ('l2norm', Normalizer(norm='l2', copy=True)),
    ('qml', qml)
])

import matplotlib.pyplot as plt

_X = pipeline.fit_transform(X, y)
_X_train = pipeline.transform(X_train)
_X_test = pipeline.transform(X_test)


plt.scatter(
    _X[:,0], _X[:,1], 
    color=['red' if yy == 0 else 'blue' for yy in y],
    s=20)
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.xlim([-1.2,1.2])
plt.ylim([-1.2,1.2])
plt.show()

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

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

PassManager([Decompose(MottonenStatePreparationGate)]).run(qml._last_predict_circuits[0]).draw(fold=-1, output="mpl")

qiskit.transpile(
    circuits=qml._last_predict_circuits[0], 
    basis_gates=['cx', 'u1', 'u2', 'u3'], 
    optimization_level=0
).draw(fold=-1, output="mpl")

qiskit.transpile(
    circuits=qml._last_predict_circuits[0], 
    basis_gates=['cx', 'u1', 'u2', 'u3'], 
    optimization_level=3
).draw(fold=-1, output="mpl")

X_train = X[45:61]
y_train = y[45:61]

X_test = X[[33, 63]]
y_test = y[[33, 63]]

_X = pipeline.fit_transform(X, y)
_X_train = pipeline.transform(X_train)
_X_test = pipeline.transform(X_test)

plt.scatter(
    _X[:,0], _X[:,1], 
    color=['red' if yy == 0 else 'blue' for yy in y], s=10)
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.show()

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

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

PassManager([Decompose(MottonenStatePreparationGate)]).run(qml._last_predict_circuits[0]).draw(fold=-1, output="mpl")

qiskit.transpile(circuits=qml._last_predict_circuits[0], optimization_level=0, backend=execution_backend).draw(fold=120, output="mpl")