import qiskit
from qiskit_aer.backends.aerbackend import AerBackend

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

initial_state_builder = QmlGenericStateCircuitBuilder(MottonenStatePreparation())

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

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

X, y = load_wine(return_X_y=True)

from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.10, random_state=42)

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

pipeline.fit(X_train, y_train)
prediction = pipeline.predict(X_test)

prediction

"Test Accuracy: {}".format(
    sum([1 if p == t else 0 for p, t in zip(prediction, y_test)])/len(prediction)
)

prediction_train = pipeline.predict(X_train)

"Train Accuracy: {}".format(
    sum([1 if p == t else 0 for p, t in zip(prediction_train, y_train)])/len(prediction_train)
)

import matplotlib.pyplot as plt

colors = ['red', 'blue', 'green', 'orange']

plt.scatter(
    pipeline.transform(X_train[:,0]), pipeline.transform(X_train[:,1]), 
    color=[colors[yy] for yy in y_train],
    marker='.', s=50)
plt.show()

plt.scatter(
    pipeline.transform(X_test[:,0]), pipeline.transform(X_test[:,1]), 
    color=[colors[yy] for yy in prediction],
    marker='.', s=50)
plt.show()

plt.scatter(
    pipeline.transform(X_test[:,0]), pipeline.transform(X_test[:,1]), 
    color=[colors[yy] for yy in y_test],
    marker='.', s=20)
plt.show()

_X_train = pipeline.transform(X_train)
_X_test = pipeline.transform(X_test)

for i in range(len(_X_test)):
    p_acc_theory = QmlHadamardNeighborClassifier.p_acc_theory(X_train, y_train, X_test[i])
    print(f"{qml.last_predict_p_acc[i]:.4f} ~~ {p_acc_theory:.4f}")

for i in range(len(X_test)):
    p_label_theory = QmlHadamardNeighborClassifier.p_label_theory(X_train, y_train, X_test[i], prediction[i])
    print(f"{qml.last_predict_probability[i]:.4f} ~~ {p_label_theory:.4f}")

print(qml._last_predict_circuits[0].qasm())