import FluidGridData;

void linearSolve(int b, FluidGridData<float> grid, const FluidGridData<float> grid0, float a, float c, int iterations)
{
    float cRecip = 1.0f / c;
    for(int k = 0; k < iterations; k++)
    {
        for(int z = 1; z < gridSize.z - 1; z++)
        {
            for(int y = 1; y < gridSize.y - 1; y++)
            {
                for(int x = 1; x < gridSize.x - 1; x++)
                {
                    grid[x, y, z] = (grid0[x, y, z] + a * (grid[x + 1, y, z] + grid[x - 1, y, z] + grid[x, y + 1, z] + grid[x, y - 1, z] + grid[x, y, z + 1] + grid[x, y, z - 1])) * cRecip;
                }
            }
        }
        setBound(b, grid);
    }
}

void diffuse(int b, FluidGridData<float> grid, const FluidGridData<float> grid0, float diff, float dt, int iterations)
{
    float a = dt * diff * (gridSize.x - 2) * (gridSize.y - 2) * (gridSize.z - 2);
    linearSolve(b, grid, grid0, a, 1 + 6 * a, iterations);
}

void advect(int b, FluidGridData<float> density, FluidGridData<float> density0, FluidGridData<float> velocityX, FluidGridData<float> velocityY, FluidGridData<float> velocityZ, float dt)
{
    float dt0x = dt * (gridSize.x - 2);
    float dt0y = dt * (gridSize.y - 2);
    float dt0z = dt * (gridSize.z - 2);

    for(int z = 1; z < gridSize.z - 1; z++)
    {
        for(int y = 1; y < gridSize.y - 1; y++)
        {
            for(int x = 1; x < gridSize.x - 1; x++)
            {
                float3 pos = float3(x - dt0x * velocityX[x, y, z], y - dt0y * velocityY[x, y, z], z - dt0z * velocityZ[x, y, z]);
                pos.x = clamp(pos.x, 0.5f, gridSize.x - 1.5f);
                pos.y = clamp(pos.y, 0.5f, gridSize.y - 1.5f);
                pos.z = clamp(pos.z, 0.5f, gridSize.z - 1.5f);

                int3 i0 = int3(pos);
                int3 i1 = i0 + int3(1, 1, 1);

                float3 s1 = pos - i0;
                float3 s0 = 1 - s1;

                density[x, y, z] = 
                    s0.x * (s0.y * (s0.z * density0[i0.x, i0.y, i0.z] + s1.z * density0[i0.x, i0.y, i1.z]) + 
                    s1.y * (s0.z * density0[i0.x, i1.y, i0.z] + s1.z * density0[i0.x, i1.y, i1.z])) + 
                    s1.x * (s0.y * (s0.z * density0[i1.x, i0.y, i0.z] + s1.z * density0[i1.x, i0.y, i1.z]) + 
                    s1.y * (s0.z * density0[i1.x, i1.y, i0.z] + s1.z * density0[i1.x, i1.y, i1.z]));
            }
        }
    }
    setBound(b, density);
}

void project(FluidGridData<float> velocityX, FluidGridData<float> velocityY, FluidGridData<float> velocityZ, FluidGridData<float> pressure, FluidGridData<float> divergence, int iter) 
{
    for(int z = 1; z < gridSize.z - 1; z++)
    {
        for(int y = 1; y < gridSize.y - 1; y++)
        {
            for(int x = 1; x < gridSize.x - 1; x++)
            {
                divergence[x, y, z] = -0.5f * (velocityX[x + 1, y, z] - velocityX[x - 1, y, z] + velocityY[x, y + 1, z] - velocityY[x, y - 1, z] + velocityZ[x, y, z + 1] - velocityZ[x, y, z - 1]) / gridSize.x;
                pressure[x, y, z] = 0;
            }
        }
    }
    setBound(0, divergence);
    setBound(0, pressure);
    linearSolve(0, pressure, divergence, 1, 6, iter);
    for(int z = 1; z < gridSize.z - 1; z++)
    {
        for(int y = 1; y < gridSize.y - 1; y++)
        {
            for(int x = 1; x < gridSize.x - 1; x++)
            {
                velocityX[x, y, z] -= 0.5f * (pressure[x + 1, y, z] - pressure[x - 1, y, z]) * gridSize.x;
                velocityY[x, y, z] -= 0.5f * (pressure[x, y + 1, z] - pressure[x, y - 1, z]) * gridSize.y;
                velocityZ[x, y, z] -= 0.5f * (pressure[x, y, z + 1] - pressure[x, y, z - 1]) * gridSize.z;
            }
        }
    }
    setBound(1, velocityX);
    setBound(2, velocityY);
    setBound(3, velocityZ);
}

[shader("compute")]
void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
{
    diffuse(1, pComputeParams.velocityX0, pComputeParams.velocityX, pComputeParams.viscosity, pComputeParams.dt, 20);
    diffuse(2, pComputeParams.velocityY0, pComputeParams.velocityY, pComputeParams.viscosity, pComputeParams.dt, 20);
    diffuse(3, pComputeParams.velocityZ0, pComputeParams.velocityZ, pComputeParams.viscosity, pComputeParams.dt, 20);
    project(pComputeParams.velocityX0, pComputeParams.velocityY0, pComputeParams.velocityZ0, pComputeParams.scratch, pComputeParams.density, 20);
    advect(1, pComputeParams.velocityX, pComputeParams.velocityX0, pComputeParams.velocityX0, pComputeParams.velocityY0, pComputeParams.velocityZ0, pComputeParams.dt); 
    advect(2, pComputeParams.velocityY, pComputeParams.velocityY0, pComputeParams.velocityX0, pComputeParams.velocityY0, pComputeParams.velocityZ0, pComputeParams.dt);
    advect(3, pComputeParams.velocityZ, pComputeParams.velocityZ0, pComputeParams.velocityX0, pComputeParams.velocityY0, pComputeParams.velocityZ0, pComputeParams.dt);

    diffuse(0, pComputeParams.scratch, pComputeParams.density, pComputeParams.diffusion, pComputeParams.dt, 20);
    advect(0, pComputeParams.density, pComputeParams.scratch, pComputeParams.velocityX0, pComputeParams.velocityY0, pComputeParams.velocityZ0, pComputeParams.dt);
}