import numpy as np
import matplotlib.pyplot as plt
import time, copy, warnings
import networkx as nx


n = 17                                # = number of qubits
dim = 2**n                             # = dimension of the n-qubit Hilbert space

Q = np.random.randint(-10, 10, (n, n))  
Q = np.triu(Q)

print(f'{Q} = Q matrix')

def plot_Q_matrix():
    plt.matshow(Q, cmap='bwr')
    plt.xticks(np.arange(n))
    plt.yticks(np.arange(n))
    cb = plt.colorbar(shrink=0.7)
    plt.title(f'Q matrix') 
    plt.show()

plot_Q_matrix()

#---------------------------------------------------------------------------- 
'In the graph, a node represents a qubit (binary variables), width of the edge (i,j) is proportional to Q_ij'

plt.subplots(figsize=(10,8))  
g = nx.Graph()

edges = []
for i in range(n):
    for j in range(i, n):
        if Q[i,j]!=0:
            edges.append((i, j, Q[i,j]))
            
g.add_weighted_edges_from(edges)

pos = nx.spring_layout(g, seed=0)
nx.draw_networkx(g, pos, node_size= 700, node_color='magenta', 
                 width=list(nx.get_edge_attributes(g, "weight").values()))


%%time

def energy(x, Q):                                              # x in an integer, one-to-one correspondence: x <--> x_bin 
    x_bin = np.array([int(i) for i in np.binary_repr(x, n)])   # x_bin has n components, each one is either 0 (spin-up) or 1(spin-down)                                      
    E = np.linalg.multi_dot([x_bin, Q, x_bin])                 # E_x = energy for a given x
    return E
    

def energies(Q):
    n = Q.shape[0]                  # number of qubits
    dim = 2**n                      # dimension of the n-qubit Hilbert space
    xE = []
    for x in range(dim):                                                                                
        E = energy(x, Q)              
        xE.append([x, E])
    return np.array(xE)             # collection of all possible 2**n energies for different x

#----------------------------------------------------------------------------  
        
xE = energies(Q=Q)
engs = xE[:,1] 
print(f'energies for different x = [...E_x...] = {engs}')

#==============================================================================

min_eng = np.min(engs)
max_eng = np.max(engs)

plt.rcParams["figure.figsize"] = (12,10)
plt.scatter(np.arange(dim), engs, s=1, c='blue')
plt.axhline(y = 0, color = 'black', linestyle = '-')
plt.axhline(y = min_eng, color = 'red', linestyle = '-')
plt.ylabel(r'energy = $E_x = \langle x|H|x\rangle$', fontsize=20)
plt.xlabel('x', fontsize=20)
plt.show()

xmins = np.where(engs==min_eng)[0]
xmins_bin = [[int(i) for i in np.binary_repr(xmin, n)] for xmin in xmins]

print(f'global min_energy = {min_eng}')
print(f'global min_energy eigenstates = {xmins} = {xmins_bin}\n')
print(f'global max_energy = {max_eng}')
    


from qiskit_optimization import QuadraticProgram                         #!pip install qiskit[optimization]

from qiskit.algorithms.minimum_eigensolvers import NumPyMinimumEigensolver, SamplingVQE, QAOA
from qiskit_optimization.algorithms import MinimumEigenOptimizer


from qiskit.primitives import Sampler
from qiskit.circuit.library import TwoLocal
from qiskit.algorithms.optimizers import SPSA, COBYLA

qp = QuadraticProgram()

for i in range(n):
    qp.binary_var(f'x{i}')

    
linear = [Q[i,i] for i in range(n)]

quadratic = {}
for i in range(n):
    for j in range(i+1, n):
        quadratic[(f'x{i}', f'x{j}')] = Q[i,j]
        

qp.minimize(linear=linear, quadratic=quadratic)
qp


%%time

exact = MinimumEigenOptimizer(min_eigen_solver=NumPyMinimumEigensolver())
result = exact.solve(qp)
print(result.prettyprint())
print()



%%time
ansatz = TwoLocal(n, "ry", "cx", reps=3, entanglement="linear")           # construct ansatz for VQE
vqe = SamplingVQE(sampler=Sampler(), ansatz=ansatz, optimizer=COBYLA(maxiter=10))    # construct VQE


meo_vqe = MinimumEigenOptimizer(min_eigen_solver=vqe)                                # run VQE
result_vqe = meo_vqe.solve(qp)


print(result_vqe.prettyprint())
print()



%%time
qaoa = QAOA(sampler=Sampler(), optimizer=COBYLA(maxiter=10), reps=3)      # construct QAOA


meo_qaoa = MinimumEigenOptimizer(min_eigen_solver=qaoa)                   # run QAOA
result_qaoa = meo_qaoa.solve(qp)


print(result_qaoa.prettyprint())
print()
