from qutip_qip.circuit import QubitCircuit, Measurement, Gate
from qutip import basis, tensor, identity
from math import sqrt

"""
Explanation: Quantum Teleportation Circuit
End of explanation
"""


teleportation = QubitCircuit(3, num_cbits = 2, input_states = ["\psi", "0", "0", "c0", "c1"])

"""
Explanation: Introduction
This notebook introduces the basic quantum teleportation circuit (http://www.physics.udel.edu/~msafrono/425-2011/Lecture%2025-1.pdf), complete with measurements and classical control. This notebook also serves as an example on how to add measurement gates and classical controls to a circuit.
We will describe the circuit that enables quantum teleportation. We will use two classical wires and three qubit wires. The first qubit wire represents the quantum state $| q0 ⟩ = | \psi ⟩$ that needs to be transferred from Alice to Bob (so the first qubit is in the possession of Alice).
End of explanation
"""


teleportation.add_gate("SNOT", targets=[1])
teleportation.add_gate("CNOT", targets=[2], controls=[1])

"""
Explanation: First, Alice and Bob need to create the shared EPR pair ($\frac{| 00 ⟩ + | 11 ⟩} {2}$) between the second and third qubit by using the hadamard gate on Alice's qubit followed by an entangling CNOT gate.
End of explanation
"""


teleportation.add_gate("CNOT", targets=[1], controls=[0])
teleportation.add_gate("SNOT", targets=[0])

teleportation.add_measurement("M0", targets=[0], classical_store=1)
teleportation.add_measurement("M1", targets=[1], classical_store=0)

"""
Explanation: Following this, Alice makes the qubit $| q0 ⟩$ interact with Alice's EPR qubit, followed by measuring on the two qubits belonging to Alice. The measurement results for the first qubit is stored in classical register $c1$ and the second qubit is stored in classical register $c0$.
End of explanation
"""


teleportation.add_gate("X", targets=[2], classical_controls=[0])
teleportation.add_gate("Z", targets=[2], classical_controls=[1])

"""
Explanation: Now, we apply the $X$ gate on Bob's qubit based on the classical control $c0$ and $Z$ gate based on classical control $c1$. These operations correspond to the following operations based on the state of Alice's measurement. 
$|00⟩ \rightarrow $ no operation \
$|01⟩ \rightarrow Z$ \
$|10⟩ \rightarrow X$ \
$|11⟩ \rightarrow ZX$ 
The final circuit mathematically must result in the third qubit taking the state $|\psi⟩$.
End of explanation
"""


teleportation.gates

"""
Explanation: Finally, our teleportation circuit is ready to run, we can view the circuit structure using the following command.
End of explanation
"""


teleportation

"""
Explanation: The circuit can also be visualized:
End of explanation
"""


a = 1/sqrt(2) * basis(2, 0) + 1/sqrt(2) * basis(2, 1)
state = tensor(a, basis(2,0), basis(2,0))

"""
Explanation: The first qubit is user-specified $|\psi ⟩$ state and the other two must be $|0⟩$. 
Example 1
$|\psi⟩ = |+ ⟩$
End of explanation
"""


initial_measurement = Measurement("start", targets=[0])
initial_measurement.measurement_comp_basis(state)

"""
Explanation: We can confirm our state is initialized correctly by observing the measurment statistics on the first qubit, followed by which we run the circuit.
End of explanation
"""


state_final = teleportation.run(state)
print(state_final)

"""
Explanation: We can run the circuit using the QubitCircuit.run() function which provided the initial state-vector (or density matrix) initiates one run on the circuit (including sampling any intermediate measurements) and providing the final results (any classical bits can also be explicitly set using the argument cbits). The results are returned as a Result object. The result states can be accessed through the get_states() function where the argument index=0 specifies the first(only) result should be returned
End of explanation
"""


final_measurement = Measurement("start", targets=[2])
final_measurement.measurement_comp_basis(state_final)

"""
Explanation: After running, we can see the measurement statistics on the last qubit to see that the qubit is teleported correctly.
End of explanation
"""


state = tensor(basis(2,1), basis(2,0), basis(2,0))
initial_measurement = Measurement("start", targets=[0])
initial_measurement.measurement_comp_basis(state)

state_final = teleportation.run(state)
print(state_final)

final_measurement = Measurement("start", targets=[2])
final_measurement.measurement_comp_basis(state_final)

"""
Explanation: Example 2
$|\psi⟩ = |1 ⟩$
End of explanation
"""


results = teleportation.run_statistics(state)
results.probabilities

results.final_states

from qutip.ipynbtools import version_table
version_table()

"""
Explanation: Another useful feature of the circuit module is the QubitCircuit.run_statistics() feature which provides the opportunity to gather all the possible output states of the circuit along with their output probabilities. Again, the results are returned as a Result object. The result states and respective probabilites can be accessed through the get_results() function.
End of explanation
"""


