# Mount Google Drive
from google.colab import drive # import drive from google colab
 
ROOT = "/content/drive"     # default location for the drive
print(ROOT)                 # print content of ROOT (Optional)
 
drive.mount(ROOT)           # we mount the google drive at /content/drive

!pip install qiskit
from IPython.display import clear_output
clear_output()

import os

def restart_runtime():
  os.kill(os.getpid(), 9)
restart_runtime()

# %matplotlib inline
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable

import numpy as np
import tensorflow as tf
from tensorflow.keras.utils import to_categorical

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()

x_train_flatten = x_train.reshape(x_train.shape[0], x_train.shape[1]*x_train.shape[2])/255.0
x_test_flatten = x_test.reshape(x_test.shape[0], x_test.shape[1]*x_test.shape[2])/255.0

# Sanity check

print(x_train_flatten.shape, y_train.shape)
print(x_test_flatten.shape, y_test.shape)

x_train_0 = x_train_flatten[y_train == 0]
x_train_1 = x_train_flatten[y_train == 1]
x_train_2 = x_train_flatten[y_train == 2]
x_train_3 = x_train_flatten[y_train == 3]
x_train_4 = x_train_flatten[y_train == 4]
x_train_5 = x_train_flatten[y_train == 5]
x_train_6 = x_train_flatten[y_train == 6]
x_train_7 = x_train_flatten[y_train == 7]
x_train_8 = x_train_flatten[y_train == 8]
x_train_9 = x_train_flatten[y_train == 9]

x_train_list = [x_train_0, x_train_1, x_train_2, x_train_3, x_train_4, x_train_5, x_train_6, x_train_7, x_train_8, x_train_9]

print(x_train_0.shape)
print(x_train_1.shape)
print(x_train_2.shape)
print(x_train_3.shape)
print(x_train_4.shape)
print(x_train_5.shape)
print(x_train_6.shape)
print(x_train_7.shape)
print(x_train_8.shape)
print(x_train_9.shape)

x_test_0 = x_test_flatten[y_test == 0]
x_test_1 = x_test_flatten[y_test == 1]
x_test_2 = x_test_flatten[y_test == 2]
x_test_3 = x_test_flatten[y_test == 3]
x_test_4 = x_test_flatten[y_test == 4]
x_test_5 = x_test_flatten[y_test == 5]
x_test_6 = x_test_flatten[y_test == 6]
x_test_7 = x_test_flatten[y_test == 7]
x_test_8 = x_test_flatten[y_test == 8]
x_test_9 = x_test_flatten[y_test == 9]

x_test_list = [x_test_0, x_test_1, x_test_2, x_test_3, x_test_4, x_test_5, x_test_6, x_test_7, x_test_8, x_test_9]

print(x_test_0.shape)
print(x_test_1.shape)
print(x_test_2.shape)
print(x_test_3.shape)
print(x_test_4.shape)
print(x_test_5.shape)
print(x_test_6.shape)
print(x_test_7.shape)
print(x_test_8.shape)
print(x_test_9.shape)

def binary_classification_data_generator(num_sample, mult_test, n_class, class_list):
    X_train = x_train_list[class_list[0]][:num_sample, :]
    X_test = x_test_list[class_list[0]][:int(mult_test*num_sample), :]

    Y_train = np.zeros((n_class*X_train.shape[0],), dtype=int)
    Y_test = np.zeros((n_class*X_test.shape[0],), dtype=int)

    for i in range(n_class-1):
      X_train = np.concatenate((X_train, x_train_list[class_list[i+1]][:num_sample, :]), axis=0)
      Y_train[num_sample*(i+1):num_sample*(i+2)] = int(i+1)

      X_test = np.concatenate((X_test, x_test_list[class_list[i+1]][:int(mult_test*num_sample), :]), axis=0)
      Y_test[int(mult_test*num_sample*(i+1)):int(mult_test*num_sample*(i+2))] = int(i+1)


    Y_train = to_categorical(Y_train)
    Y_test = to_categorical(Y_test)

    return X_train, Y_train, X_test, Y_test

num_sample = 200
n_class = 2
mult_test = 2.5

X_train01, Y_train01, X_test01, Y_test01 = binary_classification_data_generator(num_sample, mult_test, n_class, [0,1])

# Sanity check
print(X_train01.shape, Y_train01.shape)
print(X_test01.shape, Y_test01.shape)

plt.imshow

plt.imshow(np.concatenate((X_train01[0].reshape(28,28), X_train01[-1].reshape(28,28)), axis=1), cmap='gray')

X_train36, Y_train36, X_test36, Y_test36 = binary_classification_data_generator(num_sample, mult_test, n_class, [3,6])

# Sanity check
print(X_train36.shape, Y_train36.shape)
print(X_test36.shape, Y_test36.shape)

plt.imshow

plt.imshow(np.concatenate((X_train36[0].reshape(28,28), X_train36[-1].reshape(28,28)), axis=1), cmap='gray')

def normalize(X, use_params=False, params=None):
    """Normalize the given dataset X
    Args:
        X: ndarray, dataset
    
    Returns:
        (Xbar, mean, std): tuple of ndarray, Xbar is the normalized dataset
        with mean 0 and standard deviation 1; mean and std are the 
        mean and standard deviation respectively.
    
    Note:
        You will encounter dimensions where the standard deviation is
        zero, for those when you do normalization the normalized data
        will be NaN. Handle this by setting using `std = 1` for those 
        dimensions when doing normalization.
    """
    if use_params:
        mu = params[0]
        std_filled = [1]
    else:
        mu = np.mean(X, axis=0)
        std = np.std(X, axis=0)
        #std_filled = std.copy()
    #std_filled[std==0] = 1.
    Xbar = (X - mu)/(std + 1e-8)
    return Xbar, mu, std


X_train01, mu_train01, std_train01 = normalize(X_train01)
X_test01 = (X_test01 - mu_train01)/(std_train01 + 1e-8)

# Sanity check
print(X_train01.shape, Y_train01.shape)
print(X_test01.shape, Y_test01.shape)

X_train36, mu_train36, std_train36 = normalize(X_train36)
X_test36 = (X_test36 - mu_train36)/(std_train36 + 1e-8)

# Sanity check
print(X_train36.shape, Y_train36.shape)
print(X_test36.shape, Y_test36.shape)

from sklearn.decomposition import PCA
from matplotlib import pyplot as plt

num_component = 6
pca01 = PCA(n_components=num_component, svd_solver='full')
pca36 = PCA(n_components=num_component, svd_solver='full')

pca01.fit(X_train01)
pca36.fit(X_train36)

print(np.cumsum(pca01.explained_variance_ratio_))
print(np.cumsum(pca36.explained_variance_ratio_))

X_train01 = pca01.transform(X_train01)
X_test01 = pca01.transform(X_test01)

# Sanity check
print(X_train01.shape, Y_train01.shape)
print(X_test01.shape, Y_test01.shape)

X_train36 = pca36.transform(X_train36)
X_test36 = pca36.transform(X_test36)

# Sanity check
print(X_train36.shape, Y_train36.shape)
print(X_test36.shape, Y_test36.shape)

X_train01 = (X_train01.T / np.sqrt(np.sum(X_train01 ** 2, -1))).T
X_test01 = (X_test01.T / np.sqrt(np.sum(X_test01 ** 2, -1))).T

plt.scatter(X_train01[:100, 0], X_train01[:100, 1])
plt.scatter(X_train01[100:200, 0], X_train01[100:200, 1])
plt.scatter(X_train01[200:300, 0], X_train01[200:300, 1])

X_train36 = (X_train36.T / np.sqrt(np.sum(X_train36 ** 2, -1))).T
X_test36 = (X_test36.T / np.sqrt(np.sum(X_test36 ** 2, -1))).T

plt.scatter(X_train36[:100, 0], X_train36[:100, 1])
plt.scatter(X_train36[100:200, 0], X_train36[100:200, 1])
plt.scatter(X_train36[200:300, 0], X_train36[200:300, 1])

from qiskit import *
import numpy as np

# Set a random seed
np.random.seed(42)

def drc_circuit(x_sample, parameters, num_features, layer=1):
    q_reg = QuantumRegister(1)
    c_reg = ClassicalRegister(1)
    circuit = QuantumCircuit(q_reg, c_reg)
    
    # Layer iteration
    for i in range(layer):
        # Gate iteration
        for j in range(int(num_features/3)):
            circuit.u(*(parameters[3*j:3*(j+1)] * x_sample[3*j:3*(j+1)] + parameters[3*j+num_features:3*(j+1)+num_features]), 0)
    
    circuit.measure(0,0)
    
    return circuit

num_features = 6
num_layer = 1

params = np.random.uniform(size=(2*num_features*num_layer + 2))

drc_circuit(X_train01[0], params, num_features).draw()

def weighted_fidelity_cost_function(params, X_train, Y_train, backend, L=1, num_features=6, shots=2000):
    loss = 0
    
    # Training sample iteration
    for k in range(len(X_train)):
        circuit = drc_circuit(X_train[k], params, num_features, layer=L)
        
        t_circ = transpile(circuit, backend)
        qobj = assemble(t_circ, shots=shots)
        job = backend.run(qobj)
        result = job.result().get_counts(circuit)

        counts = np.array(list(result.values()))
        # Pauli-Z eigenvalue
        states = np.array([0, -1])
        # Compute probabilities for each state
        probabilities = counts / shots
        # Get state expectation
        z_expectation = np.sum(states * probabilities)
        # Get weighted fidelity
        fidel_0 = params[-2]*(1 + z_expectation)/2
        fidel_1 = params[-1]*(1 - z_expectation)/2
        # Calculate the loss from the sample
        loss += ((fidel_0 - Y_train[k,0])**2 + (fidel_1 - Y_train[k,1])**2)/2

    loss = loss/len(X_train)
    
    return loss


def drc_pred(params, X, backend, L=1, num_features=6, shots=2000):
    pred = []

    # Sample iteration
    for k in range(len(X)):
        circuit = drc_circuit(X[k], params, num_features, layer=L)
        
        t_circ = transpile(circuit, backend)
        qobj = assemble(t_circ, shots=shots)
        job = backend.run(qobj)
        result = job.result().get_counts(circuit)

        counts = np.array(list(result.values()))
        # Pauli-Z eigenvalue
        states = np.array([0, -1])
        # Compute probabilities for each state
        probabilities = counts / shots
        # Get state expectation
        z_expectation = np.sum(states * probabilities)
        # Get weighted fidelity
        fidel_0 = params[-2]*(1 + z_expectation)/2
        fidel_1 = params[-1]*(1 - z_expectation)/2

        pred += [[fidel_0, fidel_1]]

    return np.array(pred)

shots=2000
backend = Aer.get_backend('qasm_simulator')
num_layer = 1
num_features = 6
params = np.random.uniform(size=(2*num_features*num_layer + 2))

from qiskit.aqua.components.optimizers import COBYLA
import scipy

opt_params_01 = scipy.optimize.minimize(fun=weighted_fidelity_cost_function, x0=params, method='COBYLA',
                                            args=(X_train01, Y_train01, backend, num_layer, num_features, shots),
                                            options={'maxiter':1000})

print(opt_params_01)

Y_pred_train = drc_pred(opt_params_01['x'], X_train01, backend, L=num_layer, num_features=num_features, shots=shots)
Y_pred_test = drc_pred(opt_params_01['x'], X_test01, backend, L=num_layer, num_features=num_features, shots=shots)

train_acc = (np.argmax(Y_pred_train, axis=1) == np.argmax(Y_train01, axis=1)).sum()/len(Y_train01)
test_acc = (np.argmax(Y_pred_test, axis=1) == np.argmax(Y_test01, axis=1)).sum()/len(Y_test01)

print("Train Accuracy:", train_acc)
print("Test Accuracy:", test_acc)

train_acc_list01 = []
test_acc_list01 = []
train_loss_list01 = []
test_loss_list01 = []

for i in range(5):
    opt_params_01 = scipy.optimize.minimize(fun=weighted_fidelity_cost_function, x0=params, method='COBYLA',
                                            args=(X_train01, Y_train01, backend, i+1, num_features, shots),
                                            options={'maxiter':1000})
    
    print("Number of layer: " + str(i+1) + ", status: finished.")
    
    Y_pred_train = drc_pred(opt_params_01['x'], X_train01, backend, L=num_layer, num_features=num_features, shots=shots)
    Y_pred_test = drc_pred(opt_params_01['x'], X_test01, backend, L=num_layer, num_features=num_features, shots=shots)

    train_acc = (np.argmax(Y_pred_train, axis=1) == np.argmax(Y_train01, axis=1)).sum()/len(Y_train01)
    test_acc = (np.argmax(Y_pred_test, axis=1) == np.argmax(Y_test01, axis=1)).sum()/len(Y_test01)

    train_acc_list01 += [train_acc]
    test_acc_list01 += [test_acc]

    train_loss = weighted_fidelity_cost_function(opt_params_01['x'], X_train01, Y_train01, backend, L=num_layer, num_features=num_features, shots=shots)
    test_loss = weighted_fidelity_cost_function(opt_params_01['x'], X_test01, Y_test01, backend, L=num_layer, num_features=num_features, shots=shots)

    train_loss_list01 += [train_loss]
    test_loss_list01 += [test_loss]

plt.plot(train_acc_list01, label='Train Acc')
plt.plot(test_acc_list01, label='Test Acc')

plt.xlabel('Number of Layer')
plt.ylabel('Classification Accuracy')
plt.legend()
plt.show()

plt.plot(train_loss_list01, label='Train Loss')
plt.plot(test_loss_list01, label='Test Loss')

plt.xlabel('Number of Layer')
plt.ylabel('Loss')
plt.legend()
plt.show()

import pandas as pd

d = {'Training Acc': [0.9875, 0.9175], 'Testing Acc': [0.992, 0.915]}
pd.DataFrame(data=d, index=["Qiskit (COBYLA, QASM Simulator)", "PennyLane (Adam, PennyLane's Statevector Simulator)"])

opt_params_36 = scipy.optimize.minimize(fun=weighted_fidelity_cost_function, x0=params, method='COBYLA',
                                            args=(X_train36, Y_train36, backend, num_layer, num_features, shots),
                                            options={'maxiter':1000})

print(opt_params_36)

Y_pred_train = drc_pred(opt_params_36['x'], X_train36, backend, L=num_layer, num_features=num_features, shots=shots)
Y_pred_test = drc_pred(opt_params_36['x'], X_test36, backend, L=num_layer, num_features=num_features, shots=shots)

train_acc = (np.argmax(Y_pred_train, axis=1) == np.argmax(Y_train36, axis=1)).sum()/len(Y_train36)
test_acc = (np.argmax(Y_pred_test, axis=1) == np.argmax(Y_test36, axis=1)).sum()/len(Y_test36)

print("Train Accuracy:", train_acc)
print("Test Accuracy:", test_acc)

