#Imports
from itertools import product
import matplotlib.pyplot as plt
%matplotlib inline

from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, Result
from qiskit import available_backends, get_backend, execute, register, least_busy
from qiskit.tools.visualization import matplotlib_circuit_drawer, qx_color_scheme #require qiskit>0.5.5
%config InlineBackend.figure_format = 'svg'
style = qx_color_scheme()
style['compress'] = True
style['cregbundle'] = True
style['plotbarrier'] = True
circuit_drawer = lambda circuit: matplotlib_circuit_drawer(circuit, style=style)

# Creating registers
qa = QuantumRegister(1, 'alice')
qb = QuantumRegister(1, 'bob')
c = ClassicalRegister(2, 'c')

# Bell Measurement
bell_measurement = QuantumCircuit(qa, qb, c)
bell_measurement.cx(qa, qb)
bell_measurement.h(qa)
bell_measurement.measure(qa[0], c[0])
bell_measurement.measure(qb[0], c[1])
circuit_drawer(bell_measurement)

alice = {}
bob = {}

# eigenvalue
alice['+'] = QuantumCircuit(qa, qb, c) # do nothing
alice['-'] = QuantumCircuit(qa, qb, c)
alice['-'] .x(qa)

bob['+'] = QuantumCircuit(qa, qb, c) # do nothing
bob['-'] = QuantumCircuit(qa, qb, c)
bob['-'].x(qb)

# matrix
alice['Z'] = QuantumCircuit(qa, qb, c) # do nothing

alice['X'] = QuantumCircuit(qa, qb, c)
alice['X'].h(qa)

bob['W'] = QuantumCircuit(qa, qb, c)
bob['W'].h(qb)  
bob['W'].t(qb)
bob['W'].h(qb)
bob['W'].s(qb)

bob['V'] = QuantumCircuit(qa, qb, c)
bob['V'].h(qb)  
bob['V'].tdg(qb)
bob['V'].h(qb)
bob['V'].s(qb)

qc = alice['+'] + alice['Z'] + bob['+'] + bob['V']
circuit_drawer(qc)

qc = alice['-'] + alice['X'] + bob['+'] + bob['W']
circuit_drawer(qc)

backend = 'local_qasm_simulator'
shots = 1024
circuits = {}

for a_sign, a_basis, b_sign, b_basis in product(['+', '-'], ['Z', 'X'], ['+', '-'], ['V', 'W']):
    name = a_sign + a_basis + b_sign + b_basis
    circuits[name] = QuantumCircuit(qa, qb, c)
    circuits[name] += alice[a_sign]
    circuits[name] += alice[a_basis]
    circuits[name] += bob[b_sign]
    circuits[name] += bob[b_basis]  
    circuits[name].barrier(qa)
    circuits[name].barrier(qb)
    circuits[name] += bell_measurement

# Example
print('A quantum circuit of -X+V.')
circuit_drawer(circuits['-X+V'])

job = execute(circuits.values(), backend=backend, shots=shots)
result = job.result()

completely_mixed = {'00': 0, '01': 0, '10': 0, '11': 0}
for a_sign, a_basis, b_sign, b_basis in product(['+', '-'], ['Z', 'X'], ['+', '-'], ['V', 'W']):
    name = a_sign + a_basis + b_sign + b_basis
    completely_mixed['00'] += result.average_data(circuits[name], {'00': 1})/16
    completely_mixed['01'] += result.average_data(circuits[name], {'01': 1})/16
    completely_mixed['10'] += result.average_data(circuits[name], {'10': 1})/16
    completely_mixed['11'] += result.average_data(circuits[name], {'11': 1})/16

plt.bar(completely_mixed.keys(), completely_mixed.values(), width=0.5)
plt.xlabel('Outcome', fontsize=15)
plt.ylabel('Probability for the completely mixed state', fontsize=10)
plt.axhline(y=0.25, color='red', linewidth=1)

def E(result: Result, observable: dict, a_basis: str, b_basis: str) -> float:
    val = 0
    str2int = lambda x: 1 if x == '+' else -1
    for a_sign, b_sign in product(['+', '-'], ['+', '-']):
        name = a_sign + a_basis + b_sign + b_basis
        sign = str2int(a_sign) * str2int(b_sign)
        val += sign * result.average_data(circuits[name], observable)
    return val

def D(result: Result, outcome: str) -> float:
    val = 0
    for a_basis, b_basis in product(['Z', 'X'], ['V', 'W']):
        if outcome[0] == outcome[1] and a_basis == 'X' and b_basis == 'V':
            sign = -1
        elif outcome[0] != outcome[1] and a_basis == 'X' and b_basis == 'W':
            sign = -1
        else:
            sign = 1
        val += sign * E(result, {outcome: 1}, a_basis=a_basis, b_basis=b_basis)
    return val

for outcome in ['00', '01', '10', '11']:
    if completely_mixed[outcome] <= 1/4: # check the condition
        print('D is equal to {} for the outcome {}.'.format(D(result, outcome), outcome))

# Connecting to the IBM Quantum Experience
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.')

register(qx_config['APItoken'], qx_config['url'])

device_shots = 1024
device_name = least_busy(available_backends({'simulator': False, 'local': False}))
device = get_backend(device_name)
device_coupling = device.configuration['coupling_map']

print("the best backend is " + device_name + " with coupling " + str(device_coupling))

job = execute(circuits.values(), backend=device_name, shots=device_shots)
result = job.result()

completely_mixed = {'00': 0, '01': 0, '10': 0, '11': 0}
for a_sign, a_basis, b_sign, b_basis in product(['+', '-'], ['Z', 'X'], ['+', '-'], ['V', 'W']):
    name = a_sign + a_basis + b_sign + b_basis
    completely_mixed['00'] += result.average_data(circuits[name], {'00': 1})/16
    completely_mixed['01'] += result.average_data(circuits[name], {'01': 1})/16
    completely_mixed['10'] += result.average_data(circuits[name], {'10': 1})/16
    completely_mixed['11'] += result.average_data(circuits[name], {'11': 1})/16

plt.bar(completely_mixed.keys(), completely_mixed.values(), width=0.5)
plt.xlabel('Outcome', fontsize=15)
plt.ylabel('Probability for the completely mixed state', fontsize=10)
plt.axhline(y=0.25, color='red', linewidth=1)

for outcome in ['00', '01', '10', '11']:
    if completely_mixed[outcome] <= 1/4: # check the condition
        print('D is equal to {} for the outcome {}.'.format(D(result, outcome), outcome))

