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 i = dispatchThreadID.x + 1;
    int j = dispatchThreadID.y + 1;

    // X-faces: indices range over (gridSize.y, gridSize.z)
    if(i < gridParams.gridSize.y - 1 && j < gridParams.gridSize.z - 1)
    {
        grid[0, i, j] = b == 1 ? -grid[1, i, j] : grid[1, i, j];
        grid[gridParams.gridSize.x - 1, i, j] = b == 1 ? -grid[gridParams.gridSize.x - 2, i, j] : grid[gridParams.gridSize.x - 2, i, j];
    }

    // Y-faces: indices range over (gridSize.x, gridSize.z)
    if(i < gridParams.gridSize.x - 1 && j < gridParams.gridSize.z - 1)
    {
        grid[i, 0, j] = b == 2 ? -grid[i, 1, j] : grid[i, 1, j];
        grid[i, gridParams.gridSize.y - 1, j] = b == 2 ? -grid[i, gridParams.gridSize.y - 2, j] : grid[i, gridParams.gridSize.y - 2, j];
    }

    // Z-faces: indices range over (gridSize.x, gridSize.y)
    if(i < gridParams.gridSize.x - 1 && j < gridParams.gridSize.y - 1)
    {
        grid[i, j, 0] = b == 3 ? -grid[i, j, 1] : grid[i, j, 1];
        grid[i, j, gridParams.gridSize.z - 1] = b == 3 ? -grid[i, j, gridParams.gridSize.z - 2] : grid[i, j, gridParams.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;

    // X-axis edges: x ranges [1, gridSize.x-2]
    if(x < gridParams.gridSize.x - 1)
    {
        grid[x, 0, 0] = 0.5f * (grid[x, 1, 0] + grid[x, 0, 1]);
        grid[x, gridParams.gridSize.y - 1, 0] = 0.5f * (grid[x, gridParams.gridSize.y - 2, 0] + grid[x, gridParams.gridSize.y - 1, 1]);
        grid[x, 0, gridParams.gridSize.z - 1] = 0.5f * (grid[x, 1, gridParams.gridSize.z - 1] + grid[x, 0, gridParams.gridSize.z - 2]);
        grid[x, gridParams.gridSize.y - 1, gridParams.gridSize.z - 1] = 0.5f * (grid[x, gridParams.gridSize.y - 2, gridParams.gridSize.z - 1] + grid[x, gridParams.gridSize.y - 1, gridParams.gridSize.z - 2]);
    }

    // Y-axis edges: x ranges [1, gridSize.y-2]
    if(x < gridParams.gridSize.y - 1)
    {
        grid[0, x, 0] = 0.5f * (grid[1, x, 0] + grid[0, x, 1]);
        grid[0, x, gridParams.gridSize.z - 1] = 0.5f * (grid[1, x, gridParams.gridSize.z - 1] + grid[0, x, gridParams.gridSize.z - 2]);
        grid[gridParams.gridSize.x - 1, x, 0] = 0.5f * (grid[gridParams.gridSize.x - 2, x, 0] + grid[gridParams.gridSize.x - 1, x, 1]);
        grid[gridParams.gridSize.x - 1, x, gridParams.gridSize.z - 1] = 0.5f * (grid[gridParams.gridSize.x - 2, x, gridParams.gridSize.z - 1] + grid[gridParams.gridSize.x - 1, x, gridParams.gridSize.z - 2]);
    }

    // Z-axis edges: x ranges [1, gridSize.z-2]
    if(x < gridParams.gridSize.z - 1)
    {
        grid[0, 0, x] = 0.5f * (grid[1, 0, x] + grid[0, 1, x]);
        grid[0, gridParams.gridSize.y - 1, x] = 0.5f * (grid[1, gridParams.gridSize.y - 1, x] + grid[0, gridParams.gridSize.y - 2, x]);
        grid[gridParams.gridSize.x - 1, 0, x] = 0.5f * (grid[gridParams.gridSize.x - 2, 0, x] + grid[gridParams.gridSize.x - 1, 1, x]);
        grid[gridParams.gridSize.x - 1, gridParams.gridSize.y - 1, x] = 0.5f * (grid[gridParams.gridSize.x - 2, gridParams.gridSize.y - 1, x] + grid[gridParams.gridSize.x - 1, gridParams.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 : (gridParams.gridSize.x - 1);
    uint y = ((threadIdx & 2) == 0) ? 0 : (gridParams.gridSize.y - 1);
    uint z = ((threadIdx & 4) == 0) ? 0 : (gridParams.gridSize.z - 1);

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

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