#import Aer here, before calling qiskit_ionq_provider
from qiskit import Aer

from qiskit_ionq import IonQProvider 

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

provider.backends()

from qiskit import QuantumCircuit, QuantumRegister
from math import pi

qr = QuantumRegister(6,'q')
qc_par = QuantumCircuit(qr)

for i in range(3):
    qc_par.rxx(pi/2,i,i+3)

qc_par.draw()

qr = QuantumRegister(6,'q')
qc_star = QuantumCircuit(qr)

for i in range(1,6):
    qc_star.rxx(pi/2,0,i)

qc_star.draw()

from qiskit import ClassicalRegister

# Add the measurement register
circs = [qc_par,qc_star]
cr = ClassicalRegister(6,'c')
for qc in circs:
    qc.add_register(cr)
    qc.measure(range(6),range(6))

from qiskit.providers.jobstatus import JobStatus
from qiskit import Aer, execute

# Choose the simulator backend 
backend = provider.get_backend("ionq_simulator")
#backend = Aer.get_backend("qasm_simulator")

# Run the circuit:
def run_jobs(backend,circs,nshots):
    jobs = []
    job_ids = []
    qcs = []

    for qc in circs:
        qcs.append(qc)
        job = backend.run(qc, shots=nshots)
        #job = execute(qc, backend, shots=nshots, memory=True)
        jobs.append(job)
        #job_ids.append(job.job_id())
    
    return jobs

jobs = run_jobs(backend,circs,1000)

# Calculate output state populations
def get_pops(res,nn,n):
    #print(res)
    pops = [0 for i in range(2**nn)]
    for key in res.keys():
        pops[int(key,16)] = res[key]/n
        #pops[int(key,2)] = res[key]/n
    return pops

# Fetch the result
def get_jobs(jobs,nshots):
    results = []
    for i in range(len(jobs)):
        result = jobs[i].result()
        print(result.data()['counts'])
        print(get_pops(result.data()['counts'],6,nshots))
        results.append(get_pops(result.data()['counts'],6,nshots))
    
    return results

results = get_jobs(jobs,1000)

def get_ion(res,ion):
    p1 = 0
    for x in range(2**6):
        if (x&(2**ion)):
            p1 += res[x]
    return p1

def get_pair(res,pair):
    p00 = 0
    p01 = 0
    p10 = 0
    p11 = 0
    for x in range(2**6):
        if (x&(2**pair[0])>0) and (x&(2**pair[1])>0):
            p11 += res[x]
        elif (x&(2**pair[0])>0) and (x&(2**pair[1])==0):
            p01 += res[x]
        elif (x&(2**pair[1])>0) and (x&(2**pair[0])==0):
            p10 += res[x]
        elif (x&(2**pair[0])==0) and (x&(2**pair[1])==0):
            p00 += res[x]
    return [p00,p01,p10,p11]

prsq = [get_pair(results[0],[i,i+3]) for i in range(3)]
prbi = [get_ion(results[1],i) for i in range(6)]
print(prbi)
print(prsq)

import matplotlib.pyplot as plt
for i in range(3):
    plt.bar([x-0.3+0.3*i for x in range(4)],prsq[i],width=0.3,label="Pair "+str(i)+"-"+str(i+3))

plt.ylim([0,1])
plt.xticks(range(4), [bin(x)[2:].zfill(2) for x in range(4)],rotation=30)
plt.ylabel("Output probability")
plt.xlabel("Output state")
plt.legend()

plt.bar(range(6),prbi)
plt.ylim([0,1])
plt.ylabel("State |1> probability")
plt.xlabel("Qubit")



# Switch the backend to run circuits on a quantum computer
qpu_backend = provider.get_backend("ionq_qpu")

jobs = run_jobs(qpu_backend,circs,1000)

#Check if jobs are done
for i in range(len(jobs)):
    print(jobs[i].status())

# Fetch the result
results = get_jobs(jobs,1000)
prsq_m = [get_pair(results[0],[i,i+3]) for i in range(3)]
prbi_m = [get_ion(results[1],i) for i in range(6)]

for i in range(3):
    plt.bar([x-0.3+0.3*i for x in range(4)],prsq_m[i],width=0.3,label="Pair "+str(i)+"-"+str(i+3))

plt.ylim([0,1])
plt.xticks(range(4), [bin(x)[2:].zfill(2) for x in range(4)],rotation=30)
plt.ylabel("Output probability")
plt.xlabel("Output state")
plt.legend()

plt.bar(range(6),prbi_m)
plt.ylim([0,1])
plt.ylabel("State |1> probability")
plt.xlabel("Qubit")

