import numpy as np

def cnot_matrix_ab(d):
        cn = np.zeros((d**2,d**2),dtype=int)
        for p in range(0,d): # d^2 elementos não nulos
            for q in range(0,d):
                cn[p*d+(p+q)%d,p*d+q] = 1
        return cn

d=2; cn = cnot_matrix_ab(d); print(cn)

d=3; cn = cnot_matrix_ab(d); print(cn)

d=4; cn = cnot_matrix_ab(d); print(cn)

import numpy as np

def cnot_matrix(d, ctr):
    cn = np.zeros((d**2,d**2),dtype=int)
    if ctr == 'a':        
        for p in range(0,d):
            for q in range(0,d):
                cn[p*d+(p+q)%d,p*d+q] = 1
    elif ctr == 'b':
        for p in range(0,d):
            for q in range(0,d):
                cn[((q+p)%d)*d+q,p*d+q] = 1
    return cn

d=2; print(cnot_matrix(d, 'b')) # ok!

d=3; print(cnot_matrix(d, 'b'))

def den2dec(local,d):
    # local = list with the local computational base state values 
    # d = individual qudit dimension
    N = len(local)
    j = 0
    for k in range(0,N):
        j += local[k]*d**(N-1-k)
    return j # value of the global computational basis index

# example: 3 qutrits
d = 3; N = 3
local = [0,0,0]
for j0 in range(0,3):
    for j1 in range(0,3):
        for j2 in range(0,3):
            local = [j0,j1,j2]
            j = den2dec(local,d)
            print('j0 j1 j2 =',j0,j1,j2,', j =',j)

def dec2den(j,N,d):
    den = [0 for k in range(0,N)]
    jv = j
    for k in range(0,N):
        if jv >= d**(N-1-k):
            den[k] = jv//(d**(N-1-k))
            jv = jv - den[k]*d**(N-1-k)
    return den

N = 3; d = 3
for j in range(0,d**N):
    print('j=',j,', den=', dec2den(j,N,d))

def dec2den_opt(j,N,d): # optmized version
    den = [0 for k in range(0,N)]
    if j < d:
        den[N-1] = j; return den
    rem = [] # list for the remainders
    k = j
    while k >= d:
        rem.append(k%d)
        k = k//d
    rem.append(k%d)
    for j in range(0,len(rem)):
        den[N-1-j] = rem[j]
    return den

N = 3; d = 3
for j in range(0,d**N):
    print('j=',j,', den=', dec2den_opt(j,N,d))

def cnot_qudits(d,N,t,c):
    # d = qudit dimension; N = No. of qudits; t = target qudit (0,...,N-1); c = control qudit (0,...,N-1)
    D = d**N; cnot = np.zeros((D,D), dtype=int) # set all d^2N elements to 0
    for k in range (0,d**N):
        localr = dec2den_opt(k,N,d) # sequence os values for the local basis
        locall = localr.copy() # copy the right list to the left list
        locall[t] = (localr[t]+localr[c])%d
        globall = den2dec(locall,d)
        cnot[globall,k] = 1
    return cnot    

d=2; N=2; t=1; c=0; cnot = cnot_qudits(d,N,t,c); print(cnot) # ok!

d=2; N=2; t=0; c=1; cnot = cnot_qudits(d,N,t,c); print(cnot) # ok!

d=3; N=2; t=1; c=0; cnot = cnot_qudits(d,N,t,c); print(cnot) # ok!

d=2; N=3; t=2; c=0; cnot = cnot_qudits(d,N,t,c); print(cnot) # ok

from sympy import init_printing, Matrix
init_printing(use_unicode=True)
from sympy.physics.quantum import TensorProduct
def tp(x,y):
    return TensorProduct(x,y)
s1 = Matrix([[0,1],[1,0]]); s0 = Matrix([[1,0],[0,1]]); P0 = Matrix([[1,0],[0,0]]); P1 = Matrix([[0,0],[0,1]])
cnot03 = tp(tp(P0,s0),s0) + tp(tp(P1,s0),s1); cnot03                                                                                             

d=3; N=3; t=1; c=0; cnot = cnot_qudits(d,N,t,c); print(cnot) # ok

