# 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 pennylane
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

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)

num_sample = 200
n_class = 2
mult_test = 2.5

X_train = x_train_list[0][:num_sample, :]
X_test = x_test_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[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[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)

print(X_train.shape, Y_train.shape)
print(X_test.shape, Y_test.shape)

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_train, mu_train, std_train = normalize(X_train)
X_train.shape, Y_train.shape

X_test = (X_test - mu_train)/(std_train + 1e-8)
X_test.shape, Y_test.shape

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

num_component = 6
pca = PCA(n_components=num_component, svd_solver='full')

pca.fit(X_train)

np.cumsum(pca.explained_variance_ratio_)

X_train = pca.transform(X_train)
X_test = pca.transform(X_test)

print(X_train.shape, Y_train.shape)
print(X_test.shape, Y_test.shape)

X_train = (X_train.T / np.sqrt(np.sum(X_train ** 2, -1))).T
X_test = (X_test.T / np.sqrt(np.sum(X_test ** 2, -1))).T

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

import pennylane as qml
from pennylane import numpy as np
from pennylane.optimize import AdamOptimizer, GradientDescentOptimizer

qml.enable_tape()


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

# Define output labels as quantum state vectors

# def density_matrix(state):
#     """Calculates the density matrix representation of a state.

#     Args:
#         state (array[complex]): array representing a quantum state vector

#     Returns:
#         dm: (array[complex]): array representing the density matrix
#     """
#     return state * np.conj(state).T


label_0 = [[1], [0]]
label_1 = [[0], [1]]


def density_matrix(state):
    """Calculates the density matrix representation of a state.

    Args:
        state (array[complex]): array representing a quantum state vector

    Returns:
        dm: (array[complex]): array representing the density matrix
    """
    return np.outer(state,  np.conj(state))

state_labels = [label_0, label_1]
#state_labels = np.loadtxt('./tetra_states.txt', dtype=np.complex_)

dm_labels = [density_matrix(state_labels[i]) for i in range(2)]
len(dm_labels)

dm_labels

n_qubits = 2  # number of class
dev_fc = qml.device("default.qubit", wires=n_qubits)


@qml.qnode(dev_fc)
def q_fc(params, inputs):
    """A variational quantum circuit representing the DRC.

    Args:
        params (array[float]): array of parameters
        inputs = [x, y]
        x (array[float]): 1-d input vector
        y (array[float]): single output state density matrix

    Returns:
        float: fidelity between output state and input
    """
    
    # layer iteration
    for l in range(len(params[0])):
        # qubit iteration
        for q in range(n_qubits):
            # gate iteration
            for g in range(int(len(inputs)/3)):
                qml.Rot(*(params[0][l][3*g:3*(g+1)] * inputs[3*g:3*(g+1)] + params[1][l][3*g:3*(g+1)]), wires=q)
    
    return [qml.expval(qml.Hermitian(dm_labels[i], wires=[i])) for i in range(n_qubits)]


X_train[0].shape

a = np.random.uniform(size=(2, 1, 6))

q_fc(a, X_train[0])

tetra_class = np.loadtxt('./tetra_class_label.txt')
binary_class = np.array([[1, 0], [0, 1]])

class_labels = binary_class

n_class = 2

temp = np.zeros((len(Y_train), n_class))
for i in range(len(Y_train)):
    temp[i, :] = class_labels[Y_train[i]]
Y_train = temp


temp = np.zeros((len(Y_test), n_class))
for i in range(len(Y_test)):
    temp[i, :] = class_labels[Y_test[i]]
Y_test = temp

Y_train.shape, Y_test.shape

from keras import backend as K


# Alpha Custom Layer
class class_weights(tf.keras.layers.Layer):
    def __init__(self):
        super(class_weights, self).__init__()
        w_init = tf.random_normal_initializer()
        self.w = tf.Variable(
            initial_value=w_init(shape=(1, 2), dtype="float32"),
            trainable=True,
        )

    def call(self, inputs):
        return (inputs * self.w)

n_component = 6

X = tf.keras.Input(shape=(n_component,), name='Input_Layer')


# Quantum FC Layer, trainable params = 18*L*n_class + 2, output size = 2
num_fc_layer = 3
q_fc_layer_0 = qml.qnn.KerasLayer(q_fc, {"params": (2, num_fc_layer, n_component)}, output_dim=2)(X)

# Alpha Layer
alpha_layer_0 = class_weights()(q_fc_layer_0)

model = tf.keras.Model(inputs=X, outputs=alpha_layer_0)

model(X_train[0:32])

opt = tf.keras.optimizers.Adam(learning_rate=0.1)
model.compile(opt, loss='mse', metrics=["accuracy"])

filepath = "./Model_2/2_PCA_QFC_saved-model-{epoch:02d}.hdf5"
checkpoint = tf.keras.callbacks.ModelCheckpoint(filepath, monitor='val_acc', verbose=1,
                                                save_weights_only=True, save_best_only=False, mode='auto')

H = model.fit(X_train, Y_train, epochs=10, batch_size=32, initial_epoch=0,
              validation_data=(X_test, Y_test), verbose=1,
              callbacks=[checkpoint])

# another run to get 10 different set of weights
# 3 layer, 6 components
model.weights

# another run to get 10 different set of weights
# 3 layer, 6 components
H.history

# 3 layer, 3 components
H.history

# 2 layer, 3 components
H.history

# 1 layer, 3 components
H.history

# 3 layer, 6 components
H.history

# 2 layer, 6 components
H.history

# 1 layer, 6 components
H.history

model.summary()

# 3 layer, 6 components
model.weights

# random weights
model.weights

QFC_weights = np.array([[[-0.23825717, -0.6927138 ,  1.9275627 , -0.6338903 ,
          -1.5635191 , -0.54419893],
         [-1.7943736 ,  0.6555091 ,  0.7756721 , -0.52433693,
           0.33494785,  0.5896128 ],
         [-1.93925   ,  0.26347128, -1.0761861 ,  0.1370001 ,
           0.7956914 ,  0.43674225]],
 
        [[ 0.04991536,  0.4438185 ,  0.28958392,  0.17212974,
          -0.5826888 ,  0.1689779 ],
         [-0.05101966, -0.45951375, -0.7024131 , -0.51385176,
          -1.2483467 , -0.46147949],
         [-0.37164477, -0.980824  , -0.43812418, -0.7645581 ,
          -0.8557718 , -0.01652514]]])

alpha_weights = np.array([[1.0457449, 1.0058646]])


random_weights = np.array([[[-0.01116788, -0.24572456, -0.34748596,  0.12041545,
          -0.26595542, -0.18240207],
         [ 0.3004377 ,  0.5260234 ,  0.35336393,  0.19536716,
           0.01010054, -0.40111852],
         [ 0.5502764 , -0.18851322, -0.25431797, -0.24531147,
           0.07113147, -0.37417513]],
 
        [[-0.44977266,  0.34296083,  0.51045287, -0.2749984 ,
           0.41424918, -0.4837346 ],
         [ 0.10718679,  0.507282  ,  0.07516974, -0.4324274 ,
           0.36671913,  0.3184935 ],
         [ 0.49623525,  0.49480498,  0.4315616 ,  0.46681094,
          -0.12224993, -0.14642656]]])

QFC_weights.shape, alpha_weights.shape, random_weights.shape

# sanity check

pred = []
random_pred = []
for i in range(len(X_test)):
    pred += [np.argmax(q_fc(QFC_weights, X_test[i]).numpy())]
    random_pred += [np.argmax(q_fc(random_weights, X_test[i]).numpy())]

n_qubits = 1
dev_state = qml.device("default.qubit", wires=n_qubits)


@qml.qnode(dev_state)
def q_fc_state(params, inputs):
    """A variational quantum circuit representing the DRC.

    Args:
        params (array[float]): array of parameters
        inputs = [x, y]
        x (array[float]): 1-d input vector
        y (array[float]): single output state density matrix

    Returns:
        float: fidelity between output state and input
    """
    
    # layer iteration
    for l in range(len(params[0])):
        # qubit iteration
        for q in range(n_qubits):
            # gate iteration
            for g in range(int(len(inputs)/3)):
                qml.Rot(*(params[0][l][3*g:3*(g+1)] * inputs[3*g:3*(g+1)] + params[1][l][3*g:3*(g+1)]), wires=q)
    
    return qml.expval(qml.Hermitian(dm_labels[0], wires=[0]))


q_fc_state(QFC_weights, X_test[0])

dev_state._state

train_state = np.zeros((len(X_train), 2), dtype=np.complex_)
test_state = np.zeros((len(X_test), 2), dtype=np.complex_)

for i in range(len(train_state)):
    
    q_fc_state(QFC_weights, X_train[i])
    temp = np.flip(dev_state._state)
    train_state[i, :] = temp

for i in range(len(test_state)):
    
    q_fc_state(QFC_weights, X_test[i])
    temp = np.flip(dev_state._state)
    test_state[i, :] = temp

np.conj(train_state[150]) @ dm_labels[0] @ train_state[150]

train_state[0] @ np.conj(train_state[0])

np.conj(test_state[0]) @ dm_labels[0] @ test_state[0]

np.savetxt('./2_PCA_QFC-State_Train.txt', train_state)
np.savetxt('./2_PCA_QFC-State_Test.txt', test_state)

train_state = np.zeros((len(X_train), 2), dtype=np.complex_)
test_state = np.zeros((len(X_test), 2), dtype=np.complex_)

for i in range(len(train_state)):
    
    q_fc_state(random_weights, X_train[i])
    temp = np.flip(dev_state._state)
    train_state[i, :] = temp

for i in range(len(test_state)):
    
    q_fc_state(random_weights, X_test[i])
    temp = np.flip(dev_state._state)
    test_state[i, :] = temp

np.conj(train_state[0]) @ dm_labels[0] @ train_state[0]

train_state[0] @ np.conj(train_state[0])

np.savetxt('./2_PCA_QFC-RandomState_Train.txt', train_state)
np.savetxt('./2_PCA_QFC-RandomState_Test.txt', test_state)

random_weights = np.array(model.get_weights()[0], dtype=float)

QFC_weights_list = [random_weights]

for i in range(10):
    if i == 9:
        model.load_weights('./Model_2/2_PCA_QFC_saved-model-' + str(i+1) + '.hdf5')
    else:
        model.load_weights('./Model_2/2_PCA_QFC_saved-model-0' + str(i+1) + '.hdf5')
    
    QFC_weights_list += [np.array(model.get_weights()[0], dtype=float)]
    
len(QFC_weights_list)

n_qubits = 1
dev_state = qml.device("default.qubit", wires=n_qubits)


@qml.qnode(dev_state)
def q_fc_state(params, inputs):
    """A variational quantum circuit representing the DRC.

    Args:
        params (array[float]): array of parameters
        inputs = [x, y]
        x (array[float]): 1-d input vector
        y (array[float]): single output state density matrix

    Returns:
        float: fidelity between output state and input
    """
    
    # layer iteration
    for l in range(len(params[0])):
        # qubit iteration
        for q in range(n_qubits):
            # gate iteration
            for g in range(int(len(inputs)/3)):
                qml.Rot(*(params[0][l][3*g:3*(g+1)] * inputs[3*g:3*(g+1)] + params[1][l][3*g:3*(g+1)]), wires=q)
    
    return qml.expval(qml.Hermitian(dm_labels[0], wires=[0]))


q_fc_state(QFC_weights_list[0], X_test[0])


for k in range(10+1):

    train_state = np.zeros((len(X_train), 2), dtype=np.complex_)
    test_state = np.zeros((len(X_test), 2), dtype=np.complex_)

    for i in range(len(train_state)):

        q_fc_state(QFC_weights_list[k], X_train[i])
        temp = np.flip(dev_state._state)
        train_state[i, :] = temp

    for i in range(len(test_state)):

        q_fc_state(QFC_weights_list[k], X_test[i])
        temp = np.flip(dev_state._state)
        test_state[i, :] = temp
        
    print(train_state.shape, test_state.shape)  
    np.savetxt('./Model_2/2_PCA_QFC-State_Train_Epoch=' + str(k) + '.txt', train_state)
    np.savetxt('./Model_2/2_PCA_QFC-State_Test_Epoch=' + str(k) + '.txt', test_state)

