#Import packages

import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
import time

# import from qiskit
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, execute, compile
from qiskit.converters import qobj_to_circuits
from qiskit import Aer, IBMQ
from qiskit.providers.aer import noise

# import tomography library
import qiskit.tools.qcvv.tomography as tomo

# useful additional packages 
from qiskit.tools.visualization import plot_state, plot_histogram        
from qiskit.tools.qi.qi import state_fidelity, outer
from qiskit.tools.qi.qi import outer
from qiskit.quantum_info import state_fidelity
from qiskit.tools.monitor import job_monitor, backend_overview
from qiskit.providers.ibmq import least_busy

Aer.backends() # No need for credentials for running the next cells

# Determine the job
n = int(input("type number of qubits + enter: "))

# create a n-qubit quantum register
qr = QuantumRegister(n)
cr = ClassicalRegister(n)
my_state = QuantumCircuit(qr, cr, name='my_state')

# desired vector for a n-qubit W state
desired_vector = []
qr_vector = []
n_vector = 2**n
for i in range(n_vector):
        desired_vector.append(0)
for j in range(n) :           
        desired_vector[2**j] = 1/np.sqrt(n)

        
# choice of the circuit building method (arbitrary or specific)        
print("Do you want to use the specific method?")
W_state_circuit = input("    Answer by (y/n) + enter\n").upper()

if (W_state_circuit == "N") :
    method = "arbitrary"
# Initialize a n-qubit W quantum state using the arbitrary method   
    for j in range(n) :                   
        qr_vector.append(qr[j])
    my_state.initialize(desired_vector, qr_vector)
    
else:
# Quantum circuit to make a n-qubit W state using the specific method
    method = "specific"
    my_state.x(qr[n-1]) #start is |10...0>
    for i in range(1,n) :
        theta = np.arccos(np.sqrt(1/(n-i+1)))
        my_state.ry(-theta,qr[n-i-1])       
        my_state.cz(qr[n-i],qr[n-i-1])
        my_state.ry(theta,qr[n-i-1])
    for i in range(1,n) :
        my_state.cx(qr[n-i-1],qr[n-i])

# Measurement circuit
measuring = QuantumCircuit(qr, cr, name='measuring')
for i in range(n) :
    measuring.measure(qr[i] , cr[i])
test = my_state+measuring

# Test circuit "my_state" : Noise free model on simulator
backend_sim = Aer.get_backend('qasm_simulator')
shots = 1024
job_noisefree = execute(test, backend_sim, shots=shots, max_credits=5)
job_monitor(job_noisefree)

noisefree_count = job_noisefree.result().get_counts(test)
time_exp = time.strftime('%d/%m/%Y %H:%M:%S')
print("Date (DMY):", time_exp)
print(noisefree_count)
plot_histogram(noisefree_count, color=['purple'], 
               title=str(n) + '- qubit W state, noise free simulation, ' +  method + " method")

# QASM from test
QASM_source = test.qasm()
print(QASM_source)

# Draw the circuit 
test.draw(output='mpl')

# Execute state tomography using noise free quantum device simulation

# obtain the final state vector
backend_stvct = Aer.get_backend('statevector_simulator')
job = execute(my_state, backend_stvct)
my_state_psi = job.result().get_statevector(my_state)

# construct state tomography set for measurement of qubits [0, ..., n-1] in the Pauli basis
qubit_set = []
for i in range(0,n) :
    qubit_set.append(i)

my_state_tomo_set = tomo.state_tomography_set(qubit_set) # default value for meas_basis ='Pauli'.

# add the state tomography measurement circuits to the Quantum Program
my_state_tomo_circuits = tomo.create_tomography_circuits(my_state, qr, cr, my_state_tomo_set)

backend_tomo = Aer.get_backend('qasm_simulator') # for simulation

# take 1024 shots for each measurement basis
# note: reduce this number for larger number of qubits
shots = 1024
mode = "noise free simulation"
my_state_job = execute(my_state_tomo_circuits, backend_tomo, shots=shots)
job_monitor(my_state_job)
my_state_tomo_result = my_state_job.result()    
        
# extract tomography data from results
my_state_tomo_data = tomo.tomography_data(my_state_tomo_result, my_state.name, my_state_tomo_set)

# Quantum fidelity

# reconstruct experimentally measured density matrix 
rho_fit = tomo.fit_tomography_data(my_state_tomo_data)

# calculate fidelity of fitted state:
time_exp = time.strftime('%d/%m/%Y %H:%M:%S')
print("Date (DMY):", time_exp)
print('Tomography',str(n)+'-qubit W state on', backend_tomo,
      ", shots:", shots, ", method:", method, ", mode:", mode)
F_fit = state_fidelity(rho_fit, my_state_psi)
print('Fidelity with theoretical ideal state')
print('F =', F_fit)

IBMQ.load_accounts()

backend_overview() # you may skip running this cell if you want

backend_real = least_busy(IBMQ.backends(operational=True, simulator=False))
print(backend_real)

# Prepare device noise simulation (DNS)

device = backend_real
print("device: ", device)
properties = device.properties()
coupling_map = device.configuration().coupling_map
prepared = False
if device.name() == 'ibmq_16_melbourne' :
    gate_times = [
        ('u1', None, 0), ('u2', None, 100), ('u3', None, 200),
        ('cx', [1, 0], 678), ('cx', [1, 2], 547), ('cx', [2, 3], 721),
        ('cx', [4, 3], 733), ('cx', [4, 10], 721), ('cx', [5, 4], 800),
        ('cx', [5, 6], 800), ('cx', [5, 9], 895), ('cx', [6, 8], 895),
        ('cx', [7, 8], 640), ('cx', [9, 8], 895), ('cx', [9, 10], 800),
        ('cx', [11, 10], 721), ('cx', [11, 3], 634), ('cx', [12, 2], 773),
        ('cx', [13, 1], 2286), ('cx', [13, 12], 1504), ('cx', [], 800)
    ]
    prepared = True
    
elif device.name() == 'ibmqx4' :
    gate_times = [
        ('u1', None, 0), ('u2', None, 60), ('u3', None, 120),
        ('cx', [1, 0], 340), ('cx', [2, 0], 424), ('cx', [2, 1], 520),
        ('cx', [3, 2], 620), ('cx', [3, 4], 420), ('cx', [4, 2], 920)    
    ]
    prepared = True

else :
    print("No gate times yet defined in this notebook for: ", device)
        
if prepared :
    # Construct the noise model from backend properties and custom gate times
    noise_model = noise.device.basic_device_noise_model(properties, gate_times=gate_times)

    # Get the basis gates for the noise model
    basis_gates = noise_model.basis_gates
    
    print("noise model prepared for", device)


# Execute test using device noise simulation (DNS)

backend_noise = Aer.get_backend('qasm_simulator') # for simulation (DNS)

shots = 1024
mode = "DNS"
job_noise = execute(test, backend_noise, 
                       noise_model=noise_model,
                       coupling_map=coupling_map,
                       basis_gates=basis_gates)
job_monitor(job_noise)

print(job_noise.status)
time_exp = time.strftime('%d/%m/%Y %H:%M:%S')
print("Date (DMY):", time_exp)
noisy_count = job_noise.result().get_counts(test) 
print(noisy_count)
plot_histogram(noisy_count, color=['cyan'],
               title= str(n) + '- qubit W state, ' + mode + ': {}, '.format(device.name()) + method + " method")

# Execute state tomography using device noise simulation (DNS)

# obtain the final state vector
backend_stvct = Aer.get_backend('statevector_simulator')
job = execute(my_state, backend_stvct)
my_state_psi = job.result().get_statevector(my_state)

# construct state tomography set for measurement of qubits [0, ..., n-1] in the Pauli basis
qubit_set = []
for i in range(0,n) :
    qubit_set.append(i)
my_state_tomo_set = tomo.state_tomography_set(qubit_set) # default value for meas_basis ='Pauli'.    

# add the state tomography measurement circuits to the Quantum Program
my_state_tomo_circuits = tomo.create_tomography_circuits(my_state, qr, cr, my_state_tomo_set)

backend_tomo = Aer.get_backend('qasm_simulator') # for simulation

# take 1024 shots for each measurement basis
# note: reduce this number for larger number of qubits
shots = 1024
mode = "DNS"
my_state_job = execute(my_state_tomo_circuits, backend_tomo, 
                               noise_model=noise_model,
                               coupling_map=coupling_map,
                               basis_gates=basis_gates)
job_monitor(my_state_job)
my_state_tomo_result = my_state_job.result()    

# extract tomography data from results
my_state_tomo_data = tomo.tomography_data(my_state_tomo_result, my_state.name, my_state_tomo_set)

# Quantum fidelity

# reconstruct experimentally measured density matrix 
rho_fit = tomo.fit_tomography_data(my_state_tomo_data)

# calculate fidelity of fitted state:
time_exp = time.strftime('%d/%m/%Y %H:%M:%S')
print("Date (DMY):", time_exp)
print('Tomography',str(n)+'-qubit W state on', backend_tomo,
      ", shots:", shots, ", method:", method, ", mode:", mode, "of", device)
F_fit = state_fidelity(rho_fit, my_state_psi)
print('Fidelity with theoretical ideal state')
print('F =', F_fit)

plt.xlabel('Number of Entangled Qubits')
plt.ylabel('Quantum Fidelity')
plt.axis([1.4, 6, 0, 1])
plt.grid()
plt.plot([2,3,4,5], [0.788,0.734,0.628,0.542], 'mo-', label="Specific Algorithm - DNS")
plt.plot([2,3,4,5], [0.773,0.682,0.386,0.148], 'cs-', label="Arbitrary Initialization - DNS")
plt.legend()
plt.show()

# Execute test using superconducting quantum computing device (SQC)

#Choose the backend
#backend_noise = Aer.get_backend('qasm_simulator') # for optional test before final experiment 
backend_noise = device # for least busy SQC device

# Execute on SQC and get counts
shots = 1024
if backend_noise.name() == "qasm_simulator" : # optional test before final experiment 
    mode = "DNS"
    job_noise = execute(test, backend_noise, 
                           noise_model=noise_model,
                           coupling_map=coupling_map,
                           basis_gates=basis_gates)
else: # final experiment on real device
    mode = "SQC"
    job_noise = execute(test, backend_noise)
    job_monitor(job_noise)

print(job_noise.status)
time_exp = time.strftime('%d/%m/%Y %H:%M:%S')
noisy_count = job_noise.result().get_counts(test) 
print(noisy_count)
plot_histogram(noisy_count, color=['orange'], 
               title= str(n) + '- qubit W state, ' + mode + ': {}, '.format(device.name()) + method + " method")

# Execute state tomography on superconducting quantum computing device (SQC)

# obtain the final state vector
backend_stvct = Aer.get_backend('statevector_simulator')
job = execute(my_state, backend_stvct)
my_state_psi = job.result().get_statevector(my_state)

# construct state tomography set for measurement of qubits [0, ..., n-1] in the Pauli basis
qubit_set = []
for i in range(0,n) :
    qubit_set.append(i)

my_state_tomo_set = tomo.state_tomography_set(qubit_set) # default value for meas_basis ='Pauli'.

# add the state tomography measurement circuits to the Quantum Program
my_state_tomo_circuits = tomo.create_tomography_circuits(my_state, qr, cr, my_state_tomo_set)

#Choose the backend
#backend_tomo = Aer.get_backend('qasm_simulator') # optional test before final experiment
backend_tomo = device # for least busy SQC device

# take 1024 shots for each measurement basis
# note: reduce this number for larger number of qubits
shots = 1024

# loop: 27 circuits maximum per job to avoid exceeding the allowed limit for the real device.  
n_circ = 3**n
i_max = min(27,n_circ)
my_jobs = []
index_job = -1
for i in range(0,n_circ,i_max) :
    circs =[]
    for j in range(i, i+i_max):
        circs.append(my_state_tomo_circuits[j])
    if backend_tomo.name() == "qasm_simulator" : # optional test before final experiment
        mode = "DNS"
        my_state_job = execute(circs, backend_tomo, 
                               noise_model=noise_model,
                               coupling_map=coupling_map,
                               basis_gates=basis_gates)
    else: # final experiment on real device
        mode = "SQC" 
        my_state_job = execute(circs, backend_tomo, shots=shots)
    my_jobs.append(my_state_job)
    index_job = index_job + 1 
    job_monitor(my_jobs[index_job], monitor_async = True)
    
    my_state_new_result = my_state_job.result()
    if i == 0:
        my_state_tomo_result = my_state_new_result
    else:
        my_state_tomo_result = my_state_tomo_result +  my_state_new_result

# extract tomography data from results
my_state_tomo_data = tomo.tomography_data(my_state_tomo_result, my_state.name, my_state_tomo_set)

# Quantum fidelity

# reconstruct experimentally measured density matrix 
rho_fit = tomo.fit_tomography_data(my_state_tomo_data)

# calculate fidelity of fitted state:
time_exp = time.strftime('%d/%m/%Y %H:%M:%S')
print("Date (DMY):", time_exp)
print('Tomography',str(n)+'-qubit W state on', backend_tomo,
      ", shots:", shots, ", method:", method, ", mode:", mode, device)
F_fit = state_fidelity(rho_fit, my_state_psi)
print('Fidelity with theoretical ideal state')
print('F =', F_fit)

plt.xlabel('Number of Entangled Qubits')
plt.ylabel('Quantum Fidelity')
plt.axis([1.4, 6, 0, 1])
plt.grid()
plt.plot([2,3,4,5], [0.812,0.735,0.592,0.457], 'ro-', label="Specific Algorithm")
plt.plot([2,3,4,5], [0.838,0.539,0.144,0.060], 'bs-', label="Arbitrary Initialization")
plt.legend()
plt.show()

# DSN vs SQC, count histograms
# Shots = 1024
count_noisefree = {'010': 344, '001': 335, '100': 345}
count_DNS = {'110': 42, '010': 258, '111': 9, '011': 38, '000': 76, '001': 281, '100': 286, '101': 34}
count_real = {'001': 274, '000': 129, '010': 224, '011': 20, '100': 282, '111': 14, '101': 46, '110': 35} 
plot_histogram([count_noisefree, count_DNS, count_real], 
               title= 'Determistic W state generation, specific algorithm, device: ibmqx4', 
               color=['purple','cyan', 'orange'], bar_labels=False,
               legend = ['Noise free simulation', 'Device noise simulation','Real quantum device']) 

plt.xlabel('Number of Entangled Qubits')
plt.ylabel('Quantum Fidelity')
plt.axis([1.4, 6, 0, 1])
plt.grid()
plt.plot([2,3,4,5], [0.812,0.735,0.592,0.457], 'ro-', label="Specific Algorithm - SQC")
plt.plot([2,3,4,5], [0.838,0.539,0.144,0.060], 'bs-', label="Arbitrary Initialization - SQC")
plt.plot([2,3,4,5], [0.788,0.734,0.628,0.542], 'mo-', label="Specific Algorithm - DNS")
plt.plot([2,3,4,5], [0.773,0.682,0.386,0.148], 'cs-', label="Arbitrary Initialization - DNS")
plt.legend()
plt.show()