from __future__ import print_function

import numpy as np


def commutation(mat1, mat2):
    return mat1 @ mat2 - mat2 @ mat1


# Define Sx, Sy, Sz
sx = np.array([[0, 1], [1, 0]]) / 2
sy = np.array([[0, -1j], [1j, 0]]) / 2
sz = np.array([[1, 0], [0, -1]]) / 2

print("Sx =\n", sx)
print("Sy =\n", sy)
print("Sz =\n", sz)

# Commutation relation
print("\nCommutation relations")
print("[Sx, Sy] =\n", commutation(sx, sy))  # = i Sz
print("[Sy, Sz] =\n", commutation(sy, sz))  # = i Sx
print("[Sz, Sx] =\n", commutation(sz, sx))  # = i Sy

# Diagonalize Sx
eigval, eigvec = np.linalg.eigh(sx)
print("\nEigenvalues =\n", eigval)
print("\nEigenvectors (column vectors)=\n", eigvec)

# U = <z|x>
#   |x> : eigenvectors of the Sx operator
#   |z> : eigenvectors of the Sz operator
# u = eigvec
u = eigvec @ np.diag([-1j, 1])
print(u)
print(u.conj().T @ u)

# Transform ops from |z> to |x>
sx_2 = u.conj().T @ sx @ u
sy_2 = u.conj().T @ sy @ u
sz_2 = u.conj().T @ sz @ u
print("\nAfter basis transform")
print("Sx' =\n", sx_2.round(10))
print("Sy' =\n", sy_2.round(10))
print("Sz' =\n", sz_2.round(10))