//N-Body simulation
//Algorithm based on: http://http.developer.nvidia.com/GPUGems3/gpugems3_ch31.html

#define BLOCK_SIZE 128

cbuffer data                        : register(b0)
{
    uint num_particles;
    uint num_blocks;
    float deltaTime;
    float damping;
};

struct PosVel
{
    float4 pos; //xyz = pos, w = weight
    float4 vel; //only xyz
};

StructuredBuffer<PosVel>    oldPar  : register(t0);
RWStructuredBuffer<PosVel>  newPar  : register(u0);

static float g_FG = 6.67300e-11f;
static float g_softeningFactorSq =  0.0012500000*0.0012500000;
static float g_fParticleMass = 10000.0f * 10000.0f * 10000.0f;

groupshared float4 sh_Positions[BLOCK_SIZE];

void body_body_interaction(inout float3 ai, float4 bi, float4 bj)
{
    float3 r = bj.xyz - bi.xyz;

    float distSqr = dot(r, r);
    distSqr += g_softeningFactorSq;

    float distInvCube = 1.0f / sqrt(distSqr * distSqr * distSqr);

    //ai += g_FG * bj.w * distInvCube * r;
    ai += g_FG *g_fParticleMass * distInvCube * r;
}

[numthreads(BLOCK_SIZE,1,1)]
void CS(uint3 DTid : SV_DispatchThreadID, uint indexGroup : SV_GroupIndex, uint3 Gid : SV_GroupID)
{
    float3 accel = float3(0.0f, 0.0f, 0.0f);

    PosVel myParticle = oldPar[DTid.x];

    [loop]
    for(uint block=0; block< num_blocks; ++block)
    {
        //Fetch positions to shared cache
        sh_Positions[indexGroup] = oldPar[block * BLOCK_SIZE + indexGroup].pos;
        GroupMemoryBarrierWithGroupSync();

        [unroll]
        for(uint i = 0; i<BLOCK_SIZE; i+=8)
        {
            body_body_interaction(accel, myParticle.pos, sh_Positions[i]);
            body_body_interaction(accel, myParticle.pos, sh_Positions[i+1]);
            body_body_interaction(accel, myParticle.pos, sh_Positions[i+2]);
            body_body_interaction(accel, myParticle.pos, sh_Positions[i+3]);
            body_body_interaction(accel, myParticle.pos, sh_Positions[i+4]);
            body_body_interaction(accel, myParticle.pos, sh_Positions[i+5]);
            body_body_interaction(accel, myParticle.pos, sh_Positions[i+6]);
            body_body_interaction(accel, myParticle.pos, sh_Positions[i+7]);
        }

        GroupMemoryBarrierWithGroupSync();
    }

    //CS returns 0 for memory access out of range (write out of memory are NOPs)
    //So we calculated gravity for many particles from position (0,0,0), because num_particles != spawned_threads in certain scenarios
    //We need to subtract it now
    uint particleDifference = (num_blocks * BLOCK_SIZE) - num_particles;
    float3 a = float3(0.0f, 0.0f, 0.0f);

    body_body_interaction(a, myParticle.pos, float4(0.0f, 0.0f, 0.0f, 0.0f));
    accel -= particleDifference * a;

    //Calculate new position and velocity
    myParticle.vel.xyz += accel * deltaTime;
    myParticle.pos.xyz += myParticle.vel.xyz * deltaTime;

    //Save to global memory
    if(DTid.x < num_particles)
    {
        newPar[DTid.x].pos = myParticle.pos;
        newPar[DTid.x].vel = myParticle.vel;
    }
}