# general imports
import math
import time
import pickle
import random
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import minimize
from datetime import datetime

# Noisy optimization package — you could also use scipy's optimization functions, 
# but this is a little better suited to the noisy output of NISQ devices.
%pip install noisyopt
import noisyopt

# magic invocation for producing visualizations in notebook
%matplotlib inline

# Qiskit imports
from qiskit import Aer

from qiskit_ionq import IonQProvider 

#Call provider and set token value
provider = IonQProvider(token='My token')

provider.backends()

# fix random seed for reproducibility — this allows us to re-run the process and get the same results
seed = 42
np.random.seed(seed)
random.seed(a=seed)

#Generate BAS, only one bar or stripe allowed
def generate_target_distribution(rows, columns):        
    #Stripes
    states=[]
    for i in range(rows):
        s=['0']*rows*columns
        for j in range(columns):
            s[j+columns*i]='1'
        states.append(''.join(s))
                
        
    #Bars
    for j in range(columns):
        s=['0']*rows*columns
        for i in range(rows):
            s[j+columns*i]='1'
        states.append(''.join(s))
        
    return states

#Parameters for driver
def generate_beta(ansatz_type,random_init):
    if ansatz_type[0]==0: #No driver
        beta=[]
    elif ansatz_type[0]==1: #Rz(t1)Rx(t2)Rz(t3), angles different for each qubit
        beta=(
            [random.uniform(0.0,2.*math.pi) for i in range(3*n*(layers-1))] 
            if random_init 
            else [0]*3*n*(layers-1)
        )
    elif ansatz_type[0]==2: #Rz(t1)Rx(t2)Rz(t3), angles same for all qubits
        beta=(
            [random.uniform(0.0,2.*math.pi) for i in range(3*(layers-1))]
            if random_init
            else [0]*3*(layers-1)
        )
    elif ansatz_type[0]==3: #Rz(t1), angles different for each qubit
        beta=(
            [random.uniform(0.0,2.*math.pi) for i in range(n*(layers-1))]
            if random_init
            else [0]*n*(layers-1)
        )
    else:
        raise Exception("Undefined driver type")
            
    return beta

#Parameters for entangler
def generate_gamma(ansatz_type,random_init, n, conn):
    length_gamma=int(n*conn-conn*(conn+1)/2.)
    if ansatz_type[1]==0: #No entangler
        gamma=[]
    elif ansatz_type[1]==1: #XX(t1), angles different for each qubit
        gamma=(
            [random.uniform(0.0,2.*math.pi) for i in range(length_gamma*layers)]
            if random_init
            else [0]*length_gamma*layers
        )
    elif ansatz_type[1]==2: #XX(t1), angles same for all qubits
        gamma=(
            [random.uniform(0.0,2.*math.pi) for i in range(layers)]
            if random_init
            else gamma[0]*layers
        )
    else:
        raise Exception("Undefined entangler type")
            
    return gamma


from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister
from qiskit.circuit import Gate

def driver(circ,qr,beta,n,ansatz_type):
    #qr=QuantumRegister(n)
    #circ = QuantumCircuit(qr)
    
    if ansatz_type==0:
        pass
    elif ansatz_type==1:
        for i_q in range(n):
            circ.rz(beta[3*i_q], qr[i_q])
            circ.rx(beta[3*i_q+1], qr[i_q])
            circ.rz(beta[3*i_q+2], qr[i_q])
    elif ansatz_type==2:
        for i_q in range(n):
            circ.rz(beta[0], qr[i_q])
            circ.rx(beta[1], qr[i_q])
            circ.rz(beta[2], qr[i_q])
    elif ansatz_type==3:
        for i_q in range(n):
            circ.rz(beta[i_q], qr[i_q])        
    
    return

def entangler(circ,qr,gamma,n,conn,ansatz_type):
    #qr=QuantumRegister(n)
    #cr=ClassicalRegister(n)
    #circ = QuantumCircuit(qr)

    if ansatz_type==0:
        pass
    elif ansatz_type==1:
        i_gamma=0
        for i_conn in range(1,conn+1):
            for i_q in range(0,n-i_conn):  
                circ.cx(qr[i_q],qr[i_q+i_conn])
                circ.rx(gamma[i_gamma],qr[i_q])
                circ.cx(qr[i_q],qr[i_q+i_conn])
                #circ.rxx(gamma[i_gamma], qr[i_q], qr[i_q+i_conn])
                i_gamma+=1            
                
    elif ansatz_type==2:
        for i_conn in range(1,conn+1):
            for i_q in range(0,n-i_conn): 
                circ.cx(qr[i_q],qr[i_q+i_conn])
                circ.rx(gamma[0],qr[i_q])
                circ.cx(qr[i_q],qr[i_q+i_conn])                                                       #circ.rxx(gamma[0], qr[i_q], qr[i_q+i_conn])
    
    #circ.measure(qr,cr)
    return
    
#Define circuit ansatz
def circuit_ansatz(n,params,conn=1, layers=1, ansatz_type=[1,1]):
    qr=QuantumRegister(n)
    cr=ClassicalRegister(n)
    circ = QuantumCircuit(qr,cr)
    if ansatz_type[0]==0:
        length_beta=0
    elif ansatz_type[0]==1:
        length_beta=3*n
    elif ansatz_type[0]==2:
        length_beta=3
    elif ansatz_type[0]==3:
        length_beta=n
        
    if ansatz_type[1]==0:
        length_gamma=0
    elif ansatz_type[1]==1:
        length_gamma=int(n*conn-conn*(conn+1)/2.)            
    elif ansatz_type[1]==2:
        length_gamma=1

 
    for i_layer in range(layers-1):       
        
        beta=params[(length_beta+length_gamma)*i_layer:(length_beta+length_gamma)*i_layer+length_beta]
        gamma=params[(length_beta+length_gamma)*i_layer+length_beta:(length_beta+length_gamma)*(i_layer+1)]
        
        entangler(circ,qr,gamma,n,conn,ansatz_type[1])
        driver(circ,qr,beta,n,ansatz_type[0])
    
    gamma=params[(length_beta+length_gamma)*(layers-1):]
    entangler(circ,qr,gamma,n,conn,ansatz_type[1])
    circ.measure(qr, cr)
    
    return circ

def cost(counts, shots, target_states,tol=0.0001):
    cost=0
    for state in target_states:
        if state in counts:
            cost-=1./len(target_states)*np.log2(max(tol,counts[state]/shots))
        else:
            cost-=1./len(target_states)*np.log2(tol)
    return cost

from qiskit.providers.jobstatus import JobStatus

def run_iteration(circ, num_qubits, shots=100): 

    # submit task: define task (asynchronous)
    task_status=''
    while task_status != JobStatus.DONE:
        
            task = backend.run(circ, shots=shots)  
            #print("Job submitted")      
    
            # Get ID of submitted task
            task_id = task.job_id()
            #print('Task ID :', task_id)
            while (task_status == JobStatus.INITIALIZING) or (task_status == JobStatus.QUEUED)   or (task_status == JobStatus.VALIDATING) or (task_status == JobStatus.RUNNING) or (task_status==''):  
                time.sleep(1)
                try:
                    task_status=task.status()
                except:
                    print("Error querying status. Trying again.")
                    pass                    

        
        #print('Task status is', task_status)

    # get result
    counts = task.result().get_counts()

    return counts


#Run an iteration of the circuit and calculate its cost
def run_circuit_and_calc_cost(params, *args):
    n, conn, layers, ansatz_type, target_states, shots = args
    circ=circuit_ansatz(n, params, conn, layers,ansatz_type)
    counts=run_iteration(circ, n, shots=shots)
    iter_cost=cost(counts, shots, target_states)
    cost_history.append(iter_cost)
    print("Current cost:", iter_cost)
    return iter_cost

backend = provider.get_backend("ionq_simulator")

#Set problem parameters
r=2 #Number of rows — make sure r*c is less than qubit count
c=2 #Number of columns — make sure r*c is less than qubit count
n=r*c #Qubits needed
conn=2 #Connectivity of the entangler.

#Check on qubit count
if n>backend.configuration().n_qubits:
    raise Exception("Too many qubits")

#Check on conn
if conn>(n-1):
    raise Exception("Connectivity is too large")
    
target_states=generate_target_distribution(r,c)

# Check expected output
print('Bitstrings that should be generated by trained circuit are', target_states, 'corresponding to the following BAS patterns:\n')
for state in target_states:
    for i in range(r):
        print(state[c*i:][:c].replace('0','□ ').replace('1','■ '))
    print('')

# Choose entangler and driver type
layers=1 # Number of layers in each circuit
shots=100 # Number of shots per iteration
ansatz_type=[1,1]
random_init=False # Set to true for random initialization; will be slower to converge

# Set up args
beta=generate_beta(ansatz_type, random_init)
gamma=generate_gamma(ansatz_type, random_init, n, conn)
params=beta+gamma

base_bounds=(0,2.*math.pi)
bnds=((base_bounds , ) * len(params)) 

#Set up list to track cost history
cost_history=[]

max_iter=100 # max number of optimizer iterations

args=(n, conn, layers, ansatz_type, target_states, shots)
opts = {'disp': True, 'maxiter': max_iter, 'maxfev': max_iter, 'return_all': True}

#Visualize the circuit
circ=circuit_ansatz(n, params, conn,layers,ansatz_type)
circ.draw()

from qiskit.visualization import plot_histogram

#Train Circuit
result=noisyopt.minimizeSPSA(run_circuit_and_calc_cost, params, args=args, bounds=bnds, niter=max_iter, disp=False, paired=False)

print("Success: ", result.success)
print(result.message)
print("Final cost function is ", result.fun)
print("Min possible cost function is ", np.log2(len(target_states)))
print("Number of iterations is ", result.nit)
print("Number of function evaluations is ", result.nfev)
print("Number of parameters was ", len(params))
print("Approximate cost on hardware: $", 0.01*len(cost_history)*shots)

#Plot the evolution of the cost function
plt.figure()
plt.plot(cost_history)
plt.ylabel('Cost')
plt.xlabel('Iteration')


#Visualize the result
result_final=run_iteration(circuit_ansatz(n, result.x, conn,layers,ansatz_type), shots)
plot_histogram(result_final)

#Train Circuit
backend = provider.get_backend("ionq_qpu")
result=noisyopt.minimizeSPSA(run_circuit_and_calc_cost, params, args=args, bounds=bnds, niter=max_iter, disp=False, paired=False)

print("Success: ", result.success)
print(result.message)
print("Final cost function is ", result.fun)
print("Min possible cost function is ", np.log2(len(target_states)))
print("Number of iterations is ", result.nit)
print("Number of function evaluations is ", result.nfev)
print("Number of parameters was ", len(params))
print("Approximate cost on hardware: $", 0.01*len(cost_history)*shots)

#Plot the evolution of the cost function
plt.figure()
plt.plot(cost_history)
plt.ylabel('Cost')
plt.xlabel('Iteration')

#Visualize the result
result_final=run_iteration(circuit_ansatz(n, result.x, conn,layers,ansatz_type), shots)
plot_histogram(result_final)

