import numpy as np
import matplotlib.pyplot as plt


def update_border(voxel_values, idx=None):

    voxel_values = voxel_values > 0

    if idx is None:
        x_m1 = voxel_values[1:, :, :]
        x_m1 = np.concatenate((x_m1, np.zeros((1, x_m1.shape[1], x_m1.shape[2]))), axis=0)

        x_p1 = voxel_values[:-1, :, :]
        x_p1 = np.concatenate((np.zeros((1, x_p1.shape[1], x_p1.shape[2])), x_p1), axis=0)

        y_m1 = voxel_values[:, 1:, :]
        y_m1 = np.concatenate((y_m1, np.zeros((y_m1.shape[0], 1, y_m1.shape[2]))), axis=1)

        y_p1 = voxel_values[:, :-1, :]
        y_p1 = np.concatenate((np.zeros((y_p1.shape[0], 1, y_p1.shape[2])), y_p1), axis=1)

        z_m1 = voxel_values[:, :, 1:]
        z_m1 = np.concatenate((z_m1, np.zeros((z_m1.shape[0], z_m1.shape[1], 1))), axis=2)

        z_p1 = voxel_values[:, :, :-1]
        z_p1 = np.concatenate((np.zeros((z_p1.shape[0], z_p1.shape[1], 1)), z_p1), axis=2)

        return np.logical_or.reduce((voxel_values != x_m1, voxel_values != x_p1,
                                     voxel_values != y_m1, voxel_values != y_p1,
                                     voxel_values != z_m1, voxel_values != z_p1))

    # TODO: update only concidered voxels (idx)


if __name__ == '__main__':
    voxel_values = np.array([[[np.sqrt(x**2 + y**2 + z**2) < 20 for z in np.arange(-10, 10, 1.0)] for y in np.arange(-10, 10, 1.0)] for x in np.arange(-10, 10, 1.0)])

    border = update_border(voxel_values)
    
    # Plot voxel grid that are the border
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(np.where(voxel_values)[0], np.where(voxel_values)[1], np.where(voxel_values)[2], c='r', marker='o')
    ax.scatter(np.where(border)[0], np.where(border)[1], np.where(border)[2], c='b', marker='o', s=1)
    plt.show()