// CelerisLab – operators/macro.cuh
// Macroscopic quantity computation (rho, u) and equilibrium distribution (feq).
// Supports DDF-shifting (USE_DDF_SHIFTING=1) and standard mode (=0).
//
// D2Q9 paired ordering:
//   cx = {0, 1,-1, 0, 0, 1,-1, 1,-1}
//   cy = {0, 0, 0, 1,-1, 1,-1,-1, 1}
// ============================================================================

#ifndef CELERIS_OPERATORS_MACRO_CUH
#define CELERIS_OPERATORS_MACRO_CUH

#ifndef USE_DDF_SHIFTING
#define USE_DDF_SHIFTING 0
#endif

// ---------------------------------------------------------------------------
// compute_rho_u:  f[NQ] → (rho, ux, uy [, uz])
// ---------------------------------------------------------------------------
#if LATTICE_MODEL == LATTICE_D2Q9

__device__ __forceinline__ void compute_rho_u(const float* __restrict__ f,
                                              float& rho,
                                              float& ux, float& uy)
{
    // Density: sum of all populations
    rho = f[0] + f[1] + f[2] + f[3] + f[4] + f[5] + f[6] + f[7] + f[8];
#if USE_DDF_SHIFTING
    rho += 1.0f;   // DDF-shifting: stored f_tilde = f - w, sum(w) = 1
#endif

    // Momentum (works identically with or without DDF-shifting,
    // because sum(w_i * cx_i) = 0)
    float inv_rho = 1.0f / rho;
    //  cx:  0  1 -1  0  0  1 -1  1 -1
    ux = (f[1] - f[2] + f[5] - f[6] + f[7] - f[8]) * inv_rho;
    //  cy:  0  0  0  1 -1  1 -1 -1  1
    uy = (f[3] - f[4] + f[5] - f[6] - f[7] + f[8]) * inv_rho;
}

#elif LATTICE_MODEL == LATTICE_D3Q19

__device__ __forceinline__ void compute_rho_u(const float* __restrict__ f,
                                              float& rho,
                                              float& ux, float& uy, float& uz)
{
    rho = f[0];
    for (int i = 1; i < 19; i++) rho += f[i];
#if USE_DDF_SHIFTING
    rho += 1.0f;
#endif

    float inv_rho = 1.0f / rho;
    ux = (f[1]-f[2]  + f[7]-f[8]  + f[9]-f[10] + f[13]-f[14] + f[15]-f[16]) * inv_rho;
    uy = (f[3]-f[4]  + f[7]-f[8]  + f[11]-f[12] + f[14]-f[13] + f[17]-f[18]) * inv_rho;
    uz = (f[5]-f[6]  + f[9]-f[10] + f[11]-f[12] + f[16]-f[15] + f[18]-f[17]) * inv_rho;
}

#endif  // LATTICE_MODEL (compute_rho_u)

// ---------------------------------------------------------------------------
// compute_feq:  (rho, u) → feq[NQ]
// ---------------------------------------------------------------------------
#if LATTICE_MODEL == LATTICE_D2Q9

__device__ __forceinline__ void compute_feq(float rho, float ux, float uy,
                                            float* __restrict__ feq)
{
#if USE_DDF_SHIFTING
    const float rhom1 = rho - 1.0f;          // perturbation density
#endif
    const float ux3 = 3.0f * ux;
    const float uy3 = 3.0f * uy;
    const float c3  = -1.5f * (ux * ux + uy * uy);  // = -|u|²/(2cs²)

    // Helper: feq_i = w_i * rho * (1 + 3 ci·u + 4.5 (ci·u)² - 1.5 u²)
    // With DDF-shifting: feq_tilde_i = feq_i - w_i
    //   = w_i * (rho * (1 + ...) - 1) = w_i * (rho*(...) + rho - 1)

    #define FEQ_CORE(cu)  (rho * fmaf(0.5f * 3.0f, (cu) * (cu), (cu) + c3))

#if USE_DDF_SHIFTING
    #define FEQ(w, cu)  ((w) * fmaf(rho, fmaf(0.5f, (cu)*(cu), (cu) + c3), rhom1))
#else
    #define FEQ(w, cu)  ((w) * rho * (1.0f + (cu) + 0.5f * (cu) * (cu) + c3))
#endif

    // i=0: rest
    feq[0] = FEQ(W0_VAL, 0.0f);

    // (1,2) ±x
    feq[1] = FEQ(WS_VAL,  ux3);
    feq[2] = FEQ(WS_VAL, -ux3);

    // (3,4) ±y
    feq[3] = FEQ(WS_VAL,  uy3);
    feq[4] = FEQ(WS_VAL, -uy3);

    // Pre-compute diagonal velocity combos
    const float u0 = ux3 + uy3;    // +x+y
    const float u1 = ux3 - uy3;    // +x-y

    // (5,6) ±(x+y)
    feq[5] = FEQ(WE_VAL,  u0);
    feq[6] = FEQ(WE_VAL, -u0);

    // (7,8) ±(x-y)
    feq[7] = FEQ(WE_VAL,  u1);
    feq[8] = FEQ(WE_VAL, -u1);

    #undef FEQ
    #undef FEQ_CORE
}

#elif LATTICE_MODEL == LATTICE_D3Q19

__device__ __forceinline__ void compute_feq(float rho, float ux, float uy, float uz,
                                            float* __restrict__ feq)
{
#if USE_DDF_SHIFTING
    const float rhom1 = rho - 1.0f;
#endif
    const float ux3 = 3.0f * ux;
    const float uy3 = 3.0f * uy;
    const float uz3 = 3.0f * uz;
    const float c3  = -1.5f * (ux*ux + uy*uy + uz*uz);

#if USE_DDF_SHIFTING
    #define FEQ(w, cu)  ((w) * fmaf(rho, fmaf(0.5f, (cu)*(cu), (cu) + c3), rhom1))
#else
    #define FEQ(w, cu)  ((w) * rho * (1.0f + (cu) + 0.5f * (cu) * (cu) + c3))
#endif

    feq[0]  = FEQ(W0_VAL, 0.0f);

    // Straight
    feq[1]  = FEQ(WS_VAL,  ux3);   feq[2]  = FEQ(WS_VAL, -ux3);
    feq[3]  = FEQ(WS_VAL,  uy3);   feq[4]  = FEQ(WS_VAL, -uy3);
    feq[5]  = FEQ(WS_VAL,  uz3);   feq[6]  = FEQ(WS_VAL, -uz3);

    // Diagonal
    const float u0 = ux3+uy3, u1 = ux3+uz3, u2 = uy3+uz3;
    const float u3 = ux3-uy3, u4 = ux3-uz3, u5 = uy3-uz3;

    feq[7]  = FEQ(WE_VAL,  u0);   feq[8]  = FEQ(WE_VAL, -u0);
    feq[9]  = FEQ(WE_VAL,  u1);   feq[10] = FEQ(WE_VAL, -u1);
    feq[11] = FEQ(WE_VAL,  u2);   feq[12] = FEQ(WE_VAL, -u2);
    feq[13] = FEQ(WE_VAL,  u3);   feq[14] = FEQ(WE_VAL, -u3);
    feq[15] = FEQ(WE_VAL,  u4);   feq[16] = FEQ(WE_VAL, -u4);
    feq[17] = FEQ(WE_VAL,  u5);   feq[18] = FEQ(WE_VAL, -u5);

    #undef FEQ
}

#endif  // LATTICE_MODEL (compute_feq)

// ---------------------------------------------------------------------------
// compute_pressure (diagnostic, from state equation p = cs² ρ)
// ---------------------------------------------------------------------------
__device__ __forceinline__ float compute_pressure(float rho) {
    return CS2 * rho;
}

__device__ __forceinline__ float compute_pressure_perturbation(float rho, float rho_ref) {
    return CS2 * (rho - rho_ref);
}

#endif // CELERIS_OPERATORS_MACRO_CUH
