#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

def get_qc(p):
    qr = QuantumRegister(11,'q')
    qc = QuantumCircuit(qr)

    for i in range(1,12):
        qc.rx(i*p,i-1)
    
    return qc

qci = get_qc(pi/3)
qci.draw()

qca = get_qc(pi/3+pi/18)
qcr = get_qc(pi/3-pi/18)

from qiskit import ClassicalRegister

# Add the measurement register
circs = [qci,qca,qcr]
cr = ClassicalRegister(11,'c')
for qc in circs:
    qc.add_register(cr)
    qc.measure(range(11),range(11))

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'],11,nshots))
        results.append(get_pops(result.data()['counts'],11,nshots))
    
    return results

results = get_jobs(jobs,1000)

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

prbi = [get_ion(results[0],i) for i in range(11)]
prba = [get_ion(results[1],i) for i in range(11)]
prbr = [get_ion(results[2],i) for i in range(11)]
avres = [sum([results[i][j] for i in range(3)])/3 for j in range(2**11)]
prbs = [get_ion(avres,i) for i in range(11)]

import matplotlib.pyplot as plt
plt.plot(range(12),[0]+prbi,label="Ideal")
plt.plot(range(12),[0]+prba,label="Over-rotation")
plt.plot(range(12),[0]+prbr,label="Under-rotation")
plt.plot(range(12),[0]+prbs,label="Average",linewidth=3)
plt.ylim([0,1])
plt.ylabel("State |1> probability")
plt.xlabel("Qubit (*time*)")
plt.legend()

# 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)
prbi_m = [get_ion(results[0],i) for i in range(11)]
prba_m = [get_ion(results[1],i) for i in range(11)]
prbr_m = [get_ion(results[2],i) for i in range(11)]
avres = [sum([results[i][j] for i in range(3)])/3 for j in range(2**11)]
prbs_m = [get_ion(avres,i) for i in range(11)]

plt.plot(range(12),[0]+prbi_m,label="Ideal")
plt.plot(range(12),[0]+prba_m,label="Over-rotation")
plt.plot(range(12),[0]+prbr_m,label="Under-rotation")
plt.plot(range(12),[0]+prbs_m,label="Average",linewidth=3)
plt.ylim([0,1])
plt.ylabel("State |1> probability")
plt.legend()