# !pip3 install qiskit

# !pip3 install pylatexenc

import qiskit as q

import numpy as np

import matplotlib.pyplot as plt

qr = q.QuantumRegister(3, name="q")
crz, crx = q.ClassicalRegister(1, name="crz"), q.ClassicalRegister(1, name="crx")
teleport_circuit = q.QuantumCircuit(qr, crz, crx)
teleport_circuit.draw()

def epr_pair(qc, a, b):
  '''
  qc is the quantum circuit
  a, b are two different qubits
  '''
  qc.h(a)
  qc.cx(a, b)

epr_pair(teleport_circuit, qr[1], qr[2]) # we leave 0th qubit (this will be used later as Alice's qubit the she wants to teleport)
teleport_circuit.draw()

# assume alice owns q_1 and bob owns q_2
# now we are ready to take the transformations that Alice applies
# the qubit she wants to teleport is control bit and the epr member bit is the target bit 

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

teleport_circuit.barrier() # barrier ensures the epr pair creation part isn't hampered/optimized by whatever happens next

alice_steps(teleport_circuit, qr[0], qr[1])
teleport_circuit.draw()

# after this alice applies the measurement 
def alice_measure(qc, a, b):
  qc.barrier()
  qc.measure(a, 0) # measures the first qubit and put it into first classical register 
  # measure goes as 
  # Measure quantum bit into classical bit (tuples).

  # Args:
  #     qubit (QuantumRegister|list|tuple): quantum register
  #     cbit (ClassicalRegister|list|tuple): classical register
  qc.measure(b, 1)  




alice_measure(teleport_circuit, 0, 1) # 0th qubit is the one she wants to send and 1st qubit is one of the EPR pairs
teleport_circuit.draw() 

def bob_transform(qc, bob_bit, crz, crx):
  qc.x(bob_bit).c_if(crx, 1) #c_if Add classical condition on register classical and value val c_if(self, classical, val) 
  qc.z(bob_bit).c_if(crz, 1) # apply x and z gate only when if any of then are in the state '1'

teleport_circuit.barrier()
bob_transform(teleport_circuit, qr[2], crz, crx)
teleport_circuit.draw()

from qiskit.quantum_info import random_statevector
psi = random_statevector(2, seed=10)
print ('check the qubit Alice wants to teleport: ', psi, type(psi),  type(psi.data))

q.visualization.plot_bloch_multivector(psi)

from qiskit.extensions import Initialize
init_qubit = Initialize(psi.data)
init_qubit.label = "init"

inverse_init_qubit = init_qubit.gates_to_uncompute() 
#initialization start from zero to psi, we reverse this initialization so it takes us back to zero 

qc_new_teleport = q.QuantumCircuit(qr, crz, crx)
qc_new_teleport.append(init_qubit, [qr[0]])
qc_new_teleport.barrier()
epr_pair(qc_new_teleport, qr[1], qr[2])
qc_new_teleport.barrier()
alice_steps(qc_new_teleport, qr[0], qr[1])
alice_measure(qc_new_teleport, qr[0], qr[1])
bob_transform(qc_new_teleport, qr[2], crz, crx)
qc_new_teleport.draw()

from google.colab import drive
drive.mount('/content/drive')

path = '/content/drive/My Drive/'

style = {'backgroundcolor': 'lavender', 'dpi':200, 'subfontsize':10}

qc_new_teleport.draw(output='mpl',scale=0.8, style=style, filename=path+'teleport_1.png')

statevec_sim = q.Aer.get_backend('statevector_simulator')
qobj = q.assemble(qc_new_teleport)

out_vector = statevec_sim.run(qobj).result().get_statevector()
q.visualization.plot_bloch_multivector(out_vector,)

qc_teleport_check = q.QuantumCircuit(qr, crz, crx)
qc_teleport_check.append(init_qubit, [qr[0]])
qc_teleport_check.barrier()
epr_pair(qc_teleport_check, qr[1], qr[2])
qc_teleport_check.barrier()
alice_steps(qc_teleport_check, qr[0], qr[1])
alice_measure(qc_teleport_check, qr[0], qr[1])
bob_transform(qc_teleport_check, qr[2], crz, crx)
qc_teleport_check.append(inverse_init_qubit, [qr[2]])

style = {'backgroundcolor': 'lavender'}

qc_teleport_check.draw(output='mpl',scale=0.8, style=style)

cr_result = q.ClassicalRegister(1) # we need another register to measure if we get 0 or not 
qc_teleport_check.add_register(cr_result)
qc_teleport_check.measure(qr[2], 2) # measure the last qubit and put it into the last classical register
qc_teleport_check.draw(output='mpl', scale=0.8, style={'backgroundcolor': 'lavender', 'subfontsize':12})

simulator = q.Aer.get_backend(name='qasm_simulator')
results = q.execute(qc_teleport_check, backend=simulator,).result()
q.visualization.plot_histogram(results.get_counts(qc_teleport_check)) 

