// CelerisLab – operators/forcing_guo.cuh
// Guo forcing scheme  (Guo, Zheng & Shi, Phys. Rev. E 65, 2002).
//
// Velocity correction:   u* = u + F·dt / (2ρ)
// Forcing term:
//   Fin[i] = (1 - ω/2) · w_i · [ (c_i - u*)/cs² + (c_i · u*)/cs⁴ · c_i ] · F
//   simplified: Fin[i] = 9·w_i · [ (c_i·F)(c_i·u* + 1/3) - u*·F/3 ]
// ============================================================================

#ifndef CELERIS_OPERATORS_FORCING_GUO_CUH
#define CELERIS_OPERATORS_FORCING_GUO_CUH

// ---------------------------------------------------------------------------
// D2Q9 Guo forcing terms
// ---------------------------------------------------------------------------
#if NQ == 9

__device__ __forceinline__ void apply_guo_velocity_correction(
    float& ux, float& uy,
    float fx, float fy, float rho)
{
    float rho2 = 0.5f / rho;
    ux = fmaf(fx, rho2, ux);
    uy = fmaf(fy, rho2, uy);
}

__device__ __forceinline__ void compute_guo_forcing(
    float ux, float uy,           // velocity (already corrected: u*)
    float fx, float fy,           // force density
    float* __restrict__ Fin)
{
    const float uF = -0.33333334f * fmaf(ux, fx, uy * fy);   // -u*·F / 3

    Fin[0] = 9.0f * W0_VAL * uF;

    // Straight directions
    //   cx = {0, 1,-1, 0, 0, 1,-1, 1,-1}
    //   cy = {0, 0, 0, 1,-1, 1,-1,-1, 1}
    Fin[1] = 9.0f * WS_VAL * fmaf( fx,  ux + 0.33333334f, uF);   // +x
    Fin[2] = 9.0f * WS_VAL * fmaf(-fx, -ux + 0.33333334f, uF);   // -x
    Fin[3] = 9.0f * WS_VAL * fmaf( fy,  uy + 0.33333334f, uF);   // +y
    Fin[4] = 9.0f * WS_VAL * fmaf(-fy, -uy + 0.33333334f, uF);   // -y

    // Diagonal directions
    const float f5x = fx, f5y = fy;   // +x+y
    const float c5u = ux + uy;
    const float c5F = f5x + f5y;
    Fin[5] = 9.0f * WE_VAL * fmaf(c5F, c5u + 0.33333334f, uF);

    const float c6u = -ux - uy;        // -x-y
    const float c6F = -fx - fy;
    Fin[6] = 9.0f * WE_VAL * fmaf(c6F, c6u + 0.33333334f, uF);

    const float c7u = ux - uy;         // +x-y
    const float c7F = fx - fy;
    Fin[7] = 9.0f * WE_VAL * fmaf(c7F, c7u + 0.33333334f, uF);

    const float c8u = -ux + uy;        // -x+y
    const float c8F = -fx + fy;
    Fin[8] = 9.0f * WE_VAL * fmaf(c8F, c8u + 0.33333334f, uF);
}

#elif NQ == 19

__device__ __forceinline__ void apply_guo_velocity_correction(
    float& ux, float& uy, float& uz,
    float fx, float fy, float fz, float rho)
{
    float rho2 = 0.5f / rho;
    ux = fmaf(fx, rho2, ux);
    uy = fmaf(fy, rho2, uy);
    uz = fmaf(fz, rho2, uz);
}

__device__ __forceinline__ void compute_guo_forcing(
    float ux, float uy, float uz,
    float fx, float fy, float fz,
    float* __restrict__ Fin)
{
    const float uF = -0.33333334f * fmaf(ux, fx, fmaf(uy, fy, uz * fz));

    Fin[0] = 9.0f * W0_VAL * uF;

    // Generic loop for D3Q19 (compiler will fully unroll)
    for (int i = 1; i < 19; i++) {
        float ci_dot_F = d_cx[i]*fx + d_cy[i]*fy + d_cz[i]*fz;
        float ci_dot_u = d_cx[i]*ux + d_cy[i]*uy + d_cz[i]*uz;
        Fin[i] = 9.0f * d_w[i] * fmaf(ci_dot_F, ci_dot_u + 0.33333334f, uF);
    }
}

#endif // NQ

// ---------------------------------------------------------------------------
// Zero forcing helper (when no external force)
// ---------------------------------------------------------------------------
__device__ __forceinline__ void zero_forcing(float* __restrict__ Fin) {
    #pragma unroll
    for (int i = 0; i < NQ; i++) Fin[i] = 0.0f;
}

#endif // CELERIS_OPERATORS_FORCING_GUO_CUH
