import FluidGridData;

struct Params
{
    float dtau;
    // read-write: phi being reinitialized
    FluidGridData<float> phi;
    // read-only: snapshot of phi before reinitialization (for sign)
    FluidGridData<float> phi0;
};
ParameterBlock<Params> params;

// Smooth sign function to avoid discontinuity at zero
float smoothSign(float x, float dx)
{
    return x / sqrt(x * x + dx * dx);
}

[shader("compute")]
[numthreads(32, 8, 1)]
void reinitialize(uint3 dispatchThreadID : SV_DispatchThreadID)
{
    uint x = dispatchThreadID.x + 1;
    uint y = dispatchThreadID.y + 1;
    uint z = dispatchThreadID.z + 1;
    FluidGridData<float> phi = params.phi;
    FluidGridData<float> phi0 = params.phi0;
    if(x >= gridParams.gridSize.x - 1 || y >= gridParams.gridSize.y - 1 || z >= gridParams.gridSize.z - 1) return;

    float dtau = params.dtau;
    float dx = 1.0f;

    // Current value and original sign
    float p = phi[x, y, z];
    float s = smoothSign(phi0[x, y, z], dx);

    // Upwind finite differences for |grad phi|
    // Forward and backward differences
    float dxp = phi[x + 1, y, z] - p;
    float dxm = p - phi[x - 1, y, z];
    float dyp = phi[x, y + 1, z] - p;
    float dym = p - phi[x, y - 1, z];
    float dzp = phi[x, y, z + 1] - p;
    float dzm = p - phi[x, y, z - 1];

    // Godunov upwind scheme
    float gradPhi;
    if (s > 0)
    {
        float ax = max(max(dxm, 0.0f), -min(dxp, 0.0f));
        float ay = max(max(dym, 0.0f), -min(dyp, 0.0f));
        float az = max(max(dzm, 0.0f), -min(dzp, 0.0f));
        gradPhi = sqrt(ax * ax + ay * ay + az * az) / dx;
    }
    else
    {
        float ax = max(-min(dxm, 0.0f), max(dxp, 0.0f));
        float ay = max(-min(dym, 0.0f), max(dyp, 0.0f));
        float az = max(-min(dzm, 0.0f), max(dzp, 0.0f));
        gradPhi = sqrt(ax * ax + ay * ay + az * az) / dx;
    }

    // Update: push |grad phi| toward 1
    phi[x, y, z] = p - dtau * s * (gradPhi - 1.0f);
}