import numpy as np
import math

def Sx(j,k,d):
    # d = qudit dimension
    # j,k = identification of the Gell-Mann matrix
    # th = rotation angle
    j = j-1; k = k-1
    sx = np.zeros((d,d), dtype=complex)
    sx[j,k] = 1
    sx[k,j] = sx[j,k]
    return sx

def Sy(j,k,d):
    # d = qudit dimension
    # j,k = identification of the Gell-Mann matrix
    # th = rotation angle
    j = j-1; k = k-1
    sy = np.zeros((d,d), dtype=complex)
    sy[j,k] = -1j
    sy[k,j] = -sy[j,k]
    return sy

def Sz(j,d):
    # d = qudit dimension
    # j = identification of the Gell-Mann matrix
    # th = rotation angle
    sz = np.zeros((d,d), dtype=complex)
    for k in range(0,j+1):
        sz[k,k] = math.sqrt(2/(j*(j+1)))
        if k == j: # corrige para iniciar k em 0
            sz[k,k] = sz[k,k]*(-j)
    return sz

d = 2; j = 1; k = 2; Sx(j,k,d), Sy(j,k,d), Sz(j,d)

d = 3; Sx(1,2,d), Sx(1,3,d), Sx(2,3,d)

d = 3; Sy(1,2,d), Sy(1,3,d), Sy(2,3,d)

d = 3; Sz(1,d), Sz(2,d)

1/math.sqrt(3), -2/math.sqrt(3)

def Rx(j,k,d,th):
    # d = qudit dimension
    # j,k = identification of the Gell-Mann matrix
    # th = rotation angle
    j = j-1; k = k-1
    rx = np.eye(d, dtype=complex)
    rx[j,j] = math.cos(th/2)
    rx[k,k] = rx[j,j]
    rx[j,k] = -1j*math.sin(th/2)
    rx[k,j] = rx[j,k]
    return rx

def Ry(j,k,d,th):
    # d = qudit dimension
    # j,k = identification of the Gell-Mann matrix
    # th = rotation angle
    j = j-1; k = k-1
    ry = np.eye(d, dtype=complex)
    ry[j,j] = math.cos(th/2)
    ry[k,k] = ry[j,j]
    ry[j,k] = -math.sin(th/2)
    ry[k,j] = -ry[j,k]
    return ry

def Rz(j,d,th):
    # d = qudit dimension
    # j = identification of the Gell-Mann matrix
    # th = rotation angle
    rz = np.eye(d, dtype=complex)
    for k in range(0,j+1):
        angle = (th/2)*math.sqrt(2/(j*(j+1)))
        if k == j: # corrige para iniciar k em 0
            angle = angle*(-j)
        rz[k,k] = math.cos(angle) - 1j*math.sin(angle)
    return rz

d = 2; j = 1; k = 2; th = 0; Rx(j,k,d,th), Ry(j,k,d,th), Rz(j,d,th)

d = 3; j = 1; k = 2; th = 0; Rx(j,k,d,th), Ry(j,k,d,th), Rz(j,d,th)

