%run QECC_Shor9Q_functions.ipynb

# Qubits for our entangled state
q0 = QuantumRegister(3, name = 'q0')
q1 = QuantumRegister(3, name = 'q1')
q2 = QuantumRegister(3, name = 'q2')

# Ancillas for bit-flip errors
bf0 = AncillaRegister(2, name = 'bf0')
bf1 = AncillaRegister(2, name = 'bf1')
bf2 = AncillaRegister(2, name = 'bf2')

# Ancillas for phase-flip errors
pf = AncillaRegister(2, name = 'pf')

circuit = QuantumCircuit(q0, bf0, q1, bf1, q2, bf2, pf)

# Define our initial normalized state
initial_state = normalize_state(1.0 + 1.0j, 2.0 + 0.0j)
circuit.initialize(initial_state, 0)

# Draws the circuit
# circuit.draw(output = 'mpl', plot_barriers = False)

# Add controlled-NOTNOT to qubits 0, 3 and 6
circuit.append(cnotnot(), [q0[0], q1[0], q2[0]])

# Add Hadamard gates to qubits 0, 3 and 6
circuit.h(q0[0])
circuit.h(q1[0])
circuit.h(q2[0])

# Add controlled-NOTNOT in our three packs of three qubits
circuit.append(cnotnot(), [q0[0], q0[1], q0[2]])
circuit.append(cnotnot(), [q1[0], q1[1], q1[2]])
circuit.append(cnotnot(), [q2[0], q2[1], q2[2]])

circuit.barrier(q0, bf0, q1, bf1, q2, bf2, pf)

circuit.draw(output = 'mpl', plot_barriers = False)

_ = get_state(circuit, fancy_print = True)

# Add a source of random noise in our channels
circuit.append(random_noise(reveal_error = True), 
                            [q0[0], q0[1], q0[2],
                             q1[0], q1[1], q1[2],
                             q2[0], q2[1], q2[2]])

circuit.barrier(q0, bf0, q1, bf1, q2, bf2, pf)

circuit.draw(output = 'mpl', plot_barriers = False)

_ = get_state(circuit, fancy_print = True)

# Add Hadamard gates for the bit-flip ancillas
for index in range(2):
    circuit.h(bf0[index])
    circuit.h(bf1[index])
    circuit.h(bf2[index])

# Add controlled-ZZ gates in our ancillas
# First pack of three qubits
circuit.append(czz(), [bf0[0], q0[0], q0[1]])
circuit.append(czz(), [bf0[1], q0[1], q0[2]])

# Second pack of three qubits
circuit.append(czz(), [bf1[0], q1[0], q1[1]])
circuit.append(czz(), [bf1[1], q1[1], q1[2]])

# Third pack of three qubits
circuit.append(czz(), [bf2[0], q2[0], q2[1]])
circuit.append(czz(), [bf2[1], q2[1], q2[2]])

circuit.barrier(q0, bf0, q1, bf1, q2, bf2, pf)

# Add Hadamard gates for the bit-flip ancillas
for index in range(2):
    circuit.h(bf0[index])
    circuit.h(bf1[index])
    circuit.h(bf2[index])

# Add double controlled gates to correct bit-flip errors
# First pack of three qubits
circuit.append(invccnot(), [bf0[1], bf0[0], q0[0]])
circuit.ccx(bf0[0], bf0[1], q0[1])
circuit.append(invccnot(), [bf0[0], bf0[1], q0[2]])

# Second pack of three qubits
circuit.append(invccnot(), [bf1[1], bf1[0], q1[0]])
circuit.ccx(bf1[0], bf1[1], q1[1])
circuit.append(invccnot(), [bf1[0], bf1[1], q1[2]])

# Third pack of three qubits
circuit.append(invccnot(), [bf2[1], bf2[0], q2[0]])
circuit.ccx(bf2[0], bf2[1], q2[1])
circuit.append(invccnot(), [bf2[0], bf2[1], q2[2]])

circuit.barrier(q0, bf0, q1, bf1, q2, bf2, pf)


# BIT-FLIP ERRORS SOLVED! Let's go with the phase-flip errors

circuit.draw(output = 'mpl', plot_barriers = False)

_ = get_state(circuit, fancy_print = True)

# Add Hadamard gates for the phase-flip ancillas
for index in range(2):
    circuit.h(pf[index])

# Add controlled-6X gates in our ancillas
circuit.append(c6x(), [pf[0], q0[0], q0[1], q0[2], q1[0], q1[1], q1[2]])
circuit.append(c6x(), [pf[1], q1[0], q1[1], q1[2], q2[0], q2[1], q2[2]])

# Add Hadamard gates for the phase-flip ancillas
for index in range(2):
    circuit.h(pf[index])

# Add double controlled gates to correct phase-flip errors
circuit.append(invccz(), [pf[1], pf[0], q0[0]])
circuit.append(ccz(), [pf[1], pf[0], q1[0]])
circuit.append(invccz(), [pf[0], pf[1], q2[0]])

circuit.barrier(q0, bf0, q1, bf1, q2, bf2, pf)


# PHASE-FLIP ERRORS SOLVED! Let's go for the final manipulations in our state

circuit.draw(output = 'mpl', plot_barriers = False, fold = 40)

_ = get_state(circuit, fancy_print = True)

# Add controlled-NOTNOT in our three packs of three qubits
circuit.append(cnotnot(), [q0[0], q0[1], q0[2]])
circuit.append(cnotnot(), [q1[0], q1[1], q1[2]])
circuit.append(cnotnot(), [q2[0], q2[1], q2[2]])

# Add Hadamard gates to qubits 0, 3 and 6
circuit.h(q0[0])
circuit.h(q1[0])
circuit.h(q2[0])

# Add controlled-NOTNOT to qubits 0, 3 and 6
circuit.append(cnotnot(), [q0[0], q1[0], q2[0]])

circuit.draw(output = 'mpl', plot_barriers = False, fold = 40)

# Get the state vector
final_state = get_state(circuit, fancy_print = True)

check_error_to_ancilla(final_state)

# print(circuit.draw(output = 'latex_source', filename = '/tmp/circuit.tex'))

import qiskit.tools.jupyter
%qiskit_version_table

