// CelerisLab – core/layout.cuh
// SoA memory layout: index_f, coordinates, and neighbor computation.
// Depends on: LATTICE_MODEL, NQ, NX, NY, NZ from config_grid.h; descriptors.cuh for d_cx/d_cy.
// ============================================================================

#ifndef CELERIS_CORE_LAYOUT_CUH
#define CELERIS_CORE_LAYOUT_CUH

// ---------------------------------------------------------------------------
// Total cell count (compile-time)
// ---------------------------------------------------------------------------
#define TOTAL_CELLS ((unsigned long)(NX) * (unsigned long)(NY) * (unsigned long)(NZ))

// ---------------------------------------------------------------------------
// SoA index:  fi[ direction * N + node ]   ("Q-major")
// ---------------------------------------------------------------------------
__device__ __forceinline__ unsigned long index_f(unsigned long n, unsigned int i) {
    return (unsigned long)i * TOTAL_CELLS + n;
}

// ---------------------------------------------------------------------------
// SoA index for DIM-component field:  u[ component * N + node ]
// ---------------------------------------------------------------------------
__device__ __forceinline__ unsigned long index_u(unsigned long n, unsigned int d) {
    return (unsigned long)d * TOTAL_CELLS + n;
}

// ---------------------------------------------------------------------------
// 2D linear index → (x, y)  and  (x, y) → linear
// ---------------------------------------------------------------------------
#if DIM == 2

__device__ __forceinline__ void coordinates(unsigned long n,
                                            unsigned int& x,
                                            unsigned int& y)
{
    x = (unsigned int)(n % NX);
    y = (unsigned int)(n / NX);
}

__device__ __forceinline__ unsigned long linear_index(unsigned int x,
                                                      unsigned int y)
{
    return (unsigned long)y * NX + x;
}

__device__ __forceinline__ void index_from_thread(unsigned int& x,
                                                  unsigned int& y,
                                                  unsigned long& k)
{
    x = threadIdx.x + NT * blockIdx.x;
    y = blockIdx.y;
    k = (unsigned long)y * NX + x;
}

// Global cell of window cell (x, y) in block blockIdx.z of a row-major,
// full-domain block layout domain_nx cells wide. NX/NY are the local strides.
__device__ __forceinline__ void block_global_xy(unsigned int x, unsigned int y,
                                                unsigned int halo, unsigned int domain_nx,
                                                unsigned int& gx, unsigned int& gy)
{
    const unsigned int bx = NX - 2 * halo, by = NY - 2 * halo;
    const unsigned int nbx = domain_nx / bx;
    gx = (blockIdx.z % nbx) * bx + x - halo;
    gy = (blockIdx.z / nbx) * by + y - halo;
}

// ---------------------------------------------------------------------------
// Neighbor list (periodic BCs via modular arithmetic)
// ---------------------------------------------------------------------------
__device__ inline void compute_neighbors(unsigned long n, unsigned long* j) {
    unsigned int x, y;
    coordinates(n, x, y);

    unsigned int xp = (x + 1u) % NX;
    unsigned int xm = (x + NX - 1u) % NX;
    unsigned int yp = (y + 1u) % NY;
    unsigned int ym = (y + NY - 1u) % NY;

    // i=0: self
    j[0] = n;

#if LATTICE_MODEL == LATTICE_D2Q9
    // paired: (1,2)±x  (3,4)±y  (5,6)±(x+y)  (7,8)±(x-y)
    j[1] = linear_index(xp, y );   // +x
    j[2] = linear_index(xm, y );   // -x
    j[3] = linear_index(x , yp);   // +y
    j[4] = linear_index(x , ym);   // -y
    j[5] = linear_index(xp, yp);   // +x +y
    j[6] = linear_index(xm, ym);   // -x -y
    j[7] = linear_index(xp, ym);   // +x -y
    j[8] = linear_index(xm, yp);   // -x +y
#else
    #error "2D compute_neighbors only supports LATTICE_D2Q9."
#endif
}

// Overload: skip coordinates() when x,y already known (saves ~40 cycles)
__device__ inline void compute_neighbors(unsigned int x, unsigned int y,
                                         unsigned long* j) {
    unsigned int xp = (x + 1u) % NX;
    unsigned int xm = (x + NX - 1u) % NX;
    unsigned int yp = (y + 1u) % NY;
    unsigned int ym = (y + NY - 1u) % NY;

    j[0] = linear_index(x, y);
#if LATTICE_MODEL == LATTICE_D2Q9
    j[1] = linear_index(xp, y );
    j[2] = linear_index(xm, y );
    j[3] = linear_index(x , yp);
    j[4] = linear_index(x , ym);
    j[5] = linear_index(xp, yp);
    j[6] = linear_index(xm, ym);
    j[7] = linear_index(xp, ym);
    j[8] = linear_index(xm, yp);
#else
    #error "2D compute_neighbors only supports LATTICE_D2Q9."
#endif
}

#elif DIM == 3

__device__ __forceinline__ void coordinates(unsigned long n,
                                            unsigned int& x,
                                            unsigned int& y,
                                            unsigned int& z)
{
    x = (unsigned int)(n % NX);
    y = (unsigned int)(n / NX % NY);
    z = (unsigned int)(n / ((unsigned long)NX * NY));
}

__device__ __forceinline__ unsigned long linear_index(unsigned int x,
                                                      unsigned int y,
                                                      unsigned int z)
{
    return (unsigned long)z * NY * NX + (unsigned long)y * NX + x;
}

__device__ __forceinline__ void index_from_thread(unsigned int& x,
                                                  unsigned int& y,
                                                  unsigned int& z,
                                                  unsigned long& k)
{
    x = threadIdx.x + NT * blockIdx.x;
    y = blockIdx.y;
    z = blockIdx.z;
    k = linear_index(x, y, z);
}

__device__ inline void compute_neighbors(unsigned long n, unsigned long* j) {
    unsigned int x, y, z;
    coordinates(n, x, y, z);

    unsigned int xp = (x + 1u) % NX;
    unsigned int xm = (x + NX - 1u) % NX;
    unsigned int yp = (y + 1u) % NY;
    unsigned int ym = (y + NY - 1u) % NY;
    unsigned int zp = (z + 1u) % NZ;
    unsigned int zm = (z + NZ - 1u) % NZ;

    j[0]  = n;

#if LATTICE_MODEL == LATTICE_D3Q19
    // ±x, ±y, ±z (straight)
    j[1]  = linear_index(xp, y,  z );   // +x
    j[2]  = linear_index(xm, y,  z );   // -x
    j[3]  = linear_index(x,  yp, z );   // +y
    j[4]  = linear_index(x,  ym, z );   // -y
    j[5]  = linear_index(x,  y,  zp);   // +z
    j[6]  = linear_index(x,  y,  zm);   // -z
    // diagonal
    j[7]  = linear_index(xp, yp, z );   // +x+y
    j[8]  = linear_index(xm, ym, z );   // -x-y
    j[9]  = linear_index(xp, y,  zp);   // +x+z
    j[10] = linear_index(xm, y,  zm);   // -x-z
    j[11] = linear_index(x,  yp, zp);   // +y+z
    j[12] = linear_index(x,  ym, zm);   // -y-z
    j[13] = linear_index(xp, ym, z );   // +x-y
    j[14] = linear_index(xm, yp, z );   // -x+y
    j[15] = linear_index(xp, y,  zm);   // +x-z
    j[16] = linear_index(xm, y,  zp);   // -x+z
    j[17] = linear_index(x,  yp, zm);   // +y-z
    j[18] = linear_index(x,  ym, zp);   // -y+z
#endif
}

#endif // DIM == 3

#endif // CELERIS_CORE_LAYOUT_CUH