// CelerisLab – step/init_flow.cu
// Channel-flow initialization kernels (double-buffer and EsoPull).
// Included by kernel_v2.cu after all headers.  No standalone compilation.
// ============================================================================

#ifndef CELERIS_STEP_INIT_FLOW_CU
#define CELERIS_STEP_INIT_FLOW_CU

// ---- Helper: write equilibrium for a given velocity ----
__device__ __forceinline__ void write_equilibrium(
    fpxx* fi, unsigned long k, float u_init)
{
    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);
    }
}

__device__ __forceinline__ void write_rest_equilibrium(
    fpxx* fi, unsigned long k)
{
    for (int i = 0; i < NQ; i++) {
        float val = d_w[i] * RHO;
#if USE_DDF_SHIFTING
        val -= d_w[i];
#endif
        store_ddf(fi, index_f(k, (unsigned int)i), val);
    }
}

// ---- EsoPull layout helpers: equilibrium in EsoPull backing layout at t=0 ----
// In EsoPull t=0 (even step), store_f_esopull writes:
//   f[i_odd]  → fi[ j[i_odd],  i_odd      ]   (neighbor, forward slot)
//   f[i_even] → fi[ k,         i_even      ]   (local, reverse slot)
// These helpers write equilibrium values directly in the correct EsoPull
// layout so that load_f_esopull(t=0) reconstructs the correct f_eq[].
#if DIM == 2
__device__ __forceinline__ void write_equilibrium_esopull(
    fpxx* fi, unsigned long k, unsigned int x, unsigned int y, float u_init)
{
    unsigned long j[NQ];
    compute_neighbors(x, y, j);

    // --- rest particle: always local, slot 0 ---
    float val0 = d_w[0] * RHO * (1.0f - 1.5f * u_init * u_init);
#if USE_DDF_SHIFTING
    val0 -= d_w[0];
#endif
    store_ddf(fi, index_f(k, 0u), val0);

    // --- paired directions (i=1,3,5,7): {E,W}, {N,S}, {NE,SW}, {SE,NW} ---
    for (int i = 1; i < NQ; i += 2) {
        float cu_odd  = (float)d_cx[i] * u_init;
        float val_odd = d_w[i] * RHO
            * (1.0f + 3.0f * cu_odd + 4.5f * cu_odd * cu_odd
               - 1.5f * u_init * u_init);
        float cu_even = (float)d_cx[i + 1] * u_init;
        float val_even = d_w[i + 1] * RHO
            * (1.0f + 3.0f * cu_even + 4.5f * cu_even * cu_even
               - 1.5f * u_init * u_init);
#if USE_DDF_SHIFTING
        val_odd  -= d_w[i];
        val_even -= d_w[i + 1];
#endif
        // t=0 even step (load_f_esopull): f[i] reads from local slot i+1,
        // f[i+1] reads from neighbor slot i.
        // Therefore: fi[n, i+1] = f_eq[i]  (odd goes to local reverse),
        //            fi[j[i], i] = f_eq[i+1] (even goes to neighbor forward)
        store_ddf(fi, index_f(k,    (unsigned int)(i + 1)), val_odd);
        store_ddf(fi, index_f(j[i], (unsigned int)i),          val_even);
    }
}

__device__ __forceinline__ void write_rest_equilibrium_esopull(
    fpxx* fi, unsigned long k, unsigned int x, unsigned int y)
{
    unsigned long j[NQ];
    compute_neighbors(x, y, j);

    float val0 = d_w[0] * RHO;
#if USE_DDF_SHIFTING
    val0 -= d_w[0];
#endif
    store_ddf(fi, index_f(k, 0u), val0);

    for (int i = 1; i < NQ; i += 2) {
        float val_odd  = d_w[i] * RHO;
        float val_even = d_w[i + 1] * RHO;
#if USE_DDF_SHIFTING
        val_odd  -= d_w[i];
        val_even -= d_w[i + 1];
#endif
        store_ddf(fi, index_f(k,    (unsigned int)(i + 1)), val_odd);
        store_ddf(fi, index_f(j[i], (unsigned int)i),          val_even);
    }
}
#endif // DIM == 2

__device__ __forceinline__ uint16_t channel_flag_from_coords(
    unsigned int x, unsigned int y)
{
    if (y == 0u || y == (unsigned int)(NY - 1))
        return (uint16_t)(FLAG_SOLID | FLAG_BC_WALL);
    if (x == 0u)
        return (uint16_t)(FLAG_SOLID | FLAG_BC_INLET);
    if (x == (unsigned int)(NX - 1))
        return (uint16_t)(FLAG_SOLID | FLAG_BC_OUTLET);
    return FLAG_FLUID;
}

__device__ __forceinline__ uint16_t finalize_domain_flag(
    uint16_t fl, unsigned int x, unsigned int y)
{
    if (is_obstacle(fl))
        return fl;

    const uint16_t base = channel_flag_from_coords(x, y);
    if (is_sensor(fl) && base == FLAG_FLUID)
        return (uint16_t)(base | FLAG_SENSOR);
    if (is_force_region(fl) && base == FLAG_FLUID)
        return (uint16_t)(base | FLAG_FRC_REGION);
    return base;
}

// ============================================================================
// InitTubeFlow_v2 — double-buffer channel initialization
// ============================================================================
__global__ void InitTubeFlow_v2(uint16_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;

    const uint16_t fl = finalize_domain_flag(flag[k], x, y);
    flag[k] = fl;
    if (is_solid(fl)) {
        write_rest_equilibrium(fi, k);
    } else {
        write_equilibrium(fi, k, inlet_target_u((float)y));
    }
#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;

    const uint16_t fl = finalize_domain_flag(flag[k], x, y);
    flag[k] = fl;
    if (is_solid(fl)) {
        write_rest_equilibrium(fi, k);
    } else {
        write_equilibrium(fi, k, inlet_target_u((float)y));
    }
#endif
}

// ============================================================================
// InitEsoPull — EsoPull single-buffer channel initialization
// ============================================================================
__global__ void InitEsoPull(uint16_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;

    const uint16_t fl = finalize_domain_flag(flag[k], x, y);
    flag[k] = fl;
    // Write equilibrium directly in EsoPull t=0 backing layout so that
    // load_f_esopull(k, f, fi, j, 0) at the first EsoPullStep reconstructs
    // the correct physical f_eq[].  The neighbor scatter uses a slot
    // partitioning per cell with no write conflicts across threads.
    if (is_solid(fl)) {
        write_rest_equilibrium_esopull(fi, k, x, y);
    } else {
        write_equilibrium_esopull(fi, k, x, y, inlet_target_u((float)y));
    }
#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;

    const uint16_t fl = finalize_domain_flag(flag[k], x, y);
    flag[k] = fl;
    if (is_solid(fl)) {
        write_rest_equilibrium(fi, k);
    } else {
        write_equilibrium(fi, k, inlet_target_u((float)y));
    }
#endif
}

#endif // CELERIS_STEP_INIT_FLOW_CU