%autosave 0
%matplotlib inline
from sglib import *
set_font_size(3)
from sympy.physics.quantum.qubit import matrix_to_qubit
from sympy.physics.quantum.qubit import qubit_to_matrix, Qubit
def mtq(bits):
    return(latex(matrix_to_qubit(bits).simplify()))

# Define the single bit Identity, Not, and Hadamard Operators
I = Matrix([[1,0],[0,1]])
X = Matrix([[0,1],[1,0]])
Z = Matrix([[1,0],[0,-1]])
H = 1/sqrt(2)*Matrix([[1,1],[1,-1]])

# Define the 2bit CNOT, which I just call "C2"
C2 = Matrix([[1,0,0,0],[0,1,0,0],[0,0,0,1],[0,0,1,0]])

# C3 is a three bit operator that actually does a CNOT
# on the first two bits, and leaves the third bit alone
C3 = TP(C2,I)

# H3 is a three bit operator that actually does a Hadamard
# on the first bit only, leaving the other two bits alone.
H3 = TP(TP(H,I),I)

# Take a look at them all
Print('C2 = $%s$'%myltx(C2), font_size=2)
print
Print('C3 = $%s$'%myltx(C3), font_size=2)
print
Print('H3 = $%s$'%myltx(H3), font_size=2)

# In general, I haven't figured out how you're actually supposed to
# do operations on Sympy "Qubits" so I'm keeping my state in Vectors
# (which are actually Sympy one-column matrices). Then I can represent
# the operations as matrices and just do normal matrix multiplication
# to do the operations.

# So here is a classical zero (q0) and one (q1)
q0 = Matrix([[1],[0]]); q1 = Matrix([[0],[1]])

# Set up the three intitial qubits. Right now I just start
# off with Alice and Bob's pair already entangled.
q00=TP(q0,q0); q01=TP(q0,q1); q10=TP(q1,q0); q11=TP(q1,q1)
a=sy.Symbol('a'); b=sy.Symbol('b')
bit1=a*q0+b*q1
Print('Step 1')
Print('The bit Alice wants to transmit: $%s$'%mtq(bit1))
bits2and3 = 1/sqrt(2)*(q00 + q11)
Print('The bits Alice and Bob share: $%s$'%mtq(bits2and3))
bits_t0 = TP(bit1,bits2and3)
Print('All three bits: $%s$'%mtq(bits_t0))

# Alice applies the CNOT and Hadamard
Print('Step 2')
bits_t1 = C3*bits_t0
Print('Alice applies the cnot, resulting in:')
Print('$%s$'%mtq(bits_t1))

Print('Step 3')
bits_t2 = H3*bits_t1
Print('Alice applies the Hadamard, resulting in:')
Print('$%s$'%mtq(bits_t2))

one = TP(Qubit('00'), a*Qubit('0')+b*Qubit('1'))
two = TP(Qubit('01'), a*Qubit('1')+b*Qubit('0'))
three = TP(Qubit('10'), a*Qubit('0')-b*Qubit('1'))
four = TP(Qubit('11'), a*Qubit('1')-b*Qubit('0'))
alt = Rational(1,2)*(one+two+three+four)
Print('(Still on step 3)')
Print('The same bits can be expressed in (partially) factored form as:')
Print('$%s$'%latex(alt))
Print('The above should match (1.33) through (1.36) at the bottom of Nielsen and Chuang page 27.')

if00 = a*q0+b*q1 # (No operation needed)
if01 = X*(a*q1+b*q0)
if10 = Z*(a*q0-b*q1)
if11 = Z*X*(a*q1-b*q0)
if00, if01, if10, if11