import qiskit
qiskit.__qiskit_version__

# import warnings; warnings.filterwarnings("ignore")
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
from collections import defaultdict
from operator import itemgetter
from scipy.optimize import minimize
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, execute, Aer
from qiskit.visualization import plot_histogram

G = nx.Graph()
Lx=6
G.add_nodes_from([i for i in range(Lx)])
G.add_edges_from([[i,i+1] for i in range(Lx-1)])
G.add_edges_from([[Lx-1,0]])
options = {
    'node_color': 'y',
    'node_size': 1000,
    'width': 3,
}
nx.draw(G, with_labels=True,pos=nx.circular_layout(G),**options)

print(list(G.edges()))
print(G.number_of_nodes())

def ham_rot(G,gamma):
    L=G.number_of_nodes()
    qc=QuantumCircuit(L,L)
    for i,j in list(G.edges()):
        qc.cx(i,j)
        qc.rz(2*gamma,j)
        qc.cx(i,j)
    return qc

qc = ham_rot(G, np.pi/3)
print(G.edges())
qc.draw()

def x_field_rot(G, beta):
    N = G.number_of_nodes()
    qc = QuantumCircuit(N,N)
    for n in G.nodes():
        qc.rx(2*beta, n)
    return qc

qc = x_field_rot(G, np.pi/3)
qc.draw()

def qaoa_circuit(G,beta,gamma):
    assert(len(beta)==len(gamma))
    layer=len(beta)
    L=G.number_of_nodes()
    qc=QuantumCircuit(L,L)
    qc.h(range(L))
    for k in range(layer):
        qc=qc+ham_rot(G,gamma[k])
        qc=qc+x_field_rot(G,beta[k])
    qc.measure(range(L), range(L))
    return qc

qc = qaoa_circuit(G,[np.pi/3], [np.pi/2])
qc.draw()

def invert_counts(counts):
    return {k[::-1]:v for k, v in counts.items()}

backend = Aer.get_backend('qasm_simulator')
job = execute(qc, backend,shots=2048)          # this is just a test run for the above circuit
result = job.result()
print(result.get_counts())

def ising_energy(x,G):
    ener = 0
    for i, j in G.edges():
        if x[i] != x[j]:
            # the edge is cut
            ener += 1
        elif x[i] == x[j]:
            # the edge is cut
            ener -= 1
    return ener

ising_energy('010101', G),ising_energy('111110', G)

def compute_ising_energy(counts, G):
    energy = 0
    total_counts = 0
    for meas, meas_count in counts.items():             # meas and meas_count are the state and freq
#         print(type(meas),meas_count)
        obj_for_meas = ising_energy(meas, G)              # calling the state to compute the energy
        energy += obj_for_meas * meas_count
        total_counts += meas_count                      # energy is total energy of all the states
    return energy / total_counts                        # total_counts is total number of sampling

compute_ising_energy(invert_counts(result.get_counts()), G)

invert_counts(result.get_counts());

# here p is the number of layer used for the optimize

def get_black_box_objective(G,p,list1):
    backend = Aer.get_backend('qasm_simulator')
    def f(theta):
        # let's assume first half is betas, second half is gammas
        # theta is a 2*p length array, first 'p' entries are gamma and rest are beta
        beta = theta[:p]
        gamma = theta[p:]
        qc = qaoa_circuit(G,beta, gamma)
        counts = execute(qc, backend,shots=1024, seed_simulator=100).result().get_counts()
        # return the energy
        ener=compute_ising_energy(invert_counts(counts), G)
        list1.append(ener)
        return ener, list1
    return f

# p is the number of QAOA alternating operators
p = 5
energy_qasm=[]
obj = get_black_box_objective(G, p, energy_qasm)

init_point = np.array([0.81069872, 2.2067517 , 1.83830696, 2.15579759, 0.37060699,
       2.42068091, 6.1575306 , 2.2453419 , 0.85060091, 3.137845  ])

# We are going to limit the number of iterations to 2500
res_sample = minimize(obj, init_point, method='COBYLA', options={'maxiter':2500, 'disp': True})
res_sample

# Once we have gotten the optimum beta and gamma we run the circuit again with the 

optimal_theta = res_sample['x']
qc = qaoa_circuit(G, optimal_theta[:p], optimal_theta[p:])
counts = invert_counts(execute(qc, backend,shots=2048).result().get_counts())
counts

plot_histogram(counts,figsize=(12,4))

# get the best solution:
best_cut, best_solution = min([(ising_energy(x,G),x) for x in counts.keys()], key=itemgetter(0))
print(f"Best string: {best_solution} with cut: {-best_cut}")

# color the graph nodes by part
colors = ['c' if best_solution[node] == '0' else 'y' for node in G]
options = {
    'node_size': 1000,
    'width': 3,
}
nx.draw(G, with_labels=True,node_color = colors,pos=nx.circular_layout(G),**options)

len(energy_qasm);

def state_to_ampl_counts(vec, eps=1e-15):
    """Converts a statevector to a dictionary
    of bitstrings and corresponding amplitudes
    """
    qubit_dims = np.log2(vec.shape[0])
    if qubit_dims % 1:
        raise ValueError("Input vector is not a valid statevector for qubits.")
    qubit_dims = int(qubit_dims)
    counts = {}
    str_format = '0{}b'.format(qubit_dims)
    for kk in range(vec.shape[0]):
        val = vec[kk]
        if val.real**2+val.imag**2 > eps:
            counts[format(kk, str_format)] = val
    return counts

def compute_ising_energy_sv(sv, G):
    """Compute objective from statevector
    For large number of qubits, this is slow. 
    """
    counts = state_to_ampl_counts(sv)
    return sum(ising_energy(np.array([int(x) for x in k]), G) * (np.abs(v)**2) for k, v in counts.items())

def get_qaoa_circuit_sv(G, beta, gamma):
    assert(len(beta) == len(gamma))
    p = len(beta) # infering number of QAOA steps from the parameters passed
    N = G.number_of_nodes()
    qc = QuantumCircuit(N,N)
    # first, apply a layer of Hadamards
    qc.h(range(N))
    # second, apply p alternating operators
    for k in range(p):
        qc=qc+ham_rot(G,gamma[k])
        qc=qc+x_field_rot(G,beta[k])
    # no measurement in the end!
    return qc

def get_black_box_objective_sv(G,p,list1):
    backend = Aer.get_backend('statevector_simulator')
    def f(theta):
        # let's assume first half is betas, second half is gammas
        beta = theta[:p]
        gamma = theta[p:]
        qc = get_qaoa_circuit_sv(G,beta, gamma)
        sv = execute(qc, backend).result().get_statevector()  # the sv is np.array with complex entries
        # return the energy
        ener=compute_ising_energy_sv(sv, G)
        list1.append(ener)
        return ener, list1
    return f

# p is the number of QAOA alternating operators
energy_sv=[]
obj = get_black_box_objective_sv(G, p, energy_sv)

init_point = np.array([0.81069872, 2.2067517 , 0.83830696, 2.15579759, 0.37060699,
       2.42068091, 6.1575306 , 2.2453419 , 3.85060091, 6.137845  ])

p = len(init_point)//2

# We are going to limit the number of iterations to 2500
res_sv = minimize(obj, init_point, method='COBYLA', options={'maxiter':2500, 'disp': True})
res_sv

# Once we have gotten the optimum beta and gamma we run the circuit again with the 

backend = Aer.get_backend('qasm_simulator')

optimal_theta = res_sv['x']
qc = qaoa_circuit(G, optimal_theta[:p], optimal_theta[p:])
counts = invert_counts(execute(qc, backend,shots=2048).result().get_counts())
counts

plot_histogram(counts,figsize=(8,4))

len(energy_sv);

ms=6
x_sv=np.arange(1,len(energy_sv)+1)
x_qasm=np.arange(1,len(energy_qasm)+1)

plt.figure(figsize=(7,6))
plt.plot(x_sv,energy_sv,'b-',markersize=ms,label='Statevector')
plt.plot(x_qasm,energy_qasm,'r--',markersize=ms,label='QASM')
plt.xlabel('Iterations',fontsize=16)
plt.ylabel('Energy',fontsize=16)
plt.xticks(fontsize=16)
plt.yticks(fontsize=16)
plt.legend(fontsize=16)
# plt.title('Plot of gap for different $t_2$, at $t_1=0$ ')
plt.grid()
plt.tight_layout()
# plt.yscale('log')
# plt.legend(loc='upper right',fontsize=16)
# plt.savefig('%s/gap_t2.pdf'%(loc), bbox_inches='tight')
plt.show()

print(f"Angles found with samples: beta={res_sample['x'][:p]}, gamma={res_sample['x'][p:]}")
print(f"Angles found with full state: beta={res_sv['x'][:p]}, gamma={res_sv['x'][p:]}")