import pennylane as qml
from pennylane import numpy as np
import matplotlib.pyplot as plt

qml.drawer.use_style(style='sketch')

dev = qml.device('default.qubit', wires=['psi', 'a', 'b'])

def teleport(theta):
    # Prepare the state to be teleported
    qml.QubitStateVector([np.cos(theta), np.sin(theta)], wires='psi')
    
    # Prepare the Bell state
    qml.Hadamard(wires='a')
    qml.CNOT(wires=['a', 'b'])

    # Alice applies a reverse bell gate to her qubit, controlled by theta
    qml.CNOT(wires=['psi', 'a'])
    qml.Hadamard(wires='psi')

    qml.Barrier(['psi', 'a'], only_visual=True)

    # Measure alice's qubits
    a1 = qml.measure('psi')
    a2 = qml.measure('a')

    # Send the classical bits to Bob
    qml.cond(a2, qml.PauliX)('b')
    qml.cond(a1, qml.PauliZ)('b')

    return qml.density_matrix(wires='b')

qml.draw_mpl(teleport)(np.pi/4)
plt.show()

@qml.qnode(dev)
def teleport(theta):
    # Prepare the state to be teleported
    qml.QubitStateVector([np.cos(theta), np.sin(theta)], wires='psi')
    
    # Prepare the Bell state
    qml.Hadamard(wires='a')
    qml.CNOT(wires=['a', 'b'])

    # Alice applies a reverse bell gate to her qubit, controlled by theta
    qml.CNOT(wires=['psi', 'a'])
    qml.Hadamard(wires='psi')

    qml.Barrier(['psi', 'a'], only_visual=True)

    # Measure alice's qubits
    a1 = qml.measure('psi')
    a2 = qml.measure('a')

    # Send the classical bits to Bob
    qml.cond(a2, qml.PauliX)('b')
    qml.cond(a1, qml.PauliZ)('b')

    return qml.probs(wires='b')

qml.draw_mpl(teleport)(np.pi/4)
plt.show()

theta = np.pi/6
print("Input state: ", np.cos(theta), "|0> +", np.sin(theta), "|1>")
print()
print("Input probability vector: ")
print(np.array([np.cos(theta)**2, np.sin(theta)**2]))
print()
print("Output probability vector: ")
print(teleport(theta))