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


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


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])

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(*x_sample[3*j:3*(j+1)], 0)
            circuit.u(*parameters[3*j:3*(j+1)], 0)
    
    circuit.measure(0,0)
    
    return circuit

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

drc_circuit(X_train[0], params, num_features).draw('mpl')

shots=1000
backend = Aer.get_backend('qasm_simulator')
L = 1
num_features = 6

def weighted_fidelity_cost_function(params):
    loss = 0
    
    # Training sample iteration
    for k in range(len(X_train)):
        q_reg = QuantumRegister(1)
        c_reg = ClassicalRegister(1)
        circuit = QuantumCircuit(q_reg, c_reg)
    
        # Layer iteration
        for i in range(L):
            # Gate iteration
            for j in range(int(num_features/3)):
                circuit.u(*X_train[k][3*j:3*(j+1)], 0)
                circuit.u(*params[3*j:3*(j+1)], 0)
    
        circuit.measure(0,0)
        
        t_circ = transpile(circ, backend)
        qobj = assemble(t_circ, shots=shots)
        job = backend.run(qobj)
        result = job.result().get_counts(circ)

        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[0,0])**2 + (fidel_1-Y_train[0,1])**2)/2

    loss = loss/len(X_train)
    
    return loss

weighted_fidelity_cost_function(params)

from scipy.optimize import minimize

out = minimize(weighted_fidelity_cost_function, x0=params, method='L-BFGS-B', options={'maxiter':10000})
print(out)

params

from qiskit.aqua.components.optimizers import ADAM
from qiskit.aqua.components.optimizers import L_BFGS_B

opt_adam = ADAM(maxiter=100)
opt_lbfgsb = L_BFGS_B(maxiter=100, iprint=10)

out = opt_adam.optimize(num_vars=8, objective_function=weighted_fidelity_cost_function, initial_point=params)

out

out

out



params

weights = np.random.rand(num_features + 2)
weights = tf.Variable(weights, dtype=tf.float64, trainable=True)

weights.numpy()

shots=2000
backend = Aer.get_backend('qasm_simulator')
L = 1

def drc_pred(params):
    params = params.numpy()
    pred = np.zeros((len(X_train),2))
    
    # Training sample iteration
    for i in range(len(X_train)):
        circ = drc_circuit(X_train[i], params, num_features, layer=L)
        t_circ = transpile(circ, backend)
        qobj = assemble(t_circ)
        job = backend.run(qobj, shots=shots)
        result = job.result().get_counts(circ)

        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[i] = np.array([fidel_0, fidel_1])
        
    
    return pred

optimizer = tf.keras.optimizers.Adam(lr=0.01)

def loss(params):
    pred = drc_pred(params)
    return tf.keras.losses.MeanSquaredError()(Y_train, pred)

loss(weights)

optimizer.minimize(lambda: loss(weights), [weights])

