import FluidGridData;

struct Params
{
    int b;
    // read-write
    FluidGridData<float> grid;
};
ParameterBlock<Params> params;

[shader("compute")]
[numthreads(32, 8, 1)]
void setBound(uint3 dispatchThreadID : SV_DispatchThreadID)
{
    int b = params.b;
    FluidGridData<float> grid = params.grid;
    int x = dispatchThreadID.x + 1;
    int y = dispatchThreadID.y + 1;
    if(x >= gridSize.x - 1 || y >= gridSize.y - 1) return;
    
    grid[0, x, y] = b == 1 ? grid[1, x, y] : -grid[1, x, y];
    grid[gridSize.x - 1, x, y] = b == 1 ? grid[gridSize.x - 2, x, y] : -grid[gridSize.x - 2, x, y];
    
    grid[x, 0, y] = b == 2 ? grid[x, 1, y] : -grid[x, 1, y];
    grid[x, gridSize.y - 1, y] = b == 2 ? grid[x, gridSize.y - 2, y] : -grid[x, gridSize.y - 2, y];

    grid[x, y, 0] = b == 3 ? grid[x, y, 1] : -grid[x, y, 1];
    grid[x, y, gridSize.z - 1] = b == 3 ? grid[x, y, gridSize.z - 2] : -grid[x, y, gridSize.z - 2];
}

[shader("compute")]
[numthreads(128, 1, 1)]
void setBoundEdges(uint3 dispatchThreadID : SV_DispatchThreadID)
{
    FluidGridData<float> grid = params.grid;
    int x = dispatchThreadID.x + 1;
    if(x >= gridSize.x - 1) return;

    grid[x, 0, 0] = 0.5f * (grid[x, 1, 0] + grid[x, 0, 1]);
    grid[x, gridSize.y - 1, 0] = 0.5f * (grid[x, gridSize.y - 2, 0] + grid[x, gridSize.y - 1, 1]);
    grid[x, 0, gridSize.z - 1] = 0.5f * (grid[x, 1, gridSize.z - 1] + grid[x, 0, gridSize.z - 2]);
    grid[x, gridSize.y - 1, gridSize.z - 1] = 0.5f * (grid[x, gridSize.y - 2, gridSize.z - 1] + grid[x, gridSize.y - 1, gridSize.z - 2]);

    grid[0, x, 0] = 0.5f * (grid[1, x, 0] + grid[0, x, 1]);
    grid[0, x, gridSize.z - 1] = 0.5f * (grid[1, x, gridSize.z - 1] + grid[0, x, gridSize.z - 2]);
    grid[gridSize.x - 1, x, 0] = 0.5f * (grid[gridSize.x - 2, x, 0] + grid[gridSize.x - 1, x, 1]);
    grid[gridSize.x - 1, x, gridSize.z - 1] = 0.5f * (grid[gridSize.x - 2, x, gridSize.z - 1] + grid[gridSize.x - 1, x, gridSize.z - 2]);
    
    grid[0, 0, x] = 0.5f * (grid[1, 0, x] + grid[0, 1, x]);
    grid[0, gridSize.y - 1, x] = 0.5f * (grid[1, gridSize.y - 1, x] + grid[0, gridSize.y - 2, x]);
    grid[gridSize.x - 1, 0, x] = 0.5f * (grid[gridSize.x - 2, 0, x] + grid[gridSize.x - 1, 1, x]);
    grid[gridSize.x - 1, gridSize.y - 1, x] = 0.5f * (grid[gridSize.x - 2, gridSize.y - 1, x] + grid[gridSize.x - 1, gridSize.y - 2, x]);
}

[shader("compute")]
[numthreads(8, 1, 1)]
void setBoundCorners(uint3 dispatchThreadID : SV_DispatchThreadID)
{
    FluidGridData<float> grid = params.grid;
    uint threadIdx = dispatchThreadID.x;
    
    uint x = ((threadIdx & 1) == 0) ? 0 : (gridSize.x - 1);
    uint y = ((threadIdx & 2) == 0) ? 0 : (gridSize.y - 1);
    uint z = ((threadIdx & 4) == 0) ? 0 : (gridSize.z - 1);

    uint nx = (x == 0) ? 1 : (gridSize.x - 2);
    uint ny = (y == 0) ? 1 : (gridSize.y - 2);
    uint nz = (z == 0) ? 1 : (gridSize.z - 2);

    grid[x, y, z] = 0.33f * (grid[nx, y, z] + grid[x, ny, z] + grid[x, y, nz]);
}