import numpy as np
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister
from qiskit import IBMQ, Aer, transpile, assemble
from qiskit.visualization import plot_histogram, plot_bloch_multivector
from qiskit.extensions import Initialize
from qiskit_textbook.tools import random_state, array_to_latex
import matplotlib.pyplot
%pylab inline

#Defining srtucture of the circuit

qr = QuantumRegister(3, name="q")    
crz = ClassicalRegister(1, name="crz") 
crx = ClassicalRegister(1, name="crx") 
circuit = QuantumCircuit(qr, crz, crx)

#Creating a qubit in random state
psi = random_state(1)

# Display it nicely
array_to_latex(psi, pretext="|\\psi\\rangle =")
# Show it on a Bloch sphere
plot_bloch_multivector(psi)

init_gate = Initialize(psi)
init_gate.label = "init"

#append inti gate to the qc
circuit.append(init_gate, [0])
circuit.barrier()

#Creating entanglement in q1 and q2

circuit.h(1)
circuit.cx(1,2)
circuit.barrier()

#Defining gates for sender

circuit.cx(0,1)
circuit.h(0)
circuit.barrier()

#Measuring the bits and sending it to classical bits

circuit.measure((0,1), (0,1))

#reciever gates

circuit.x(2).c_if(crx, 1)
circuit.z(2).c_if(crz, 1)
circuit.draw('mpl')

#Using Statevector simuator

sv_sim = Aer.get_backend('statevector_simulator')
qobj = assemble(circuit)
out_vector = sv_sim.run(qobj).result().get_statevector()
plot_bloch_multivector(out_vector)

#To check the gates followed in order to confirm teleportation

circuit.barrier()
inverse_init_gate = init_gate.gates_to_uncompute()
circuit.append(inverse_init_gate, [2])

#To measure the third qubit

result_cb = ClassicalRegister(1)
circuit.add_register(result_cb)
circuit.measure(2,2)

circuit.draw('mpl')

#Using Quasm Simulator

qasm_sim = Aer.get_backend('qasm_simulator')
t_qc = transpile(circuit, qasm_sim)
qobj = assemble(t_qc)
counts = qasm_sim.run(qobj).result().get_counts()
plot_histogram(counts)