import numpy as np
from sklearn.datasets import load_iris
from sklearn.preprocessing import StandardScaler, Normalizer
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline

import qiskit

from dc_qiskit_qml.feature_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 MöttönenStatePreparation

X, y = load_iris(True)
# Only the first two features and only get two labels
# This is a toy example!
X = np.asarray([x[0:2] for x, yy in zip(X, y) if yy != 2])
y = np.asarray([yy for x, yy in zip(X, y) if yy != 2])

preprocessing_pipeline = Pipeline([
    ('scaler',  StandardScaler()),
    ('l2norm', Normalizer(norm='l2', copy=True))
])
X = preprocessing_pipeline.fit_transform(X, y)

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

# Using the generic wave function (state vector) routine using the 'Möttönen'
# state preparation algorithm
initial_state_builder = QmlGenericStateCircuitBuilder(MöttönenStatePreparation())

# The normed amplitude encoding ensures that the data is normalized
# This is a somewhat unnecessary step as above we do that already
feature_map = NormedAmplitudeEncoding()

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

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

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

prediction_train = qml.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)
)