import sys, getpass
try:
    sys.path.append("../../") # go to parent dir
    import Qconfig
    qx_config = {
        "APItoken": Qconfig.APItoken,
        "url": Qconfig.config['url']}
    print('Qconfig loaded from %s.' % Qconfig.__file__)
except:
    APItoken = getpass.getpass('Please input your token and hit enter: ')
    qx_config = {
        "APItoken": APItoken,
        "url":"https://quantumexperience.ng.bluemix.net/api"}
    print('Qconfig.py not found in qiskit-tutorial directory; Qconfig loaded using user input.')

import numpy as np

# importing the QISKit
import qiskit
from qiskit import QuantumCircuit, QuantumProgram

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

# useful additional packages 
from qiskit.tools.visualization import plot_state
from qiskit.tools.qi.qi import *

Q_program = QuantumProgram()

# Creating registers
qr = Q_program.create_quantum_register("qr", 2)
cr = Q_program.create_classical_register("cr", 2)

# hadamard on qubit-1 only
had = Q_program.create_circuit("had", [qr], [cr])
had.h(qr[1])

# CNOT gate with qubit 1 control, qubit 0 target (target for ibmqx4)
cnot = Q_program.create_circuit("cnot", [qr], [cr])
cnot.cx(qr[1], qr[0])

U_had = np.array([[1,1],[1,-1]])/np.sqrt(2)
# compute Choi-matrix from unitary
had_choi = outer(vectorize(U_had))
plot_state(had_choi)

# process tomography set for a quantum operation on qubit 1
had_tomo_set =  tomo.process_tomography_set([1])

# Generate process tomography preparation and measurement circuits
had_tomo_circuits = tomo.create_tomography_circuits(Q_program, 'had', qr, cr, had_tomo_set)
print('Tomography circuit labels for "had" circuit:')
for label in had_tomo_circuits:
    print(label)

backend = 'local_qasm_simulator'
shots = 1000
had_tomo_results = Q_program.execute(had_tomo_circuits, shots=shots, backend=backend)

had_process_data = tomo.tomography_data(had_tomo_results, 'had', had_tomo_set)
had_choi_fit = tomo.fit_tomography_data(had_process_data, options={'trace':2})
print('Process Fidelity = ', state_fidelity(vectorize(U_had)/2, had_choi_fit))
plot_state(had_choi_fit)

# Register API token for online backends
qiskit.register(qx_config['APItoken'], qx_config['url'])

# Use the IBM Quantum Experience
backend = 'ibmqx4'
# Take 500 shots for each measurement basis
# Note: reduce this number for larger number of qubits
shots = 500
# set max credits
max_credits = 8
# Run the experiment
hadqx_tomo_results = Q_program.execute(had_tomo_circuits, backend=backend, shots=shots, 
                                       max_credits=max_credits, timeout=500)
print(hadqx_tomo_results)

hadqx_process_data = tomo.tomography_data(hadqx_tomo_results, 'had', had_tomo_set)
hadqx_choi_fit = tomo.fit_tomography_data(hadqx_process_data, options={'trace':2})
print('Process Fidelity = ', state_fidelity(vectorize(U_had)/2, hadqx_choi_fit))
plot_state(hadqx_choi_fit)

#unitary matrix for CNOT with qubit 1 as control and qubit 0 as target.
U_cnot = np.array([[1,0,0,0],[0,1,0,0],[0,0,0,1],[0,0,1,0]])
# compute Choi-matrix from unitary
cnot_choi = outer(vectorize(U_cnot))
plot_state(cnot_choi)

# process tomography set for quantum operation on qubits 0 and 1
cnot_tomo_set = tomo.process_tomography_set([1, 0])

# Generate process tomography circuits
cnot_tomo_circuits = tomo.create_tomography_circuits(Q_program, 'cnot', qr, cr, cnot_tomo_set)

backend = 'local_qasm_simulator'
shots = 1000
cnot_tomo_results = Q_program.execute(cnot_tomo_circuits, backend=backend, shots=shots, timeout=300)
cnot_process_data = tomo.tomography_data(cnot_tomo_results, 'cnot', cnot_tomo_set)

cnot_choi_fit = tomo.fit_tomography_data(cnot_process_data, options={'trace':4})
print('Process Fidelity = ', state_fidelity(vectorize(U_cnot)/2, cnot_choi_fit/4))
plot_state(cnot_choi_fit)

# Register API token for online backends
qiskit.register(qx_config['APItoken'], qx_config['url'])

# Use the IBM Quantum Experience
backend = 'ibmqx4'
# Take 500 shots for each measurement basis
# Note: reduce this number for larger number of qubits
shots = 500
# set max credits
max_credits = 8
# Set batch size and complete
batch_size = 18
status = qiskit.get_backend(backend).status
if status['operational'] == False or status['pending_jobs'] > 5:
    print('Warning: the selected backend appears to be busy or unavailable at present; consider choosing a different one if possible.')
for j in range(8):
    batch_circs = cnot_tomo_circuits[j*batch_size:(j+1)*batch_size]
    if j == 0:        
        cnotqx_tomo_results = Q_program.execute(batch_circs, backend=backend, shots=shots,
                                                max_credits=max_credits, timeout=500)
    else:
        cnotqx_tomo_results += Q_program.execute(batch_circs, backend=backend, shots=shots,
                                                 max_credits=max_credits, timeout=500)
    msg = 'Batch %d/%d: %s' % (j+1, 8, cnotqx_tomo_results)
    print(msg)

cnotqx_process_data = tomo.tomography_data(cnotqx_tomo_results, 'cnot', cnot_tomo_set)
cnotqx_choi_fit = tomo.fit_tomography_data(cnotqx_process_data, options={'trace':4})
print('Process Fidelity = ', state_fidelity(vectorize(U_cnot)/2, cnotqx_choi_fit/4))
plot_state(cnotqx_choi_fit)

