import FluidGridData;
import MarchingCubesData;

struct Params
{
    RWStructuredBuffer<float> vertexBuffer;
    RWStructuredBuffer<float> normalBuffer;
    RWStructuredBuffer<uint> indexBuffer;
    RWStructuredBuffer<uint> indicesCount;
    FluidGridData<float> levelSet;
};
ParameterBlock<Params> params;

const static float isoLevel = 0.0f;
const static float3 cellSize = 1.0f / gridSize;

float3 vertexInterp(float isoLevel, float3 p1, float3 p2, float valp1, float valp2)
{
    if (abs(isoLevel - valp1) < 0.00001)
        return p1;
    if (abs(isoLevel - valp2) < 0.00001)
        return p2;
    if (abs(valp1 - valp2) < 0.00001)
        return p1;
    float mu = (isoLevel - valp1) / (valp2 - valp1);
    return p1 + mu * (p2 - p1);
}

float3 gridPos(uint x, uint y, uint z)
{
    return float3(x, y, z) * cellSize;
}

// TODO: parameterize 
const static uint32_t LOCAL_THREADGROUP_SIZE = 16 * 8 * 1; // Must match numthreads
const static uint32_t MAX_TRIS_PER_THREAD = 5; // Max triangles per cube configuration
const static uint32_t MAX_VERTS_PER_GROUP = LOCAL_THREADGROUP_SIZE * MAX_TRIS_PER_THREAD * 3;

groupshared float3 localVertices[MAX_VERTS_PER_GROUP];
groupshared float3 localNormals[MAX_VERTS_PER_GROUP];
groupshared uint localTriCounts[LOCAL_THREADGROUP_SIZE]; // per-thread triangle counts
groupshared uint localOffsets[LOCAL_THREADGROUP_SIZE];   // per-thread exclusive offsets (in vertices)
groupshared uint localBaseOffset; // global base offset for this workgroup

// Hillis-Steele inclusive scan on localTriCounts, result in localOffsets
// Then convert to exclusive scan and allocate global space
void prefixSumAndAllocate(uint groupIndex)
{
    localOffsets[groupIndex] = localTriCounts[groupIndex];
    GroupMemoryBarrierWithGroupSync();

    for (uint offset = 1; offset < LOCAL_THREADGROUP_SIZE; offset *= 2)
    {
        uint temp = 0;
        if (groupIndex >= offset)
            temp = localOffsets[groupIndex - offset];
        GroupMemoryBarrierWithGroupSync();
        localOffsets[groupIndex] += temp;
        GroupMemoryBarrierWithGroupSync();
    }

    // localOffsets is now an inclusive prefix sum
    // Last thread has the total — allocate global space
    if (groupIndex == LOCAL_THREADGROUP_SIZE - 1)
    {
        uint totalVerts = localOffsets[groupIndex] * 3;
        InterlockedAdd(params.indicesCount[0], totalVerts, localBaseOffset);
    }
    GroupMemoryBarrierWithGroupSync();

    // Convert inclusive to exclusive: shift right, first element = 0
    uint inclusive = localOffsets[groupIndex];
    GroupMemoryBarrierWithGroupSync();
    localOffsets[groupIndex] = (groupIndex == 0) ? 0 : localOffsets[groupIndex - 1];
    GroupMemoryBarrierWithGroupSync();
}

[shader("compute")]
[numthreads(16, 8, 1)]
void marchingCubes(uint3 dispatchThreadID : SV_DispatchThreadID, uint groupIndex : SV_GroupIndex)
{
    uint x = dispatchThreadID.x + 1;
    uint y = dispatchThreadID.y + 1;
    uint z = dispatchThreadID.z + 1;
    FluidGridData<float> levelSet = params.levelSet;

    // Don't early-return — all threads must participate in barriers.
    // Use a flag to skip work for out-of-bounds threads.
    bool valid = (x < gridSize.x - 2 && y < gridSize.y - 2 && z < gridSize.z - 2);

    float val[8];
    uint cubeIndex = 0;
    uint triCount = 0;

    if (valid)
    {
        val[0] = levelSet[x, y, z];
        val[1] = levelSet[x + 1, y, z];
        val[2] = levelSet[x + 1, y + 1, z];
        val[3] = levelSet[x, y + 1, z];
        val[4] = levelSet[x, y, z + 1];
        val[5] = levelSet[x + 1, y, z + 1];
        val[6] = levelSet[x + 1, y + 1, z + 1];
        val[7] = levelSet[x, y + 1, z + 1];

        if (val[0] < isoLevel) cubeIndex |= 1;
        if (val[1] < isoLevel) cubeIndex |= 2;
        if (val[2] < isoLevel) cubeIndex |= 4;
        if (val[3] < isoLevel) cubeIndex |= 8;
        if (val[4] < isoLevel) cubeIndex |= 16;
        if (val[5] < isoLevel) cubeIndex |= 32;
        if (val[6] < isoLevel) cubeIndex |= 64;
        if (val[7] < isoLevel) cubeIndex |= 128;

        for (int i = 0; triTable[cubeIndex][i] != -1; i += 3)
            triCount++;
    }

    // Prefix sum to compute per-thread offsets and allocate global space
    localTriCounts[groupIndex] = triCount;
    GroupMemoryBarrierWithGroupSync();
    prefixSumAndAllocate(groupIndex);

    // localOffsets[groupIndex] = exclusive sum of triCounts before this thread
    // localBaseOffset = global vertex offset for this workgroup

    if (valid && triCount > 0)
    {
        float3 pos[8];
        pos[0] = gridPos(x, y, z);
        pos[1] = gridPos(x + 1, y, z);
        pos[2] = gridPos(x + 1, y + 1, z);
        pos[3] = gridPos(x, y + 1, z);
        pos[4] = gridPos(x, y, z + 1);
        pos[5] = gridPos(x + 1, y, z + 1);
        pos[6] = gridPos(x + 1, y + 1, z + 1);
        pos[7] = gridPos(x, y + 1, z + 1);

        float3 vertList[12];
        if ((edgeTable[cubeIndex] & 1) != 0)
            vertList[0] = vertexInterp(isoLevel, pos[0], pos[1], val[0], val[1]);
        if ((edgeTable[cubeIndex] & 2) != 0)
            vertList[1] = vertexInterp(isoLevel, pos[1], pos[2], val[1], val[2]);
        if ((edgeTable[cubeIndex] & 4) != 0)
            vertList[2] = vertexInterp(isoLevel, pos[2], pos[3], val[2], val[3]);
        if ((edgeTable[cubeIndex] & 8) != 0)
            vertList[3] = vertexInterp(isoLevel, pos[3], pos[0], val[3], val[0]);
        if ((edgeTable[cubeIndex] & 16) != 0)
            vertList[4] = vertexInterp(isoLevel, pos[4], pos[5], val[4], val[5]);
        if ((edgeTable[cubeIndex] & 32) != 0)
            vertList[5] = vertexInterp(isoLevel, pos[5], pos[6], val[5], val[6]);
        if ((edgeTable[cubeIndex] & 64) != 0)
            vertList[6] = vertexInterp(isoLevel, pos[6], pos[7], val[6], val[7]);
        if ((edgeTable[cubeIndex] & 128) != 0)
            vertList[7] = vertexInterp(isoLevel, pos[7], pos[4], val[7], val[4]);
        if ((edgeTable[cubeIndex] & 256) != 0)
            vertList[8] = vertexInterp(isoLevel, pos[0], pos[4], val[0], val[4]);
        if ((edgeTable[cubeIndex] & 512) != 0)
            vertList[9] = vertexInterp(isoLevel, pos[1], pos[5], val[1], val[5]);
        if ((edgeTable[cubeIndex] & 1024) != 0)
            vertList[10] = vertexInterp(isoLevel, pos[2], pos[6], val[2], val[6]);
        if ((edgeTable[cubeIndex] & 2048) != 0)
            vertList[11] = vertexInterp(isoLevel, pos[3], pos[7], val[3], val[7]);

        // Write triangles into shared memory
        uint localOff = localOffsets[groupIndex] * 3; // in vertices
        for (int i = 0; triTable[cubeIndex][i] != -1; i += 3)
        {
            float3 v0 = vertList[triTable[cubeIndex][i]];
            float3 v1 = vertList[triTable[cubeIndex][i + 1]];
            float3 v2 = vertList[triTable[cubeIndex][i + 2]];
            float3 n = normalize(cross(v1 - v0, v2 - v0));

            localVertices[localOff] = v0;
            localVertices[localOff + 1] = v1;
            localVertices[localOff + 2] = v2;
            localNormals[localOff] = n;
            localNormals[localOff + 1] = n;
            localNormals[localOff + 2] = n;
            localOff += 3;
        }
    }

    GroupMemoryBarrierWithGroupSync();

    // Copy from shared memory to global buffers
    // Each thread copies a chunk of the workgroup's output
    uint totalGroupVerts = localTriCounts[LOCAL_THREADGROUP_SIZE - 1] > 0
        ? (localOffsets[LOCAL_THREADGROUP_SIZE - 1] + localTriCounts[LOCAL_THREADGROUP_SIZE - 1]) * 3
        : localOffsets[LOCAL_THREADGROUP_SIZE - 1] * 3;

    for (uint i = groupIndex; i < totalGroupVerts; i += LOCAL_THREADGROUP_SIZE)
    {
        uint globalIdx = localBaseOffset + i;
        float3 v = localVertices[i];
        float3 n = localNormals[i];
        params.vertexBuffer[globalIdx * 3 + 0] = v.x;
        params.vertexBuffer[globalIdx * 3 + 1] = v.y;
        params.vertexBuffer[globalIdx * 3 + 2] = v.z;
        params.normalBuffer[globalIdx * 3 + 0] = n.x;
        params.normalBuffer[globalIdx * 3 + 1] = n.y;
        params.normalBuffer[globalIdx * 3 + 2] = n.z;
        params.indexBuffer[globalIdx] = globalIdx;
    }
}