// CelerisLab – step/one_step_esopull.cu
// Main LBM step kernel using Esoteric-Pull single-buffer streaming.
//
// Workflow per node:
//   1. load_f_esopull: read DDF (streaming second half)
//   2. Compute macroscopic quantities (ρ, u)
//   3. Boundary conditions
//   4. Equilibrium + forcing + collision
//   5. store_f_esopull: write DDF (streaming first half)
//
// Compared to double-buffer:
//   - Uses HALF the memory (single fi array)
//   - Timestep counter t is required
//   - fi_alt pointer is unused (can be NULL)
// ============================================================================

#ifndef CELERIS_STEP_ONE_STEP_ESOPULL_CU
#define CELERIS_STEP_ONE_STEP_ESOPULL_CU

__global__ void StreamCollideEsoPull(
    fpxx*          __restrict__ fi,          // single DDF buffer (read/write)
    const uint8_t* __restrict__ flag,
    const int32_t* __restrict__ indx,
    const float*   __restrict__ delta,
    const float*   __restrict__ action,
    float*         __restrict__ obs,
    float*         __restrict__ rho_arr,
    float*         __restrict__ u_arr,
    const float*   __restrict__ force_field, // Euler force field [DIM*N], or NULL
    unsigned long                t)          // current timestep
{
    // ----- Thread → node mapping -----
#if DIM == 2
    unsigned int x, y;
    unsigned long k;
    index_from_thread(x, y, k);
    if (x >= (unsigned int)NX || y >= (unsigned int)NY) return;
#elif DIM == 3
    unsigned int x, y, z;
    unsigned long k;
    index_from_thread(x, y, z, k);
    if (x >= (unsigned int)NX || y >= (unsigned int)NY || z >= (unsigned int)NZ) return;
#endif

    uint8_t fl = flag[k];

    // Skip pure solid / gas
    if ((fl & LEGACY_SOLID) && !(fl & LEGACY_INTERFACE) && !(fl & LEGACY_SENSOR)) {
        // For inlet/outlet solid nodes, we still process below
        if (x != 0 && x != (unsigned int)(NX - 1)) return;
    }

    // ----- Neighbor indices -----
    unsigned long j[NQ];
    compute_neighbors(k, j);

    // ----- Esoteric-Pull: load DDF (streaming second half) -----
    float f[NQ];
    load_f_esopull(k, f, fi, j, t);

    // ----- Compute macroscopic quantities -----
    float rho_n, ux, uy;
#if NQ == 9
    compute_rho_u(f, rho_n, ux, uy);
#elif NQ == 19
    float uz;
    compute_rho_u(f, rho_n, ux, uy, uz);
#endif

    // ----- Boundary conditions -----
#if NQ == 9
    if (fl & LEGACY_SOLID) {
        if (x == 0) {
            float f_neb[NQ];
            unsigned long k_neb = linear_index(x + 1u, y);
            load_f_esopull(k_neb, f_neb, fi, j, t);  // approximate: reuse j
            // More accurate: compute j_neb separately
            unsigned long j_neb[NQ];
            compute_neighbors(k_neb, j_neb);
            load_f_esopull(k_neb, f_neb, fi, j_neb, t);
            apply_parabolic_inlet(f, f_neb, (float)y);
        }
        else if (x == (unsigned int)(NX - 1)) {
            unsigned long k_neb = linear_index(x - 1u, y);
            unsigned long j_neb[NQ];
            compute_neighbors(k_neb, j_neb);
            float f_neb[NQ];
            load_f_esopull(k_neb, f_neb, fi, j_neb, t);
            apply_pressure_outlet(f, f_neb, (float)y);
        }
    }

    if (y == 1 || y == (unsigned int)(NY - 2)) {
        apply_wall_bb_d2q9(y, f);
    }
#endif

    // ----- Forcing -----
    float Fin[NQ];
    float fxn = d_params.fx, fyn = d_params.fy;

    if (force_field != nullptr) {
        fxn += force_field[k];
        fyn += force_field[TOTAL_CELLS + k];
    }

    if (fxn != 0.0f || fyn != 0.0f) {
#if NQ == 9
        apply_guo_velocity_correction(ux, uy, fxn, fyn, rho_n);
        compute_guo_forcing(ux, uy, fxn, fyn, Fin);
#endif
    } else {
        zero_forcing(Fin);
    }

    // ----- Clamp velocity for stability -----
    const float CS_LIMIT = 0.57735027f;  // 1/sqrt(3)
    ux = fminf(fmaxf(ux, -CS_LIMIT), CS_LIMIT);
    uy = fminf(fmaxf(uy, -CS_LIMIT), CS_LIMIT);

    // ----- Collision -----
    if (fl & LEGACY_FLUID) {
        float feq[NQ];
#if NQ == 9
        compute_feq(rho_n, ux, uy, feq);

#if COLLISION_MODEL == 0
        collide_srt(f, feq, Fin, d_params.omega);
#elif COLLISION_MODEL == 1
        collide_trt(f, feq, Fin, d_params.omega);
#elif COLLISION_MODEL == 2
        collide_mrt(f, rho_n, ux, uy, Fin, d_params.omega);
#endif

#elif NQ == 19
        float fzn = d_params.fz;
        if (force_field != nullptr) fzn += force_field[2*TOTAL_CELLS + k];
        apply_guo_velocity_correction(ux, uy, uz, fxn, fyn, fzn, rho_n);
        compute_feq(rho_n, ux, uy, uz, feq);
        collide_srt(f, feq, Fin, d_params.omega);
#endif
    }

    // ----- Esoteric-Pull: write DDF (streaming first half) -----
    store_f_esopull(k, f, fi, j, t);

    // ----- Write macroscopic fields -----
    if (rho_arr != nullptr) rho_arr[k] = rho_n;
    if (u_arr != nullptr) {
        u_arr[k]               = ux;
        u_arr[TOTAL_CELLS + k] = uy;
    }

    // ----- Sensor observation -----
    if (fl & LEGACY_SENSOR) {
        int id_obj = indx[k];
        atomicAdd(&obs[DIM * id_obj],     ux);
        atomicAdd(&obs[DIM * id_obj + 1], uy);
    }
}

// ---------------------------------------------------------------------------
// InitializeEsoPull:  Initialize DDF via Esoteric-Pull store (t=1)
// ---------------------------------------------------------------------------
__global__ void InitializeEsoPull(uint8_t* flag, fpxx* fi)
{
#if DIM == 2
    unsigned int x, y;
    unsigned long k;
    index_from_thread(x, y, k);
    if (x >= (unsigned int)NX || y >= (unsigned int)NY) return;

    unsigned long j[NQ];
    compute_neighbors(k, j);

    float feq[NQ];
    if (y == 0 || y == NY - 1 || x == 0 || x == NX - 1) {
        flag[k] = LEGACY_SOLID;
        for (int i = 0; i < NQ; i++) feq[i] = 0.0f;
    } else {
        flag[k] = (uint8_t)LEGACY_FLUID;
        float u_init = U0 * 1.5f * (1.0f - 4.0f * ((float)y - 0.5f*(NY-1))
                                                  * ((float)y - 0.5f*(NY-1))
                                    / ((float)(NY-2) * (float)(NY-2)));
        for (int i = 0; i < NQ; i++) {
            float cu = (float)d_cx[i] * u_init;
            feq[i] = d_w[i] * RHO * (1.0f + 3.0f*cu + 4.5f*cu*cu - 1.5f*u_init*u_init);
#if USE_DDF_SHIFTING
            feq[i] -= d_w[i];
#endif
        }
    }
    // Write via Esoteric-Pull store at t=1 (odd step) for correct alignment
    store_f_esopull(k, feq, fi, j, 1ul);
#endif
}

#endif // CELERIS_STEP_ONE_STEP_ESOPULL_CU
