Determine neighboring cells in a two-dimensional grid with periodic conditions

Suppose we have the following two-dimensional network, whose cell indexes we label with integers:

20  21  22  23  2415  16  17  18  1910  11  12  13  145   6   7   8   90   1   2   3   4

What I want is a function that receives as input a cell index (cell) and the number of cells along an axis (n=5 in this case), and returns an array with its 9 neighbors (including the cell itself), taking into account the periodicity of the global box.

I show you what I've tried, which 'almost' works:

def celdas_vecinas(cell,n):    return np.mod(cell + np.array([0, -n-1, -n, -n+1, -1, 1, n-1, n, n+1], dtype=np.int64), n**2)

Where I have entered np.mod to reflect the periodic conditions. The point is that this function behaves well only for some values.

>>> celdas_vecinas(1,5) array([ 1, 20, 21, 22,  0,  2,  5,  6,  7]) # right!>>> celdas_vecinas(21,5)array([21, 15, 16, 17, 20, 22,  0,  1,  2]) # right!

But if I enter the index of one of the cells in the corners, the following happens:

>>> celdas_vecinas(0,5)array([ 0, 19, 20, 21, 24,  1,  4,  5,  6]) # should be 9 instead of 19

Also fails for cell=5 for example.

Does anyone know how I could implement this function? When the cell index doesn't touch any border it's very easy to implement, but I don't know how to include the periodic effects, although I guess it must be something related to the np.mod function