!pip install qiskit
!pip install pylatexenc--upgrade
!pip install retworkx
#!pip install numpy --upgrade
from IPython.display import clear_output
clear_output()

from qiskit import QuantumRegister, ClassicalRegister, QuantumCircuit, Aer, transpile
from qiskit.circuit.library import ZGate, XGate
from qiskit.extensions import UnitaryGate
from qiskit.aqua.algorithms import VQE, NumPyEigensolver
from qiskit.aqua.operators import WeightedPauliOperator, MatrixOperator, op_converter
from qiskit.aqua.components.optimizers import COBYLA, SPSA, SLSQP
#from qiskit.quantum_info.operators import Operator

import numpy as np
import scipy
import networkx as nx
from matplotlib import pyplot as plt
import copy

# Define the number of qubits

num_qubit = 4
qubits = range(num_qubit)

# Generating the graph

ising_graph = nx.star_graph(num_qubit-1)
print("Edges:", ising_graph.edges)
nx.draw(ising_graph)

# Set the random seed

rng = np.random.default_rng(seed=2021)

# Initialize the true weights and bias

true_weights = rng.random(size=len(ising_graph.edges)) - 0.5
true_bias = rng.random(size=num_qubit) - 0.5

print("True Weights:", true_weights)
print("True Bias", true_bias)

def create_hamiltonian_matrix(n_qubits, graph, weights, bias):

    full_matrix = np.zeros((2 ** n_qubits, 2 ** n_qubits))

    # Creates the interaction component of the Hamiltonian
    for i, edge in enumerate(graph.edges):
        interaction_term = 1
        for qubit in range(0, n_qubits):
            if qubit in edge:
                interaction_term = np.kron(interaction_term, np.array(ZGate().to_matrix(), dtype=float))
            else:
                interaction_term = np.kron(interaction_term, np.identity(2))
        full_matrix += weights[i] * interaction_term

    # Creates the bias components of the matrix
    for i in range(0, n_qubits):
        z_term = x_term = 1
        for j in range(0, n_qubits):
            if j == i:
                z_term = np.kron(z_term, np.array(ZGate().to_matrix(), dtype=float))
                x_term = np.kron(x_term, np.array(XGate().to_matrix(), dtype=float))
            else:
                z_term = np.kron(z_term, np.identity(2))
                x_term = np.kron(x_term, np.identity(2))
        full_matrix += bias[i] * z_term + x_term

    return full_matrix

# Prints a visual representation of the Hamiltonian matrix
H_matrix = create_hamiltonian_matrix(num_qubit, ising_graph, true_weights, true_bias)

print("Hamiltonian Matrix Visualization")
plt.matshow(H_matrix, cmap="cool")
plt.colorbar()
plt.show()

ground_state_energy = np.real_if_close(min(np.linalg.eig(H_matrix)[0]))
print("Ground State Energy:", ground_state_energy)

H_operator = op_converter.to_weighted_pauli_operator(MatrixOperator(H_matrix))

eig_vals = NumPyEigensolver(H_operator).run()
print("Minimum Eigenvalue:", np.real_if_close(min(eig_vals['eigenvalues'])))

opt = SLSQP(maxiter=5)

backend = Aer.get_backend("statevector_simulator")

vqe=VQE(operator=H_operator, optimizer=opt)
vqe_result=vqe.run(backend)

print("VQE Eigenvalue:", np.real_if_close(vqe_result['eigenvalue']))

#vqe.get_optimal_circuit().draw()
low_energy_state = vqe.get_optimal_vector()
print("Low Energy State:", low_energy_state)

# Sanity check, will be deleted in the final notebook

print(np.real_if_close(np.conj(vqe.get_optimal_vector())@H_matrix@vqe.get_optimal_vector()))
print(np.conj(vqe.get_optimal_vector())@vqe.get_optimal_vector())

import math
math.isclose(np.sum(np.absolute(np.kron(low_energy_state,low_energy_state)) ** 2), 1.0,
                                abs_tol=1e-11)

def state_evolve(circuit, hamiltonian, qubits, time):

    U = scipy.linalg.expm(-1j * hamiltonian * time)
    U_gate = UnitaryGate(U, label='evolved_unitary_'+str(time))
    circuit.unitary(U_gate, qubits)
    
    return circuit

def qgrnn_layer(circuit, weights, bias, qubits, graph, trotter_step):

    # Applies a layer of RZZ gates (based on a graph)
    for i, edge in enumerate(graph.edges):
        circuit.rzz(2 * weights[i] * trotter_step, edge[0], edge[1])

    # Applies a layer of RZ gates
    for i, qubit in enumerate(qubits):
        circuit.rz(2 * bias[i] * trotter_step, qubit)

    # Applies a layer of RX gates
    for qubit in qubits:
        circuit.rx(2 * trotter_step, qubit)
        
    return circuit

def swap_test(circuit, control, register1, register2):

    circuit.h(control)
    for reg1_qubit, reg2_qubit in zip(register1, register2):
        circuit.cswap(control, reg1_qubit, reg2_qubit)
    circuit.h(control)
    
    return circuit

# Defines some fixed values

reg_data = tuple(range(num_qubit))  # First qubit register for quantum data
reg_qgrnn = tuple(range(num_qubit, 2 * num_qubit))  # Second qubit register for QGRNN

control = 2 * num_qubit  # Index of control qubit
trotter_step = 0.01  # Trotter step size

# Defines the initial interaction graph

initial_ising_graph = nx.complete_graph(reg_qgrnn)

print("Edges:", initial_ising_graph.edges)
nx.draw(initial_ising_graph)

# Initialize random weights and bias

weights = rng.random(size=len(initial_ising_graph.edges)) - 0.5
bias = rng.random(size=num_qubit) - 0.5
trainable_params = np.concatenate((weights, bias))

initial_weights = copy.copy(weights)
initial_bias = copy.copy(bias)

print(trainable_params)

def qgrnn(param_weights, param_bias, time=None):

    circuit = QuantumCircuit(2*num_qubit+1, 1)
    
    # Prepares the low energy state in the two registers
    circuit.initialize(np.kron(low_energy_state, low_energy_state), reg_data + reg_qgrnn)

    # Evolves the first qubit register with the time-evolution circuit to
    # prepare a piece of quantum data
    circuit = state_evolve(circuit, H_matrix, reg_data, time)

    # Applies the QGRNN layers to the second qubit register
    depth = time / trotter_step  # P = t/Delta
    for _ in range(0, int(depth)):
        circuit = qgrnn_layer(circuit, param_weights, param_bias, reg_qgrnn, initial_ising_graph, trotter_step)

    # Applies the SWAP test between the registers
    circuit = swap_test(circuit, control, reg_data, reg_qgrnn)

    # Measure the circuit
    circuit.measure([control], [0])
    
    return circuit

qgrnn(initial_weights, initial_bias, time=0.02).draw()



def z_exp(circuit, shots=10000):
    backend = Aer.get_backend('aer_simulator_statevector')

    job = backend.run(circuit, shots=shots)
    result = job.result()
    counts = result.get_counts(0)
    
    try:
        exp_0 = counts['0']/shots
    except:
        exp_0 = 0
        
    try:
        exp_1 = counts['1']/shots
    except:
        exp_1 = 0
    
    exp = exp_0 - exp_1
    
    return exp

N = 15  # The number of pieces of quantum data that are used for each step
max_time = 0.1  # The maximum value of time that can be used for quantum data

def cost_function(params):

    # Randomly samples times at which the QGRNN runs
    times_sampled = rng.random(size=N) * max_time

    # Cycles through each of the sampled times and calculates the cost
    total_cost = 0
    for dt in times_sampled:
        result = z_exp(qgrnn(params[:len(initial_ising_graph.edges)], params[len(initial_ising_graph.edges):], time=dt))
        total_cost += -1 * result

    return total_cost / N

# Sanity check, will be deleted in the final notebook

backend = Aer.get_backend('aer_simulator_statevector')
job = backend.run(qgrnn(trainable_params[:len(initial_ising_graph.edges)], trainable_params[len(initial_ising_graph.edges):], time=0.09716818), shots=10000)
result = job.result()
counts = result.get_counts(0)
counts

# Sanity check
cost_function(trainable_params)

opt_params = scipy.optimize.minimize(fun=cost_function, x0=trainable_params, tol=1e-12,
                                     method='L-BFGS-B',
                                    #args=(),
                                    options={'maxiter':5000, 'ftol':1e-12})

print(opt_params)

trained_weights = opt_params['x'][:len(initial_ising_graph.edges)]
trained_bias = opt_params['x'][len(initial_ising_graph.edges):]

weights.shape, bias.shape

H_matrix_trained = create_hamiltonian_matrix(
    num_qubit, nx.complete_graph(num_qubit), trained_weights, trained_bias
)

H_matrix_initial = create_hamiltonian_matrix(
    num_qubit, nx.complete_graph(num_qubit), initial_weights, initial_bias
)

fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(6, 6))

axes[0].matshow(H_matrix, vmin=-7, vmax=7, cmap="cool")
axes[0].set_title("Target", y=1.13)

axes[1].matshow(H_matrix_initial, vmin=-7, vmax=7, cmap="cool")
axes[1].set_title("Initial", y=1.13)

axes[2].matshow(H_matrix_trained, vmin=-7, vmax=7, cmap="cool")
axes[2].set_title("Learned", y=1.13)

plt.subplots_adjust(wspace=0.3, hspace=0.3)
plt.show()

