import numpy as np
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, execute, BasicAer, IBMQ
from qiskit.visualization import plot_histogram, plot_bloch_multivector
from qiskit.extensions import Initialize

from qiskit.quantum_info import random_statevector


from utility import *

qr = QuantumRegister(3)    # Protocol uses 3 qubits
crz = ClassicalRegister(1) # and 2 classical bits
crx = ClassicalRegister(1) # in 2 different registers
teleportation_circuit = QuantumCircuit(qr, crz, crx)

def create_bell_pair(qc, a, b, print_state_vec=False):
    "creates a bell pair in qc using qubits a and b"
    qc.h(a)
    qc.cx(a,b)
    if print_state_vec :
        state = get_state_vector(qc)
        print (state)
        

def test_bell_pair():
    qc = QuantumCircuit(2)
    qc.h(0)
    qc.cx(0, 1)
    state = get_state_vector(qc)    
    print (state)
test_bell_pair()

qr = QuantumRegister(3)    
crz = ClassicalRegister(1) 
crx = ClassicalRegister(1)
teleportation_circuit = QuantumCircuit(qr, crz, crx)

create_bell_pair(teleportation_circuit, 1, 2)  
teleportation_circuit.draw(output='mpl')

def alice_gates(qc, psi, a):
    qc.cx(psi, a)
    qc.h(psi)

qr = QuantumRegister(3)    
crz = ClassicalRegister(1) 
crx = ClassicalRegister(1)
teleportation_circuit = QuantumCircuit(qr, crz, crx)
create_bell_pair(teleportation_circuit, 1, 2) 
teleportation_circuit.barrier()
alice_gates(teleportation_circuit, 0, 1)
teleportation_circuit.draw(output='mpl')


def measure_and_send(qc, a, b):
    """Measures qubits a & b and 'sends' the results to Bob"""
    qc.barrier()
    qc.measure(a,0)
    qc.measure(b,1)
    

#Set up
qr = QuantumRegister(3)    
crz = ClassicalRegister(1) 
crx = ClassicalRegister(1)
teleportation_circuit = QuantumCircuit(qr, crz, crx)

# step 1
create_bell_pair(teleportation_circuit, 1, 2) 

#step 2
teleportation_circuit.barrier()
alice_gates(teleportation_circuit, 0, 1)

#step 3
measure_and_send(teleportation_circuit, 0, 1)
teleportation_circuit.draw(output='mpl')


def bob_gates(qc, qubit, crz, crx):
    # Here we use c_if to control our gates with a classical
    # bit instead of a qubit
    qc.x(qubit).c_if(crx, 1) # Apply gates if the registers 
    qc.z(qubit).c_if(crz, 1) # are in the state '1'

#Set up
qr = QuantumRegister(3)    
crz = ClassicalRegister(1) 
crx = ClassicalRegister(1)
teleportation_circuit = QuantumCircuit(qr, crz, crx)

# step 1
create_bell_pair(teleportation_circuit, 1, 2) 

#step 2
teleportation_circuit.barrier()
alice_gates(teleportation_circuit, 0, 1)

#step 3
measure_and_send(teleportation_circuit, 0, 1)

# step 4
teleportation_circuit.barrier()
bob_gates(teleportation_circuit, 2, crz, crx)
teleportation_circuit.draw(output='mpl')


from qiskit.quantum_info import random_statevector
psi = random_statevector(2).data

print (psi)

plot_bloch_multivector(psi)

from IPython.display import display, Markdown, Latex
display(Latex('$\psi$'))

init_gate = Initialize(psi)

#Set up
qr = QuantumRegister(3)    
crz = ClassicalRegister(1) 
crx = ClassicalRegister(1)
qc = QuantumCircuit(qr, crz, crx)

psi = random_statevector(2).data
init_gate = Initialize(psi)
init_gate.label = "init"

# step 0
qc.append(init_gate, [0] )
qc.barrier()

# step 1
create_bell_pair(qc, 1, 2) 
qc.barrier()

#step 2
alice_gates(qc, 0, 1)

#step 3
measure_and_send(qc, 0, 1)

# step 4
#qc.barrier()
bob_gates(qc, 2, crz, crx)
qc.draw(output='mpl')


