from qiskit import Aer, QuantumCircuit
from qiskit.utils import QuantumInstance
from qiskit_machine_learning.neural_networks import CircuitQNN
from qiskit_machine_learning.algorithms.classifiers import NeuralNetworkClassifier
from qiskit.algorithms.optimizers import COBYLA
from qiskit.circuit.library import RealAmplitudes, ZZFeatureMap

from titanicibmq.titanic_data import *

# Install a pip package in the current Jupyter kernel
import sys
!{sys.executable} -m pip install sklearn

from sklearn.decomposition import PCA

# Load the data

datapath="data/"
data, __ = titanic(datapath=datapath)
data.sample(frac=1)

X = data.drop("Survived", axis=1).values
pca = PCA(n_components=2)
Xpca = pca.fit_transform(X)

split_ratio = 0.2

X_train = Xpca[int(data.shape[0]*split_ratio):,:]
X_test = Xpca[:int(data.shape[0]*split_ratio),:]

y_train = data.iloc[int(data.shape[0]*split_ratio):,:]["Survived"].values
y_test = data.iloc[:int(data.shape[0]*split_ratio),:]["Survived"].values

def get_interpret(num_classes):
        def parity(x, num_classes=num_classes):
            return '{:b}'.format(x).count('1') % num_classes
        return parity

def parity(x):
    return '{:b}'.format(x).count('1') % 2

quantum_instance = QuantumInstance(Aer.get_backend('statevector_simulator'), shots=100)
optimizer = COBYLA(maxiter=100)
feature_map = ZZFeatureMap(feature_dimension=X_train.shape[1], reps=1)
ansatz = RealAmplitudes(num_qubits=feature_map._num_qubits, reps=1)

qc = QuantumCircuit(feature_map._num_qubits)
qc.compose(feature_map, inplace=True)
qc.compose(ansatz, inplace=True)

qnn = CircuitQNN(circuit= qc, 
                input_params=feature_map.parameters, 
                weight_params=ansatz.parameters, 
                sparse=False, 
                sampling=False, 
                interpret=parity, 
                output_shape=len(np.unique(y_train, axis=0)), 
                gradient=None, 
                quantum_instance=quantum_instance)

cc = NeuralNetworkClassifier(neural_network=qnn,
                                optimizer=optimizer)

cc.fit(X_train, y_train)

# Model accuracy
acc_train = cc.score(X_train, y_train)
acc_test = cc.score(X_test, y_test)
print("Accuracy on training dataset: {}.".format(acc_train))
print("Accuracy on testing dataset: {}.".format(acc_test))