import numpy as np

# These column vectors can be stored in numpy arrays so that we can operate 
# on them with the circuit diagram's corresponding matrix (which is to be evaluated)
# as follows:
zero_zero = np.array([[1],[0],[0],[0]])
zero_one = np.array([[0],[1],[0],[0]])
one_zero = np.array([[0],[0],[1],[0]])
one_one = np.array([[0],[0],[0],[1]])
Psi = {'zero_zero': zero_zero, 'zero_one': zero_one, 'one_zero': one_zero, 'one_one': one_one}
# ^We can conveniently store all possible input states in a dictionary and then print to check the representations:
for key, val in Psi.items():
    print(key, ':', '\n', val)

# storing CNOT as a numpy array:
CNOT_1 = np.matrix([[1, 0, 0, 0],[0, 1, 0, 0],[0, 0, 0, 1],[0, 0, 1, 0]])
print(CNOT_1)

print('FINAL STATE OF i):')
#Apply CNOT to each possible state for |Psi> to find --> |Psi'>
for key, val in Psi.items():
    print(key, 'becomes..\n', CNOT_1*val)

# storing this in a numpy array:
H_2 = .5*np.matrix([[1, 1, 1, 1],[1, -1, 1, -1],[1, 1, -1, -1],[1, -1, -1, 1]])
print('H_2:')
print(H_2)

# storing this in a numpy array:
CNOT_2 = np.matrix([[1, 0, 0, 0],[0, 0, 0, 1],[0, 0, 1, 0],[0, 1, 0, 0]])

A = H_2*CNOT_2*H_2

print(A)

print(CNOT_1)

for key, val in Psi.items():
    print(key, 'becomes...\n', A*val)

