// CelerisLab – step/one_step_esopull.cu
// Esoteric-Pull single-buffer step kernel.
// Included by kernel_v2.cu.  No standalone compilation.
//
// CURVED BC: supported via PrepareEsoPullCurvedKernel + ApplyEsoPullCurvedKernel
// in aux_kernels.cu, launched by stepper.py before EsoPullStep.
// ============================================================================

#ifndef CELERIS_STEP_ONE_STEP_ESOPULL_CU
#define CELERIS_STEP_ONE_STEP_ESOPULL_CU

// ---- Helper: apply inlet/outlet BC in esopull mode ----
#if DIM == 2
__device__ __forceinline__ void apply_boundary_esopull(
    float* f, uint16_t fl, unsigned int x, unsigned int y,
    const fpxx* fi, unsigned long t)
{
    // Curved-BC nodes have been pre-processed by PrepareEsoPullCurvedKernel
    // + ApplyEsoPullCurvedKernel, so is_curved() cells here have valid backed
    // populations.  Bounce-back on obstacle cells not covered by curved
    // (is_obstacle && !is_curved) is handled in EsoPullStep.
    if (is_curved(fl)) return;

    bool interior_y = (y > 0u) && (y < (unsigned int)(NY - 1));

    if (is_inlet(fl) && interior_y) {
        apply_inlet_esopull_d2q9(f, x, y, fi, t);
    }
    else if (is_outlet(fl) && interior_y) {
        apply_outlet_esopull_d2q9(f, x, y, fi, t);
    }
    else {
        bounce_back_swap(f);
    }
}
#endif // DIM == 2

#if DIM == 3
__device__ __forceinline__ void apply_boundary_esopull_3d(
    float* f, uint16_t fl, unsigned int x, unsigned int y, unsigned int z,
    const fpxx* fi, unsigned long t)
{
    if (is_curved(fl)) return;

    bool interior_y = (y > 0u) && (y < (unsigned int)(NY - 1));

    if (is_inlet(fl) && interior_y) {
        apply_inlet_esopull_d3q19(f, x, y, z, fi, t);
    }
    else if (is_outlet(fl) && interior_y) {
        apply_outlet_esopull_d3q19(f, x, y, z, fi, t);
    }
    else {
        bounce_back_swap(f);
    }
}
#endif

// ============================================================================
// EsoPullStep — Esoteric-Pull single-buffer step kernel
// ============================================================================
__global__ __launch_bounds__(NT, 2)
void EsoPullStep(
    fpxx* fi, uint16_t* flag,
    float* action, float* obs,
    unsigned long t)
{
#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;

    uint16_t fl = flag[k];
    unsigned long j[NQ];
    const bool interior_y = (y > 0u) && (y < (unsigned int)(NY - 1));
    compute_neighbors(x, y, j);

    float f[NQ];
    load_f_esopull(k, f, fi, j, t);

    // Pure solid early exit
    if (is_solid(fl) && !has_bc(fl)) {
        bounce_back_swap(f);
        store_f_esopull(k, f, fi, j, t);
        return;
    }

    // Boundary nodes
    if (is_solid(fl))
        apply_boundary_esopull(f, fl, x, y, fi, t);

    if (is_obstacle(fl) && !is_curved(fl))
        bounce_back_swap(f);

    // ---- y-wall adjacent row correction ----
    if (is_fluid(fl) && (y == 1u || y == (unsigned int)(NY - 2))) {
#if Y_WALL_BC == 1
        apply_wall_freeslip_d2q9_y_pull(y, f, fi, k);
#else
        apply_wall_bb_d2q9_y_pull(y, f, fi, k);
#endif
    }

    // Collision
    const bool collide_inlet_ghost = is_inlet(fl) && interior_y
                                   && inlet_scheme_uses_post_collision_ghost();
    if (is_fluid(fl) || collide_inlet_ghost) {
        float rho_n, ux, uy;
        compute_rho_u(f, rho_n, ux, uy);
        collide_dispatch(f, rho_n, ux, uy);
    }

    store_f_esopull(k, f, fi, j, t);

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

    uint16_t fl = flag[k];
    unsigned long j[NQ];
    const bool interior_y = (y > 0u) && (y < (unsigned int)(NY - 1));
    compute_neighbors(k, j);

    float f[NQ];
    load_f_esopull(k, f, fi, j, t);

    if (is_solid(fl) && !has_bc(fl)) {
        bounce_back_swap(f);
        store_f_esopull(k, f, fi, j, t);
        return;
    }

    if (is_solid(fl))
        apply_boundary_esopull_3d(f, fl, x, y, z, fi, t);

    if (is_obstacle(fl) && !is_curved(fl))
        bounce_back_swap(f);

    if (is_fluid(fl) && (y == 1u || y == (unsigned int)(NY - 2)))
        apply_wall_bb_d3q19_y_pull(y, f, fi, k);

    const bool collide_inlet_ghost = is_inlet(fl) && interior_y
                                   && inlet_scheme_uses_post_collision_ghost();
    if (is_fluid(fl) || collide_inlet_ghost) {
        float rho_n, ux, uy, uz;
        compute_rho_u(f, rho_n, ux, uy, uz);
        collide_dispatch(f, rho_n, ux, uy, uz);
    }

    store_f_esopull(k, f, fi, j, t);
#endif
}

#endif // CELERIS_STEP_ONE_STEP_ESOPULL_CU
