import math
import random

from qiskit_extension.quantum_circuit2 import QuantumCircuit2 as qc2
from qiskit_extension.state_vector2 import StateVector2 as sv2

state_unknown = sv2.from_label((random.random(), "0"), (random.random(), "1"))
state_unknown.show_state()

# Create a 2bits quantum circuit
circ_init = qc2(2)

# Entangle two ground state qubits to an EPR pair
circ_init.h(0)
circ_init.cx(0,1)

# Put the ground state into the circuit to get the EPR pair
state_EPR = sv2(circ_init)

# Show state_EPR, confirm it is EPR pair
state_EPR.show_state()

# Combined state_unknown and state_EPR into a 3-qubit sv2 object
state_before_teleport = state_unknown.expand(state_EPR)
state_before_teleport.show_state()

# Create the circuit for Bell measurement
circ_bell = qc2(3)
circ_bell.cx(0,1)
circ_bell.h(0)

# Bell measurement on qubits(0,1), and store the four possible states as a list after measurement
# list[0b00]=state after measurement result is 00, 
# list[0b01]=state after measurement result is 01, 
# ...
list_state_after_measure = state_before_teleport.evolve(circ_bell).measure([0,1])

# Show the four possible states after the Bell measurement
# Display format: (|00> indicates the measured state, followed by a colon indicating [the remaining state after measurement result is 00])
# |00>: 1/2|0> + 1/2|1> ...
state_before_teleport.evolve(circ_bell).show_measure([0,1])

state_after_teleport = list_state_after_measure[0b00]
state_after_teleport.show_state(hide=[0,1])

# Correction circuit
circ_correction = qc2(3)
circ_correction.x(2)

# Applying the correction circuit
state_after_teleport = list_state_after_measure[0b01].evolve(circ_correction)
state_after_teleport.show_state(hide=[0,1])

# Correction circuit
circ_correction = qc2(3)
circ_correction.z(2)

# Applying the correction circuit
state_after_teleport = list_state_after_measure[0b10].evolve(circ_correction)
state_after_teleport.show_state(hide=[0,1])

# Correction circuit
circ_correction = qc2(3)
circ_correction.x(2)
circ_correction.z(2)

# Applying the correction circuit
state_after_teleport = list_state_after_measure[0b11].evolve(circ_correction)
state_after_teleport.show_state(hide=[0,1])