%matplotlib inline

import matplotlib
import numpy as np
import matplotlib.pyplot as plt
import sympy
from sympy import *
import itertools
from IPython.display import display
init_printing() #allows nice math displays

class QuantumMatrix():
    def __init__(self, m, coeff = 1):
        self.matrix = sympy.Matrix(m)
        self.coefficient = coeff
        self.canonize()
        
    def canonize(self):
        a = next((x for x in self.matrix if x != 0), None)
        if a is not None: #zero vector
            for i,j in itertools.product([0,1], [0,1]):
                self.matrix[i,j] = sympy.simplify(self.matrix[i,j] / a)
        self.coefficient = sympy.simplify(self.coefficient * a) 
            
    def __str__(self):
        coeff_string = ""
        if self.coefficient != 1:
            coeff_string = "{} * ".format(self.coefficient)
        return "{}[[{}, {}], [{}, {}]]".format(coeff_string,self.matrix[0], self.matrix[1], self.matrix[2], self.matrix[3])
    
    def __mul__(self, rhs):
        return QuantumMatrix(self.matrix * rhs.matrix, self.coefficient * rhs.coefficient)
    
    def __add__(self, rhs):
        temp_rhs_matrix = sympy.Matrix([[1,0],[0,1]])
        for i,j in itertools.product([0,1], [0,1]):
                 temp_rhs_matrix[i,j] = sympy.simplify((rhs.matrix[i,j] * self.coefficient) / rhs.coefficient)
        return QuantumMatrix(self.matrix + temp_rhs_matrix, self.coefficient * 1/sympy.sqrt(2))
    
    def __sub__(self, rhs):
        return self + QuantumMatrix(rhs.matrix, rhs.coefficient * -1)
    
    def __eq__(self, rhs):
        return (self.matrix == rhs.matrix and self.coefficient == rhs.coefficient)
    
    def equiv(self, rhs):
        return (self.matrix == rhs.matrix)
    
    def __iter__(self):
        for x in self.matrix:
             yield x

X = QuantumMatrix( [[0,1],  [1,0]])
Y = QuantumMatrix( [[0,-sympy.I], [sympy.I,0]])
Z = QuantumMatrix( [[1,0],  [0,-1]])
I = QuantumMatrix( [[1,0],  [0,1]])

print ("The Pauli matrices (when neglecting the global phase):")
print ("Identity: I=")
display(sympy.Matrix(I.matrix))
print ("X=")
display(sympy.Matrix(X.matrix))
print ("Y=")
display(sympy.Matrix(Y.matrix))
print ("Z=")
display(sympy.Matrix(Z.matrix))

print ("The Pauli gates X,Y,Z are of order 2:")
A=X*X
print ("X*X=")
display(sympy.Matrix(A.matrix))
A=Y*Y
print ("Y*Y=")
display(sympy.Matrix(A.matrix))
A=Z*Z
print ("Z*Z=")
display(sympy.Matrix(A.matrix))

print ("For example, one can verify that X*Y=Y*X=Z:")
A=X*Y
B=Y*X
print ("X*Y = Z =")
display(sympy.Matrix(A.matrix))
print ("Y*X = Z =")
display(sympy.Matrix(B.matrix))
print ("Similarly, X*Z=Z*X=Y, Y*Z=Z*Y=X")

H = QuantumMatrix(  [[1,1],  [1,-1]], 1/sympy.sqrt(2))
print ("Hadamard gate: H=")
display(sympy.Matrix(H.matrix))

print ("H is of order 2:")
A=H*H
print ("H*H =")
display(sympy.Matrix(A.matrix))

print ("H operates on the Pauli gates by conjugation as a reflection: X<-->Z:")
A=H*X*H
print("H*X*H = Z =")
display(sympy.Matrix(A.matrix))
A=H*Z*H
print("H*Z*H = X =")
display(sympy.Matrix(A.matrix))
A=H*Y*H
print("H*Y*H = Y =")
display(sympy.Matrix(A.matrix))

S = QuantumMatrix(  [[1,0],  [0,sympy.I]])
print ("Phase gate: S=")
display(sympy.Matrix(S.matrix))

print ("Z=S*S, thus S has order 4, and the inverse of S is S*Z:")
A=S*S
print ("S*S = Z =")
display(sympy.Matrix(A.matrix))
print ("The inverse of S: Sdg = S*Z =")
Sdg = S*Z
display(sympy.Matrix(Sdg.matrix))
print ("Indeed, S*Sdg =")
A = S*Sdg
display(sympy.Matrix(A.matrix))

print ("S operates on the Pauli gates by conjugation as a reflection: Y<-->Z:")
A = S*X*Sdg # =Y
print ("S*X*Sdg = Y =")
display(sympy.Matrix(A.matrix))
print ("S*Y*Sdg = X =")
A = S*Y*Sdg # =X
display(sympy.Matrix(A.matrix))
print ("S*Z*Sdg = Z =")
A = S*Z*Sdg # =Z
display(sympy.Matrix(A.matrix))

print ("We therefore get the commutator relations:")
print ("S*X = Y*S and Sdg*Y = X*Sdg")
print ("S*Y = X*S and Sdg*X = Y*Sdg")
print ("S*Z = Z*S = Sdg and Sdg*Z = Z*Sdg = S")

V = H*S*H*S
print ("Let: V = H*S*H*S =")
display(sympy.Matrix(V.matrix))

print ("V is of order 3:")
A=V*V*V
print ("V*V*V =")
display(sympy.Matrix(A.matrix))

print ("Hence, V = Sdg*H =")
A=Sdg*H
display(sympy.Matrix(A.matrix))

print ("The inverse of V is: W = S*Z*H*S*Z*H = H*S = ")
W=H*S
display(sympy.Matrix(W.matrix))

print ("Indeed, W = V*V =")
A=V*V
display(sympy.Matrix(A.matrix))

print ("V operates on the Pauli gates by conjugation as a rotation: Z-->X-->Y-->Z:")
A=W*X*V
print("W*X*V = Y =")
display(sympy.Matrix(A.matrix))
A=W*Y*V
print("W*Y*V = Z =")
display(sympy.Matrix(A.matrix))
A=W*Z*V
print("W*Z*V = X = ")
display(sympy.Matrix(A.matrix))

print ("I=")
display(sympy.Matrix(I.matrix))
print ("X=")
display(sympy.Matrix(X.matrix))
print ("Y=")
display(sympy.Matrix(Y.matrix))
print ("Z=")
display(sympy.Matrix(Z.matrix))

print ("V=")
display(sympy.Matrix(V.matrix))
A=V*X
print ("V*X=")
display(sympy.Matrix(A.matrix))
A=V*Y
print ("V*Y=")
display(sympy.Matrix(A.matrix))
A=V*Z
print ("V*Z=")
display(sympy.Matrix(A.matrix))

print ("W=")
display(sympy.Matrix(W.matrix))
A=W*X
print ("W*X=")
display(sympy.Matrix(A.matrix))
A=W*Y
print ("W*Y=")
display(sympy.Matrix(A.matrix))
A=W*Z
print ("W*Z=")
display(sympy.Matrix(A.matrix))

print ("H=")
display(sympy.Matrix(H.matrix))
A=H*X
print ("H*X=")
display(sympy.Matrix(A.matrix))
A=H*Y
print ("H*Y=")
display(sympy.Matrix(A.matrix))
A=H*Z
print ("H*Z=")
display(sympy.Matrix(A.matrix))

print ("H*V=")
A=H*V
display(sympy.Matrix(A.matrix))
A=H*V*X
print ("H*V*X=")
display(sympy.Matrix(A.matrix))
A=H*V*Y
print ("H*V*Y=")
display(sympy.Matrix(A.matrix))
A=H*V*Z
print ("H*V*Z=")
display(sympy.Matrix(A.matrix))

print ("H*W=")
A=H*W
display(sympy.Matrix(A.matrix))
A=H*W*X
print ("H*W*X=")
display(sympy.Matrix(A.matrix))
A=H*W*Y
print ("H*W*Y=")
display(sympy.Matrix(A.matrix))
A=H*W*Z
print ("H*W*Z=")
display(sympy.Matrix(A.matrix))

print ("H*V = H*H*S*H*S = S*H*S =")
A=S*H*S #H*V
display(sympy.Matrix(A.matrix))
A=Sdg*H*S #H*V*Y
print ("H*V*Y = H*H*S*H*S*Y = S*H*S*Y = S*H*X*S = S*Z*H*S = Sdg*H*S =")
display(sympy.Matrix(A.matrix))
A=S*H*Sdg #H*V*Z
print ("H*V*Z = H*H*S*H*S*Z = S*H*Sdg =")
display(sympy.Matrix(A.matrix))

print ("H*W = H*H*S = S =")
A=S #H*W
display(sympy.Matrix(A.matrix))
A= S*X #H*W*X
print ("H*W*X = S*X = ")
display(sympy.Matrix(A.matrix))
A=S*Y #H*W*Y
print ("H*W*Y = S*Y =")
display(sympy.Matrix(A.matrix))
A=Sdg #H*W*Z
print ("H*W*Z = Sdg =")
display(sympy.Matrix(A.matrix))

Class1Size = 2*2*3*3*4*4
print ("The size of Class 1 is: ", Class1Size)

Class2Size = 2*2*3*3*3*3*4*4
print ("The size of Class 2 is: ", Class2Size)

Class3Size = 2*2*3*3*3*3*4*4
print ("The size of Class 3 is: ", Class3Size)

Class4Size = 2*2*3*3*4*4
print ("The size of Class 4 is: ", Class4Size)

TotalSize = Class1Size + Class2Size + Class3Size + Class4Size
print ("The size of the 2-qubit Clifford group is: ", TotalSize)