# Classical computing modules
import numpy as np
import pandas as pd
import seaborn as sns
from statistics import mode
# Dataset
from sklearn import datasets, preprocessing 
from sklearn.model_selection import train_test_split
# Quantum computing modules
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, Aer, execute
from qiskit.visualization import plot_histogram
from qiskit.quantum_info import Statevector

# Load the dataset
iris = datasets.load_iris()
feature_labels = iris.feature_names
class_labels = iris.target_names

# Extract the data
X = iris.data
y = np.array([iris.target])
M = 4

# Print the features and classes
print("Features: ", feature_labels)
print("Classes: ", class_labels)

df = pd.DataFrame(data = np.concatenate((X, y.T), axis = 1), columns = feature_labels[0:M] + ["Species"])
sns.pairplot(df, hue = 'Species')

min_max_scaler = preprocessing.MinMaxScaler(feature_range=(0, 1))
min_max_scaler.fit(X)
X_normalized = min_max_scaler.transform(X)

df2 = pd.DataFrame(data = np.concatenate((X_normalized, y.T), axis = 1), columns = feature_labels[0:M] + ["Species"])
sns.pairplot(df2, hue = 'Species')

N_train = 100 # Training dataset size

X_train, X_test, y_train, y_test = train_test_split(X_normalized, y[0], train_size=N_train)

# Select a sample test sample for classification
test_index = int(np.random.rand()*len(X_test)) 
phi_test = X_test[test_index]

# Create the encoded feature vector
for i in range(M):
    phi_i = [np.sqrt(phi_test[i]), np.sqrt(1- phi_test[i])]
    if i == 0:
        phi = phi_i
    else:
        phi = np.kron(phi_i, phi)    

# Validate the statevector
print('Valid statevector: ', Statevector(phi).is_valid())

N = int(np.ceil(np.log2(N_train)))
psi = np.zeros(2**(M + N)) # Task 9

for i in range(N_train):
    # Encode |i>
    i_vec = np.zeros(2**N)
    i_vec[i] = 1

    # Encode |x>
    x = X_train[i, :]
    for j in range(M):
        dummy = [np.sqrt(x[j]), np.sqrt(1- x[j])]
        if j == 0:
            x_vec = dummy
        else:
            x_vec = np.kron(dummy, x_vec)
    psi_i = np.kron(x_vec, i_vec)
    # Task 9
    psi += psi_i
psi /= np.sqrt(N_train)

# Check the validity of the statevector
assert Statevector(psi).is_valid(), "The statevector is not square-normalized."

index_reg = QuantumRegister(N, 'i')
train_reg = QuantumRegister(M, 'train')
test_reg = QuantumRegister(M, 'test')
p = QuantumRegister(1, 'similarity')

# Create the quantum circuit
qknn = QuantumCircuit()

# Add registers to the circuit
qknn.add_register(index_reg)
qknn.add_register(train_reg)
qknn.add_register(test_reg)
qknn.add_register(p)

# Draw the quantum circuit
qknn.draw('mpl') 

# Initialize the training register
qknn.initialize(psi, index_reg[0:N] + train_reg[0:M])

# Initialize the test register
qknn.initialize(phi, test_reg)

# Draw qknn
qknn.draw('mpl')

from modules import swap_test

# Apply the Quantum SWAP Test module to the quantum circuit
qknn.append(swap_test(M), train_reg[0:M] + test_reg[0:M] + [p[0]])
# Draw qknn
qknn.draw('mpl')

# Create the classical register
meas_reg_len = N + 1
meas_reg = ClassicalRegister(meas_reg_len, 'meas')
qknn.add_register(meas_reg)

# Measure the qubits
qknn.measure(index_reg[0::] + p[0::], meas_reg)

# Draw qknn
qknn.draw('mpl', fold = -1)

backend = Aer.get_backend('qasm_simulator')
counts_knn = execute(qknn, backend, shots = 15000).result().get_counts()

result_arr = np.zeros((N_train, 3))
for count in counts_knn:
    i_dec = int(count[1::], 2)
    phase = int(count[0], 2)
    if phase == 0:
        result_arr[i_dec, 0] += counts_knn[count]
    else:
        result_arr[i_dec, 1] += counts_knn[count]

for i in range(N_train):        
    prob_1 = result_arr[i][1]/(result_arr[i][0] + result_arr[i][1])
    result_arr[i][2] = 1 - 2*prob_1

# Find the indexes of minimum distance
k = 10
k_min_dist_arr = result_arr[:, 2].argsort()[::-1][:k]

# Determine the class of the test sample
y_pred = mode(y_train[k_min_dist_arr])
y_exp = y_test[test_index]
print('Predicted class of the test sample is {}.'.format(y_pred))
print('Expected class of the test sample is {}.'.format(y_exp))

def q_knn_module(test_index, psi, k=10):
    phi_test = X_test[test_index]

    # Create the encoded feature vector
    for i in range(M):
        phi_i = [np.sqrt(phi_test[i]), np.sqrt(1- phi_test[i])]
        if i == 0:
            phi = phi_i
        else:
            phi = np.kron(phi, phi_i)    

    # Create the quantum circuit
    qknn = QuantumCircuit()

    # Add registers to circuit
    qknn.add_register(index_reg)
    qknn.add_register(train_reg)
    qknn.add_register(test_reg)
    qknn.add_register(p)

    # Initialize train reg
    qknn.initialize(psi, index_reg[0:N] + train_reg[0:M])

    # Initialize test reg
    qknn.initialize(phi, test_reg)

    # Apply the swap test module to the quantum circuit
    qknn.append(swap_test(M), train_reg[0:M] + test_reg[0:M] + [p[0]])

    # Create the classical register
    meas_reg_len = N+1
    meas_reg = ClassicalRegister(meas_reg_len, 'meas')
    qknn.add_register(meas_reg)

    # Measure the qubits
    qknn.measure(index_reg[0::] + p[0::], meas_reg)

    # Circuit execution
    backend = Aer.get_backend('qasm_simulator')
    counts_knn = execute(qknn, backend, shots = 10000).result().get_counts()

    # Decoding the results
    result_arr = np.zeros((N_train, 3))
    for count in counts_knn:
        i_dec = int(count[1::], 2)
        phase = int(count[0], 2)
        if phase == 0:
            result_arr[i_dec, 0] += counts_knn[count]
        else:
            result_arr[i_dec, 1] += counts_knn[count]

    # Computing similarity
    for i in range(N_train):        
        prob_1 = result_arr[i][1]/(result_arr[i][0] + result_arr[i][1])
        result_arr[i][2] = 1 - 2*prob_1

    # Find the indices of minimum distance
    k_min_dist_arr = result_arr[:, 2].argsort()[::-1][:k]

    # Determine the class of the test sample
    y_pred = mode(y_train[k_min_dist_arr])
    y_exp = y_test[test_index]

    return y_pred, y_exp
 
y_pred_arr = []
y_exp_arr = [] 

for test_indx in range(len(X_test)):
    y_pred, y_exp = q_knn_module(test_indx, psi)
    y_pred_arr.append(y_pred)
    y_exp_arr.append(y_exp)

t = 0
f = 0
for i in range(len(X_test)):
    if y_pred_arr[i] == y_exp_arr[i]:
        t += 1
    else:
        f += 1

print('Model accuracy is {}%.'.format(t/(t+f) *100))