// CelerisLab – step/one_step_double.cu
// Main LBM step kernel using double-buffer PULL streaming.
//
// Workflow per node:
//   1. Pull-stream: gather DDF from neighbors → local f[NQ]
//   2. Compute macroscopic quantities (ρ, u)
//   3. Apply boundary conditions (inlet/outlet/wall)
//   4. Compute equilibrium + forcing
//   5. Collision (SRT / TRT / MRT)
//   6. Write to output buffer
//
// Kernel signature maintains backward compatibility with driver.py:
//   flag*   – cell flags (uint8 for legacy, uint32 for new)
//   fi_in*  – input DDF array  (fpxx)
//   fi_out* – output DDF array (fpxx)
//   indx*   – per-cell offset into delta pool
//   delta*  – curved boundary parameter pool
//   action* – per-object control input
//   obs*    – per-object observation output (force / velocity)
// ============================================================================

#ifndef CELERIS_STEP_ONE_STEP_DOUBLE_CU
#define CELERIS_STEP_ONE_STEP_DOUBLE_CU

__global__ void StreamCollideDouble(
    const uint8_t* __restrict__ flag,
    const fpxx*    __restrict__ fi_in,
    fpxx*          __restrict__ fi_out,
    const int32_t* __restrict__ indx,
    const float*   __restrict__ delta,
    const float*   __restrict__ action,
    float*         __restrict__ obs,
    float*         __restrict__ rho_arr,    // macro output (can be NULL)
    float*         __restrict__ u_arr)      // macro output (can be NULL)
{
    // ----- Thread → node mapping (2D) -----
#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

    // ----- Read flag -----
    uint8_t fl = flag[k];

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

    // ----- Pull-stream: load DDF from neighbors -----
    float f[NQ];
    stream_pull_load(k, f, fi_in, j);

    // ----- 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 (inlet/outlet/wall) -----
#if NQ == 9
    if (fl & LEGACY_SOLID) {
        if (x == 0) {
            float f_neb[NQ];
            unsigned long k_neb = linear_index(x + 1u, y);
            for (int i = 0; i < NQ; i++) {
                f_neb[i] = load_ddf(fi_in, index_f(k_neb, (unsigned int)i));
            }
            apply_parabolic_inlet(f, f_neb, (float)y);
        }
        else if (x == NX - 1) {
            float f_neb[NQ];
            unsigned long k_neb = linear_index(x - 1u, y);
            for (int i = 0; i < NQ; i++) {
                f_neb[i] = load_ddf(fi_in, index_f(k_neb, (unsigned int)i));
            }
            apply_pressure_outlet(f, f_neb, (float)y);
        }
    }
#elif NQ == 19
    if (fl & LEGACY_SOLID) {
        if (x == 0) {
            float f_neb[NQ];
            unsigned long k_neb = linear_index(x + 1u, y, z);
            for (int i = 0; i < NQ; i++) {
                f_neb[i] = load_ddf(fi_in, index_f(k_neb, (unsigned int)i));
            }
            apply_parabolic_inlet_3d(f, f_neb, (float)y);
        }
        else if (x == NX - 1) {
            float f_neb[NQ];
            unsigned long k_neb = linear_index(x - 1u, y, z);
            for (int i = 0; i < NQ; i++) {
                f_neb[i] = load_ddf(fi_in, index_f(k_neb, (unsigned int)i));
            }
            apply_pressure_outlet_3d(f, f_neb, (float)y);
        }
    }
#endif

    // ----- Wall bounce-back (top/bottom walls) -----
#if NQ == 9
    if (y == 1 || y == NY - 2) {
        apply_wall_bb_d2q9(y, f);
    }
#elif NQ == 19
    if (y == 1 || y == NY - 2) {
        apply_wall_bb_d3q19_y(y, f);
    }
#endif

    // ----- Compute equilibrium -----
    if (fl & LEGACY_FLUID) {
        float feq[NQ];
        float Fin[NQ];

#if NQ == 9
        compute_feq(rho_n, ux, uy, feq);
        zero_forcing(Fin);

        // ----- Collision -----
#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
        compute_feq(rho_n, ux, uy, uz, feq);
        zero_forcing(Fin);
        collide_srt(f, feq, Fin, d_params.omega);
#endif
    }

    // ----- Write to output buffer -----
    stream_pull_store(k, f, fi_out);

    // ----- Write macroscopic fields (optional) -----
    if (rho_arr != nullptr) {
        rho_arr[k] = rho_n;
    }
    if (u_arr != nullptr) {
#if DIM == 2
        u_arr[k]               = ux;
        u_arr[TOTAL_CELLS + k] = uy;
#elif DIM == 3
        u_arr[k]                   = ux;
        u_arr[TOTAL_CELLS + k]     = uy;
        u_arr[2 * TOTAL_CELLS + k] = uz;
#endif
    }

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

// ---------------------------------------------------------------------------
// Curved boundary post-processing kernel (separate launch)
// Processes SOLID+INTERFACE nodes → modifies fi_out at neighbors.
// ---------------------------------------------------------------------------
__global__ void CurvedBoundaryPost(
    const uint8_t* __restrict__ flag,
    fpxx*          __restrict__ fi_out,
    const int32_t* __restrict__ indx,
    const float*   __restrict__ delta,
    const float*   __restrict__ action,
    float*         __restrict__ obs)
{
#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];

    if ((fl & LEGACY_SOLID) && (fl & LEGACY_INTERFACE)) {
#if DIM == 2
        int id_off = indx[k];
        int id_obj = *reinterpret_cast<const int*>(&delta[id_off]);

        // Wall velocity from action + normal
        float Uw = action[id_obj] * delta[id_off + NQ];
        float Vw = action[id_obj] * delta[id_off + NQ + 1];

        float* obs_fx = (obs != nullptr) ? &obs[DIM * id_obj]     : nullptr;
        float* obs_fy = (obs != nullptr) ? &obs[DIM * id_obj + 1] : nullptr;

        apply_curved_boundary(k, x, y, fi_out, delta, id_off, Uw, Vw, obs_fx, obs_fy);
#endif // DIM == 2 curved BC
    }
}

// ---------------------------------------------------------------------------
// InitTubeFlow:  Initialize parabolic tube flow  (backward compatible)
// ---------------------------------------------------------------------------
__global__ void InitTubeFlow(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;

    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)));

    if (y == 0 || y == NY - 1 || x == 0 || x == NX - 1) {
        flag[k] = LEGACY_SOLID;
        for (int i = 0; i < NQ; i++) {
            store_ddf(fi, index_f(k, (unsigned int)i), 0.0f);
        }
    } else {
        flag[k] = (uint8_t)LEGACY_FLUID;
        // feq with rho=RHO, u=(u_init, 0)
        for (int i = 0; i < NQ; i++) {
            float cu = (float)d_cx[i] * u_init;  // cy=0 for parabolic flow
            float val = d_w[i] * RHO * (1.0f + 3.0f * cu + 4.5f * cu * cu - 1.5f * u_init * u_init);
#if USE_DDF_SHIFTING
            val -= d_w[i];   // store f_tilde = f - w
#endif
            store_ddf(fi, index_f(k, (unsigned int)i), val);
        }
    }
#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;

    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)));

    // Walls: y=0, y=NY-1, x=0, x=NX-1;  z is periodic
    if (y == 0 || y == NY - 1 || x == 0 || x == NX - 1) {
        flag[k] = LEGACY_SOLID;
        for (int i = 0; i < NQ; i++) {
            store_ddf(fi, index_f(k, (unsigned int)i), 0.0f);
        }
    } else {
        flag[k] = (uint8_t)LEGACY_FLUID;
        for (int i = 0; i < NQ; i++) {
            float cu = (float)d_cx[i] * u_init;
            float val = d_w[i] * RHO * (1.0f + 3.0f * cu + 4.5f * cu * cu - 1.5f * u_init * u_init);
#if USE_DDF_SHIFTING
            val -= d_w[i];
#endif
            store_ddf(fi, index_f(k, (unsigned int)i), val);
        }
    }
#endif
}

// ---------------------------------------------------------------------------
// UpdateMacro:  Recompute rho/u from DDF (for diagnostics/output)
// ---------------------------------------------------------------------------
__global__ void UpdateMacro(
    const fpxx*    __restrict__ fi,
    float*         __restrict__ rho_arr,
    float*         __restrict__ u_arr,
    const uint8_t* __restrict__ flag)
{
#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;

    if (flag[k] & LEGACY_SOLID) return;

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

    float f[NQ];
    // For double-buffer, just read from the current buffer (no streaming)
    for (int i = 0; i < NQ; i++) {
        f[i] = load_ddf(fi, index_f(k, (unsigned int)i));
    }

    float rho_n, ux, uy;
    compute_rho_u(f, rho_n, ux, uy);

    rho_arr[k]               = rho_n;
    u_arr[k]                 = ux;
    u_arr[TOTAL_CELLS + k]   = uy;
#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;

    if (flag[k] & LEGACY_SOLID) return;

    float f[NQ];
    for (int i = 0; i < NQ; i++) {
        f[i] = load_ddf(fi, index_f(k, (unsigned int)i));
    }

    float rho_n, ux, uy, uz;
    compute_rho_u(f, rho_n, ux, uy, uz);

    rho_arr[k]                   = rho_n;
    u_arr[k]                     = ux;
    u_arr[TOTAL_CELLS + k]       = uy;
    u_arr[2 * TOTAL_CELLS + k]   = uz;
#endif
}

#endif // CELERIS_STEP_ONE_STEP_DOUBLE_CU
