# Ignore warning message
import warnings
warnings.filterwarnings("ignore")

# Import python libraries
import math, json
import numpy as np
from numpy import pi
from numpy.linalg import eigvals
from numpy.fft import fft, rfft
from scipy.interpolate import CubicSpline
from scipy.optimize import curve_fit, minimize
import networkx as nx
import matplotlib.pyplot as plt
from matplotlib import colors

# Import Qiskit libraries
import qiskit
import qiskit_aer.noise as noise
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, IBMQ, Aer, transpile, execute
from qiskit_aer import AerSimulator
from qiskit_aer.noise import NoiseModel, depolarizing_error
from qiskit.result import marginal_counts

# Load IBMQ account
provider = IBMQ.load_account()

# Problem: Max-cut problem

import networkx as nx
from scipy.optimize import minimize

def maxcut_obj(x, G):
    
    obj = 0
    for i, j in G.edges():
        if x[i] != x[j]:
            obj -= 1
            
    return obj

def compute_expectation(counts, G):
    
    avg = 0
    sum_count = 0
    for bitstring, count in counts.items():
        
        obj = maxcut_obj(bitstring[::-1], G)
        avg += obj * count
        sum_count += count
        
    return avg/sum_count

def create_qaoa_circ(G, theta):
    
    nqubits = len(G.nodes())
    p = len(theta)//2  # number of alternating unitaries
    qc = QuantumCircuit(nqubits)
    
    beta = theta[:p]
    gamma = theta[p:]
    
    # initial_state
    for i in range(0, nqubits):
        qc.h(i)
    
    for irep in range(0, p):
        
        # problem unitary
        for pair in list(G.edges()):
            qc.rzz(2 * gamma[irep], pair[0], pair[1])

        # mixer unitary
        for i in range(0, nqubits):
            qc.rx(2 * beta[irep], i)
            
    qc.measure_all()
        
    return qc

def get_expectation(G, shots=512):
    
    backend = Aer.get_backend('qasm_simulator')
    backend.shots = shots
    
    def execute_circ(theta):
        
        qc = create_qaoa_circ(G, theta)
        counts = backend.run(qc, seed_simulator=10, nshots=512).result().get_counts()
        
        return compute_expectation(counts, G)
    
    return execute_circ

G = nx.Graph()
G.add_nodes_from([0, 1, 2, 3])
G.add_edges_from([(0, 1), (1, 2), (2, 3), (3, 0)])

expectation = get_expectation(G)
result = minimize(expectation, [1.0, 1.0], method='Nelder-Mead', tol=1e-3)
result

# Problem: Rosenberg function

from numpy import atleast_1d, eye, argmin, zeros, shape, squeeze, asarray, sqrt, Inf, asfarray
import numpy as np
from scipy.optimize import minimize

def rosen(x):
    
    x = np.array(x)
    r = np.sum(100.0 * a * (x[1:] - x[:-1]**2.0)**2.0 + (1 - x[:-1])**2.0, axis=0)
    
    return r

def callback(x):
    
    global Nfeval
    print('     {0:4d}  {1:3.6f}  {2:3.6f}  {3:3.6f}  {4:3.6f}  {5:3.6f}'.format(Nfeval, x[0], x[1], x[2], x[3], rosen(x)))
    Nfeval += 1

print('{0:4s} {1:9s} {2:9s} {3:9s} {4:9s}  {5:9s}'.format('Iteration', ' X1', ' X2', ' X3', ' X4', 'f(X)'))

Nfeval = 1
a = 1
x0 = [1.3, 0.7, 0.9, 1.1]
result = minimize(rosen, x0, method='Nelder-Mead', tol=1e-3, callback=callback)
result

# Problem: Transverse-field Ising model

# Set up parameters
Nc = 5  # number of spins: N (= 2, 3, 4, 5)
g = 0.4  # dimensionless Ising coupling: J/h (= 0.2, 0.4, 0.6, 0.8)
eta = 0.3  # dimensionless broadening: \eta/h (= 0.02, 0.1, 0.2, 0.3)
dw = 0.25*eta  # dimensionless unit of discrete frequency: \delta\omega / h
L = 2*math.ceil(7.0/dw)  # count of Fourier components: L
Le = 100
dt = 2*pi/L/dw  # dimensionless unit of discrete time: h \delta t
M = 35  # Trotter depth cutoff: M
shots = 1024  # count of measurement shots

# Define the identity and Pauli matrices
s0 = np.matrix([[1,0],[0,1]])
s1 = np.matrix([[0,1],[1,0]])
s2 = np.matrix([[0,-1j],[1j,0]])
s3 = np.matrix([[1,0],[0,-1]])

# Define the model Hamitonian
Ham = np.kron(s1,s0) + np.kron(s0,s1) + g * np.kron(s3,s3)
if Nc >= 3:
    for j in range(3,Nc+1):  
        Ide = s0
        for k in range(j-3):
            Ide = np.kron(Ide,s0)
        Ham = np.kron(Ham,s0) + np.kron(np.kron(Ide,s0),s1) + g * np.kron(np.kron(Ide,s3),s3)

# Calculate the exact energy gap
evals = eigvals(Ham)
idx = evals.argsort()[::1]
evals = evals[idx]
ExactGap = evals[1] - evals[0]

# set the initial guess of energy gap
IniGap = 2*(1-(1-1/Nc)*g)

# Define the bit string (00...0) as a reference for thr initial state
s = []
for j in range (Nc):
    s.append('0')
str1 = ""
InQ = str1.join(s)

# Define discrete frequencies
frequency = [(j-L/2)*dw for j in range(L+1)]
with open('data_frequency', 'w') as f:
    f.write(json.dumps(frequency))
        
# Define the objective function for curve fitting of the Lorentzian function
def objective (x, a0, a1, a2):
    return a0 + a1 / ((x - a2)**2 + eta**2)

# Print preliminary data
print("N = %s, J/h = %.2f, M = %s" % (Nc,g,M))
print("Initial gap = %.6f" % IniGap)
print("Exact gap = %.6f" % ExactGap)

# Simulation type: ideal (noiseless)
def InvPeak_Ideal(theta):
    
    global ExactGap, IniGap
    
    # number of alternating unitaries
    P = len(theta)//2
    beta = theta[:P]
    gamma = theta[P:]
    
    # Loop for causal (i=0) and anti-causal (i=1) processes
    for i in range (2):

        # Initialize array
        propagator = []
    
        # Loop for Fourier components
        for j in range (L):

            # Initialize a set of quantum/classical registers
            qreg = QuantumRegister(Nc, 'q')
            creg = ClassicalRegister(Nc, 'c')
            qc = QuantumCircuit(qreg, creg)
            
            # Set up the unitary for trial state
            for p in range (P):
    
                for n in range (Nc-1):
                    qc.rzz(gamma[p]*pi, qreg[n], qreg[n+1])
    
                for n in range (Nc):
                    qc.ry(beta[p]*pi, qreg[n])
    
            # Time evolution: Iterate Trotter steps
            for m in range (M):

                # Apply a sequence of unitary gates for time evolution
                for n in range (Nc):
                    qc.rx(-2*j*(-1)**i*dt/M, qreg[n])
                
                for n in range (Nc-1):
                    qc.rzz(-2*g*j*(-1)**i*dt/M, qreg[n], qreg[n+1])
                    
            # Set up the inverse unitary for trial state
            for p in range (P):
        
                for n in range (Nc):
                    qc.ry(-beta[P-p-1]*pi, qreg[n])
        
                for n in range (Nc-1):
                    qc.rzz(-gamma[P-p-1]*pi, qreg[n], qreg[n+1])
        
            # Measure qubits
            qc.measure(qreg, creg)
        
            # Perform quantum simulation on the IBMQ Aer simulator
            simulator = Aer.get_backend('qasm_simulator')
            result = execute(qc, simulator, shots = shots).result()
            counts = result.get_counts()
    
            # Construct a time series of propagators with each component being the probability for (00...0)
            if InQ in counts:
                propagator.append(dt*np.exp(-eta*j*dt)*counts[InQ]/shots)
            else:
                propagator.append(0)
    
            # Print the end message for each step
            if (i==1) and (j>L-10):
                print(L-1-j,end="\r")
    
        # Flatten array
        propagator = np.ravel(propagator)

        # Find the many-body spectral function via discreate fast Fourier transform
        if i == 0:
            spectral_function_0r = np.real(np.fft.ifft(propagator))*L/2.0/pi
            spectral_function_1r = np.append(spectral_function_0r[int(L/2):L], spectral_function_0r[0:int(L/2)+1])
        else:
            spectral_function_0a = np.real(np.fft.fft(propagator))/2.0/pi
            spectral_function_1a = np.append(spectral_function_0a[int(L/2):L], spectral_function_0a[0:int(L/2)+1])
    
    # Find the many-body spectral function via discreate fast Fourier transform
    spectral_function_2 = [spectral_function_1r[j] + spectral_function_1a[j] for j in range(L+1)]

    # Truncate the lists in the range [-\eta,\eta]
    frequency_truncated = [frequency[j] for j in range(math.ceil(L/2-2*eta/dw),math.ceil(L/2+1+2*eta/dw))]
    spectral_function_truncated = [spectral_function_2[j] for j in range(math.ceil(L/2-2*eta/dw),math.ceil(L/2+1+2*eta/dw))]

    # Curve fitting for the Lorentzian function
    a2 = 0
    popt, _ = curve_fit(lambda x, a0, a1: objective(x, a0, a1, a2), frequency_truncated, spectral_function_truncated)
    a0, a1 = popt

    # Subtract the zeto peak from the manu-body spectral function
    spectral_function = [spectral_function_2[j] - a1 / (frequency[j]**2 + eta**2) for j in range(L+1)]
    
    # Truncate the lists in the range [IniGap-\eta,IniGap+\eta]
    frequency_truncated = [frequency[j] for j in range(math.ceil(L/2+(IniGap-2*eta)/dw),math.ceil(L/2+1+(IniGap+2*eta)/dw))]
    spectral_function_truncated = [spectral_function[j] for j in range(math.ceil(L/2+(IniGap-2*eta)/dw),math.ceil(L/2+1+(IniGap+2*eta)/dw))]
    
    # Curve fitting for the Lorentzian function
    popt, _ = curve_fit(objective, frequency_truncated, spectral_function_truncated)
    a0, a1, a2 = popt
    
    return -a1/eta**2, a2, abs(a2 - ExactGap)/ExactGap, spectral_function

# Define wrapper functions
def InvPeak_Ideal_1(theta):
    obj1, obj2, obj3, obj4 = InvPeak_Ideal(theta)
    return obj1

def InvPeak_Ideal_2(theta):
    obj1, obj2, obj3, obj4 = InvPeak_Ideal(theta)
    return obj2

def InvPeak_Ideal_3(theta):
    obj1, obj2, obj3, obj4 = InvPeak_Ideal(theta)
    return obj3

def InvPeak_Ideal_4(theta):
    obj1, obj2, obj3, obj4 = InvPeak_Ideal(theta)
    return obj4

# Define callback function
def callback (theta):
    
    global data, Nfeval, theta0
    Nfeval += 1
    P0 = len(theta0)//2
    beta0 = theta0[:P0]
    gamma0 = theta0[P0:]
    P = len(theta)//2
    beta = theta[:P]
    gamma = theta[P:]
    peak = -InvPeak_Ideal_1(theta)
    gap_est = InvPeak_Ideal_2(theta)
    gap_err = InvPeak_Ideal_3(theta)
    spectral_function = InvPeak_Ideal_4(theta)
    
    # Print/save data
    print('     {0:4d}  {1:3.6f}   {2:3.6f}   {3:3.6f}   {4:3.6f}   {5:3.6f}'\
          .format(Nfeval, beta[0], gamma[0], peak, gap_est, gap_err))
    data.append([Nfeval, beta[0], gamma[0], peak, gap_est, gap_err])
        
    with open('data_spectral_ideal_N%s_g%.2f_eta%.2f_M%s_beta%.2f_gamma%.2f_NM%s' % 
              (int(Nc),g,eta,int(M),beta0[0],gamma0[0],int(Nfeval)), 'w') as f:
        f.write(json.dumps(spectral_function))

# Print the start message
print("Simulation type: Ideal")
print("Optimization method: Nelder-Mead")

# Set parameters
data = []
Nfeval = 0
theta0 = [0.5, 0.1]
P0 = len(theta0)//2
beta0 = theta0[:P0]
gamma0 = theta0[P0:]
bounds = ((0,2),(0,2))

# Save data
print('{0:4s} {1:9s}   {2:9s}  {3:9s}  {4:9s}  {5:9s}'\
      .format('Iteration',' beta/pi','gamma/pi','height','gap_est','gap_error'))

# Estimate energy gap
result = minimize(InvPeak_Ideal_1, theta0, method='Nelder-Mead', tol=1e-2, bounds=bounds, callback=callback)

# Save data
with open('data_iteration_ideal_N%s_g%.2f_eta%.2f_M%s_beta%.2f_gamma%.2f_NM' % 
          (int(Nc),g,eta,int(M),beta0[0],gamma0[0]), 'w') as f:
    f.write(json.dumps(data))
            
# Print the end message
print(result)

# Build noise model from backend calibration data
backend = provider.get_backend("ibmq_manila")
noise_model = NoiseModel.from_backend(backend)
coupling_map = backend.configuration().coupling_map
basis_gates = noise_model.basis_gates

# Simulation type: Device noise model
def InvPeak_Noisy(theta):
    
    global ExactGap, IniGap
    
    # number of alternating unitaries
    P = len(theta)//2
    beta = theta[:P]
    gamma = theta[P:]
    
    # Loop for causal (i=0) and anti-causal (i=1) processes
    for i in range (2):

        # Print the message
        #print("...running")
    
        # Initialize array
        propagator = []
    
        # Loop for Fourier components
        for j in range (L):

            # Initialize a set of quantum/classical registers
            qreg = QuantumRegister(Nc, 'q')
            creg = ClassicalRegister(Nc, 'c')
            qc = QuantumCircuit(qreg, creg)
            
            # Set up input state in QAOA
            for p in range (P):
    
                for n in range (Nc-1):
                    qc.rzz(gamma[p]*pi, qreg[n], qreg[n+1])
    
                for n in range (Nc):
                    qc.ry(beta[p]*pi, qreg[n])
    
            # Time evolution: Iterate Trotter steps
            for m in range (M):

                # Apply a sequence of unitary gates for time evolution
                for n in range (Nc):
                    qc.rx(-2*j*(-1)**i*dt/M, qreg[n])
                
                for n in range (Nc-1):
                    qc.rzz(-2*g*j*(-1)**i*dt/M, qreg[n], qreg[n+1])
                    
            # Set up the state to compenstate input in QAOA
            for p in range (P):
        
                for n in range (Nc):
                    qc.ry(-beta[P-p-1]*pi, qreg[n])
        
                for n in range (Nc-1):
                    qc.rzz(-gamma[P-p-1]*pi, qreg[n], qreg[n+1])
        
            # Measure qubits
            qc.measure(qreg, creg)
        
            # Perform quantum simulation on the IBMQ Aer simulator
            simulator = Aer.get_backend('qasm_simulator')
            qc_trans = transpile(qc, backend, optimization_level = 1, initial_layout = [k for k in range(Nc)])
            result = execute(qc_trans, simulator, noise_model = noise_model, coupling_map = coupling_map,\
                         basis_gates = basis_gates, shots = shots).result()    
            counts = result.get_counts()
    
            # Construct a time series of propagators with each component being the probability for (00...0)
            if InQ in counts:
                propagator.append(dt*np.exp(-eta*j*dt)*counts[InQ]/shots)
            else:
                propagator.append(0)
    
            # Print the end message for each step
            print(j+1, end="\r")
    
        # Flatten array
        propagator = np.ravel(propagator)

        # Find the many-body spectral function via discreate fast Fourier transform
        if i == 0:
            spectral_function_0r = np.real(np.fft.ifft(propagator))*L/2.0/pi
            spectral_function_1r = np.append(spectral_function_0r[int(L/2):L], spectral_function_0r[0:int(L/2)+1])
        else:
            spectral_function_0a = np.real(np.fft.fft(propagator))/2.0/pi
            spectral_function_1a = np.append(spectral_function_0a[int(L/2):L], spectral_function_0a[0:int(L/2)+1])
    
    # Find the many-body spectral function via discreate fast Fourier transform
    spectral_function_2 = [spectral_function_1r[j] + spectral_function_1a[j] for j in range(L+1)]

    # Truncate the lists in the range [-\eta,\eta]
    frequency_truncated = [frequency[j] for j in range(math.ceil(L/2-2*eta/dw),math.ceil(L/2+1+2*eta/dw))]
    spectral_function_truncated = [spectral_function_2[j] for j in range(math.ceil(L/2-2*eta/dw),math.ceil(L/2+1+2*eta/dw))]

    # Curve fitting for the Lorentzian function
    a2 = 0
    popt, _ = curve_fit(lambda x, a0, a1: objective(x, a0, a1, a2), frequency_truncated, spectral_function_truncated)
    a0, a1 = popt
    
    # Subtract the zeto peak from the manu-body spectral function
    spectral_function = [spectral_function_2[j] - a1 / (frequency[j]**2 + eta**2) for j in range(L+1)]
    
    # Truncate the lists in the range [IniGap-\eta,IniGap+\eta]
    frequency_truncated = [frequency[j] for j in range(math.ceil(L/2+(IniGap-2*eta)/dw),math.ceil(L/2+1+(IniGap+2*eta)/dw))]
    spectral_function_truncated = [spectral_function[j] for j in range(math.ceil(L/2+(IniGap-2*eta)/dw),math.ceil(L/2+1+(IniGap+2*eta)/dw))]
    
    # Curve fitting for the Lorentzian function
    popt, _ = curve_fit(objective, frequency_truncated, spectral_function_truncated)
    a0, a1, a2 = popt
    
    return -a1/eta**2, a2, abs(a2 - ExactGap)/ExactGap, spectral_function

# Define wrapper functions
def InvPeak_Noisy_1(theta):
    obj1, obj2, obj3, obj4 = InvPeak_Noisy(theta)
    return obj1

def InvPeak_Noisy_2(theta):
    obj1, obj2, obj3, obj4 = InvPeak_Noisy(theta)
    return obj2

def InvPeak_Noisy_3(theta):
    obj1, obj2, obj3, obj4 = InvPeak_Noisy(theta)
    return obj3

def InvPeak_Noisy_4(theta):
    obj1, obj2, obj3, obj4 = InvPeak_Noisy(theta)
    return obj4

# Define callback function
def callback (theta):
    
    global data, Nfeval, theta0
    Nfeval += 1
    P0 = len(theta0)//2
    beta0 = theta0[:P0]
    gamma0 = theta0[P0:]
    P = len(theta)//2
    beta = theta[:P]
    gamma = theta[P:]
    peak = -InvPeak_Noisy_1(theta)
    gap_est = InvPeak_Noisy_2(theta)
    gap_err = InvPeak_Noisy_3(theta)
    spectral_function = InvPeak_Noisy_4(theta)
    
    # Print/save data
    print('     {0:4d}  {1:3.6f}   {2:3.6f}   {3:3.6f}   {4:3.6f}   {5:3.6f}'\
          .format(Nfeval, beta[0], gamma[0], peak, gap_est, gap_err))
    data.append([Nfeval, beta[0], gamma[0], peak, gap_est, gap_err])
        
    with open('data_spectral_noisy_N%s_g%.2f_eta%.2f_M%s_beta%.2f_gamma%.2f_NM%s' % 
              (int(Nc),g,eta,int(M),beta0[0],gamma0[0],int(Nfeval)), 'w') as f:
        f.write(json.dumps(spectral_function))

# Print the start message
print("Simulation type: Device noise model")
print("Optimization method: Nelder-Mead")

# Set parameters
data = []
Nfeval = 0
theta0 = [0.5, 0.1]
P0 = len(theta0)//2
beta0 = theta0[:P0]
gamma0 = theta0[P0:]
bounds = ((0,2),(0,2))

# Save data
print('{0:4s} {1:9s}   {2:9s}  {3:9s}  {4:9s}  {5:9s}'\
      .format('Iteration',' beta/pi','gamma/pi','height','gap_est','gap_error'))

# Estimate energy gap
result = minimize(InvPeak_Noisy_1, theta0, method='Nelder-Mead', tol=1e-2, bounds=bounds, callback=callback)

# Save data
with open('data_iteration_noisy_N%s_g%.2f_eta%.2f_M%s_beta%.2f_gamma%.2f_NM' % 
          (int(Nc),g,eta,int(M),beta0[0],gamma0[0]), 'w') as f:
    f.write(json.dumps(data))
            
# Print the end message
print(result)

