wallpaper/prim_wallpaper/pickers/diagonal.py

from math import sqrt


def get_row_num(i):
    return int((sqrt(8 * i + 1) - 1) / 2)


def get_last_whole_i(row):
    return (row * (row + 1)) // 2


class DiagonalCounter:
    def __init__(self, table_size):
        self.table_size = table_size

    def diagonal(self, i):
        row = get_row_num(i)
        last_whole_i = get_last_whole_i(row)
        if row < self.table_size:
            y = i - last_whole_i
            x = row - y
            return x, y
        else:
            x, y = self.diagonal(self.table_size ** 2 - (1 + i))
            return self.table_size - x - 1, self.table_size - y - 1

    def diagonal_reversed(self, i):
        y, x = self.diagonal(i)
        return x, y

    def diagonal_then_reversed(self, i):
        row = get_row_num(i)

        if row > self.table_size - 1:
            row = get_row_num(self.table_size ** 2 - (i + 1))
        if row % 2 != 0:
            return self.diagonal(i)
        else:
            return self.diagonal_reversed(i)


if __name__ == "__main__":
    import csv

    d = DiagonalCounter(6)
    for func in (d.diagonal, d.diagonal_reversed, d.diagonal_then_reversed):
        output = list(
            list(-1 for i in range(d.table_size)) for j in range(d.table_size)
        )
        indices = map(func, range(d.table_size ** 2))
        for i, index in enumerate(indices):
            x, y = index
            output[y][x] = i

        with open(func.__name__ + ".csv", "w", newline="") as file:
            writer = csv.writer(file)
            writer.writerows(output)