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
from time import sleep  # used for polling jobs
    
# importing the QISKit
import qiskit
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, QuantumProgram

# 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, concurrence, purity, outer

# Create a 2-qubit quantum register
qr = QuantumRegister(2)
cr = ClassicalRegister(2)

# quantum circuit to make an entangled Bell state 
bell = QuantumCircuit(qr, cr, name='bell')
bell.h(qr[1])
bell.cx(qr[1], qr[0])

job = qiskit.execute(bell, backend='local_statevector_simulator')
bell_psi = job.result().get_statevector(bell)
bell_rho = outer(bell_psi) # construct the density matrix from the state vector

# plot the state
plot_state(bell_rho,'paulivec')

rho_mixed = np.array([[1,0,0,0],[0,0,0,0],[0,0,0,0],[0,0,0,1]])/2
plot_state(rho_mixed, 'paulivec')

# Construct state tomography set for measurement of qubits [0, 1] in the Pauli basis
bell_tomo_set = tomo.state_tomography_set([0, 1])

# Create a quantum program containing the state preparation circuit
Q_program = QuantumProgram()
Q_program.add_circuit('bell', bell)

# Add the state tomography measurement circuits to the Quantum Program
bell_tomo_circuit_names = tomo.create_tomography_circuits(Q_program, 'bell', qr, cr, bell_tomo_set)

print('Created State tomography circuits:')
for name in bell_tomo_circuit_names:
    print(name)

# Use the local simulator
backend = 'local_qasm_simulator'

# Take 5000 shots for each measurement basis
shots = 5000

# Run the simulation
bell_tomo_result = Q_program.execute(bell_tomo_circuit_names, backend=backend, shots=shots)
print(bell_tomo_result)

bell_tomo_data = tomo.tomography_data(bell_tomo_result, 'bell', bell_tomo_set)

rho_fit = tomo.fit_tomography_data(bell_tomo_data)

# calculate fidelity, concurrence and purity of fitted state
F_fit = state_fidelity(rho_fit, bell_psi)
con = concurrence(rho_fit)
pur = purity(rho_fit)

# plot 
plot_state(rho_fit, 'paulivec')
print('Fidelity =', F_fit)
print('concurrence = ', str(con))
print('purity = ', str(pur))

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

# Use the IBM Quantum Experience
backend = 'ibmqx4'
# Take 1000 shots for each measurement basis
# Note: reduce this number for larger number of qubits
shots = 1000
# set max credits
max_credits = 8

# Run the experiment
status = qiskit.get_backend(backend).status
if status['operational'] == False or status['pending_jobs'] > 10:
    print('Warning: the selected backend appears to be busy or unavailable at present; consider choosing a different one if possible')
bellqx_tomo_result = Q_program.execute(bell_tomo_circuit_names, backend=backend, shots=shots,
                           max_credits=max_credits, timeout=500)
print(bellqx_tomo_result)

bellqx_tomo_data = tomo.tomography_data(bellqx_tomo_result, 'bell', bell_tomo_set)

# Reconstruct experimentally measured density matrix 
rho_fit_real = tomo.fit_tomography_data(bellqx_tomo_data)

F_fit_real = state_fidelity(rho_fit_real, bell_psi)
plot_state(rho_fit_real, 'paulivec')
print('Fidelity with ideal state')
print('F =', F_fit_real)

# calculate concurrence and purity
con = concurrence(rho_fit_real)
pur = purity(rho_fit_real)
print('concurrence = ', str(con))
print('purity = ', str(pur))

