import numpy as np

from tqdm import tqdm

def Calculate_Index(i: int, j: int, n_cols: int) -> int:

    return n_cols * i + j * (i % 2 == 0) + (n_cols - 1 - j) * (i % 2 == 1)

def Generate_Adj_Matrix(n_cols: int) -> np.ndarray:

    row_array = np.zeros(8 * 512 * n_cols, dtype = np.uint32)

    col_array = np.zeros(8 * 512 * n_cols, dtype = np.uint32)

    for i in range(0, 512):

        index = 8 * i * n_cols if i % 2 == 0 else 8 * (i + 1) * n_cols - 1

        for j in range(0, n_cols):

            mask = list(range(index, index + 8)) if i % 2 == 0 else list(range(index, index - 8, -1))

            mask.sort()

            row_array[mask] = Calculate_Index(i, j, n_cols)

            i_m = 511 if i == 0 else i - 1

            i_p = 0 if i == 511 else i + 1

            j_m = n_cols - 1 if j == 0 else j - 1

            j_p = 0 if j == n_cols - 1 else j + 1

            con_indices = np.array([Calculate_Index(i, j_m, n_cols),

                                    Calculate_Index(i_m, j_m, n_cols),

                                    Calculate_Index(i_m, j, n_cols),

                                    Calculate_Index(i_m, j_p, n_cols),

                                    Calculate_Index(i, j_p, n_cols),

                                    Calculate_Index(i_p, j_p, n_cols),

                                    Calculate_Index(i_p, j, n_cols),

                                    Calculate_Index(i_p, j_m, n_cols)],

                                    dtype = np.uint32)

            con_indices.sort()

            col_array[mask] = con_indices

            index = index + 8 if i % 2 == 0 else index - 8

    return np.array([row_array, col_array], dtype = np.uint32)

def Generate_All_Adj_Matrices(n_cols_arr: np.ndarray) -> bool:

    try:

        for n_cols in tqdm(n_cols_arr):

            adj_mat = Generate_Adj_Matrix(n_cols = n_cols)

            np.save(file = "C:\\Users\\leotu\\OneDrive\\Documents\\ImageCHD_dataset\\ADJACENCY\\Adj_Mat_" \

                    + str(n_cols) + ".npy", arr = adj_mat)

        return True

    except Exception as e:

        print(e)

        return False