// CelerisLab – step/aux_kernels.cu
// Auxiliary kernels: CurvedBoundaryKernel, PrepareEsoPullCurvedKernel,
// ApplyEsoPullCurvedKernel, SensorKernel, ForceRegionKernel.
//
// Launched on compact lists (n_curved / n_sensor / n_cells threads).
// Included by kernel_v2.cu.  No standalone compilation.
//
// Double-buffer mode:
//   CurvedBoundaryKernel runs BEFORE OneStep.
//   SensorKernel / ForceRegionKernel run after.
//
// EsoPull mode:
//   PrepareEsoPullCurvedKernel + ApplyEsoPullCurvedKernel run before EsoPullStep.
//   SensorKernel / ForceRegionKernel run after, using t = step_count+1 for
//   physical-value semantics via esopull_semantic_helpers.cuh.
//
// Verified: Kan99b K2 (200k steps, MRT) -- bit-identical to double-buffer.
// ============================================================================

#ifndef CELERIS_STEP_AUX_KERNELS_CU
#define CELERIS_STEP_AUX_KERNELS_CU

#include "../streaming/esopull_semantic_helpers.cuh"
#include "../boundary/esopull_curved_helpers.cuh"

__device__ __forceinline__ float action_omega(const float* action, int body_id)
{
    if (action == nullptr || body_id < 0) return 0.0f;
    unsigned int base = (unsigned int)body_id * (unsigned int)OBS_BODY_SLOT_FLOATS;
    return action[base + (unsigned int)(OBS_BODY_SLOT_FLOATS - 1)];
}

__global__ void CurvedBoundaryKernel(
    fpxx* fi,
    const unsigned int* cl_fluid_idx,
    const unsigned char* cl_dir,
    const float* cl_q,
    const float* cl_rx,
    const float* cl_ry,
    const float* cl_rz,
    const unsigned char* cl_fallback_class,
    const int* cl_body_id,
    const float* action,
    float* obs,
    unsigned int n_curved)
{
    unsigned int tid = threadIdx.x + blockIdx.x * blockDim.x;
    if (tid >= n_curved) return;

    const unsigned long k_f = (unsigned long)cl_fluid_idx[tid];
    const unsigned int dir = (unsigned int)cl_dir[tid];
    const float q = cl_q[tid];
    const int bid = cl_body_id[tid];
    const float rx = cl_rx[tid];
    const float ry = cl_ry[tid];
    const unsigned char fallback_class = cl_fallback_class[tid];
    const float omega = action_omega(action, bid);

#if DIM == 2
    const float Uw = -omega * ry;
    const float Vw =  omega * rx;
    (void)cl_rz;
    apply_bouzidi_link(dir, q, k_f, fi, Uw, Vw, rx, ry, fallback_class, obs, bid);
#elif DIM == 3
    const float rz = cl_rz[tid];
    const float Uw = -omega * ry;
    const float Vw =  omega * rx;
    const float Ww = 0.0f;
    apply_bouzidi_link(dir, q, k_f, fi, Uw, Vw, Ww,
                       rx, ry, rz, fallback_class, obs, bid);
#endif
}
// End of CurvedBoundaryKernel

// ---------------------------------------------------------------------------
// PrepareEsoPullCurvedKernel -- read-only curved prep for EsoPull
//
// Reads donor populations from DDF via EsoPull semantic helpers,
// computes Bouzidi reflection + moving-wall, writes to intermediate buffer.
// Does NOT write directly to fi (paired Apply kernel does the scatter).
// ---------------------------------------------------------------------------
__global__ void PrepareEsoPullCurvedKernel(
    const fpxx* fi,
    const unsigned int* cl_fluid_idx,
    const unsigned char* cl_dir,
    const float* cl_q,
    const float* cl_rx,
    const float* cl_ry,
    const unsigned char* cl_fallback_class,
    const int* cl_body_id,
    const float* action,
    float* obs,
    float* corr_value,
    unsigned int* target_cell,
    unsigned char* target_slot,
    unsigned int n_curved,
    unsigned long t)
{
    unsigned int tid = threadIdx.x + blockIdx.x * blockDim.x;
    if (tid >= n_curved) return;

    const unsigned long k_f = (unsigned long)cl_fluid_idx[tid];
    const unsigned int dir = (unsigned int)cl_dir[tid];
    const unsigned int dir_opp = (unsigned int)opp_dir((int)dir);
    const float q = cl_q[tid];
    const float rx = cl_rx[tid];
    const float ry = cl_ry[tid];
    const unsigned char fallback = cl_fallback_class[tid];
    const int bid = cl_body_id[tid];

    unsigned int x, y;
    coordinates(k_f, x, y);
    unsigned long j_f[NQ];
    compute_neighbors(x, y, j_f);
    const unsigned long k_s = j_f[dir];

    const float omega = action_omega(action, bid);
    const float Uw = -omega * ry;
    const float Vw =  omega * rx;

    const float f_toward = load_physical_dir_esopull(k_f, dir, fi, j_f, t);

    float f_reflected;
    if (fallback == CURVED_FALLBACK_BOUZIDI) {
        if (q < 0.5f) {
            const unsigned long k_ff = j_f[dir_opp];
            unsigned int xff, yff;
            coordinates(k_ff, xff, yff);
            unsigned long j_ff[NQ];
            compute_neighbors(xff, yff, j_ff);
            const float f_toward_ff = load_physical_dir_esopull(k_ff, dir, fi, j_ff, t);
            f_reflected = compute_bouzidi_reflection(f_toward, f_toward_ff, 0.0f, q, fallback);
        } else {
            const float f_opp_same = load_physical_dir_esopull(k_f, dir_opp, fi, j_f, t);
            f_reflected = compute_bouzidi_reflection(f_toward, 0.0f, f_opp_same, q, fallback);
        }
    } else {
        f_reflected = compute_bouzidi_reflection(f_toward, 0.0f, 0.0f, q, fallback);
    }
    f_reflected += compute_moving_wall_addend(dir, q, fallback, Uw, Vw);

    unsigned long dst_cell;
    unsigned int dst_slot;
    resolve_esopull_target(k_f, k_s, dir, t, dst_cell, dst_slot);

    corr_value[tid] = f_reflected;
    target_cell[tid] = (unsigned int)dst_cell;
    target_slot[tid] = (unsigned char)dst_slot;

    if (obs != nullptr) {
#if USE_DDF_SHIFTING
        const float w_i = (float)d_w[dir];
        const float fx = (float)d_cx[dir] * ((f_toward + w_i) + (f_reflected + w_i));
        const float fy = (float)d_cy[dir] * ((f_toward + w_i) + (f_reflected + w_i));
#else
        const float fx = (float)d_cx[dir] * (f_toward + f_reflected);
        const float fy = (float)d_cy[dir] * (f_toward + f_reflected);
#endif
        atomicAdd(&obs[obs_force_index(bid, 0)], fx);
        atomicAdd(&obs[obs_force_index(bid, 1)], fy);
        const float tz = rx * fy - ry * fx;
        atomicAdd(&obs[obs_torque_index(bid, 0)], tz);
    }
}

// ---------------------------------------------------------------------------
// ApplyEsoPullCurvedKernel -- scatter intermediate buffer to DDF
// Read-only on intermediate buffer, write-only on DDF.  No geometry/physics.
// ---------------------------------------------------------------------------
__global__ void ApplyEsoPullCurvedKernel(
    fpxx* fi,
    const float* corr_value,
    const unsigned int* target_cell,
    const unsigned char* target_slot,
    unsigned int n_curved)
{
    unsigned int tid = threadIdx.x + blockIdx.x * blockDim.x;
    if (tid >= n_curved) return;

    store_ddf(fi,
              index_f((unsigned long)target_cell[tid], (unsigned int)target_slot[tid]),
              corr_value[tid]);
}

// ---------------------------------------------------------------------------
// SensorKernel — accumulate area-averaged velocity from DDF.
//
// Double-buffer mode (t == 0): fi[k,i] is the physical value of direction i
// at cell k.  Direct raw reads are correct.
//
// EsoPull mode (t > 0): fi[k,i] stores the EsoPull backing slot, which may
// not be the physical direction i.  Reconstruct physical DDF using the same
// load_f_esopull semantics (paired directions, parity-dependent slots).
// Runs after EsoPullStep, so the physical interpretation of fi follows the
// store_f_esopull contract: after store, t+1 parity determines physical layout.
// ---------------------------------------------------------------------------
__global__ void SensorKernel(
    const fpxx* fi,
    const unsigned short* flag,
    const unsigned int* sensor_cells,
    const int* sensor_obj_id,
    float* obs,
    unsigned int n_sensor,
    unsigned long t)        // 0 = double_buffer, >0 = esopull step count
{
    unsigned int tid = threadIdx.x + blockIdx.x * blockDim.x;
    if (tid >= n_sensor) return;

    unsigned long k = (unsigned long)sensor_cells[tid];
    int id_obj = sensor_obj_id[tid];

    float f[NQ];
    if (t == 0ul) {
        // Double-buffer mode: raw slot = physical direction
        for (int i = 0; i < NQ; i++)
            f[i] = load_ddf(fi, index_f(k, (unsigned int)i));
    } else {
        // EsoPull mode: reconstruct physical values via shared helper
        unsigned int x, y;
        coordinates(k, x, y);
        unsigned long j[NQ];
        compute_neighbors(x, y, j);
        load_physical_node_esopull(k, fi, j, t, f);
    }

#if DIM == 2
    float rho_n, ux, uy;
    compute_rho_u(f, rho_n, ux, uy);
    atomicAdd(&obs[obs_sensor_index(id_obj, 0)], ux);
    atomicAdd(&obs[obs_sensor_index(id_obj, 1)], uy);
#elif DIM == 3
    float rho_n, ux, uy, uz;
    compute_rho_u(f, rho_n, ux, uy, uz);
    atomicAdd(&obs[obs_sensor_index(id_obj, 0)], ux);
    atomicAdd(&obs[obs_sensor_index(id_obj, 1)], uy);
    atomicAdd(&obs[obs_sensor_index(id_obj, 2)], uz);
#endif

    (void)flag;
}

// ---------------------------------------------------------------------------
// ForceRegionKernel — operator-split local Guo forcing correction
//
// Runs after the main collide+stream step.  Applies a Guo forcing term to
// cells within a force_region footprint.  This is an operator-split correction,
// not inline in OneStep's collide_dispatch.
//
// The Guo formula and c_tau = 1 - omega/2 conventions match collide_srt,
// collide_mrt, and collide_trt exactly.
// ---------------------------------------------------------------------------
__device__ __forceinline__ float action_force_x(const float* action, int body_id)
{
    if (action == nullptr || body_id < 0) return 0.0f;
    unsigned int base = (unsigned int)body_id * (unsigned int)OBS_BODY_SLOT_FLOATS;
    return action[base + 0u];
}

__device__ __forceinline__ float action_force_y(const float* action, int body_id)
{
    if (action == nullptr || body_id < 0) return 0.0f;
    unsigned int base = (unsigned int)body_id * (unsigned int)OBS_BODY_SLOT_FLOATS;
    return action[base + 1u];
}

__global__ void ForceRegionKernel(
    fpxx* fi,
    const unsigned short* flag,
    const unsigned int* fr_cells,
    const int* fr_obj_id,
    const float* action,
    unsigned int n_cells,
    unsigned long t)        // 0 = double_buffer, >0 = esopull step count
{
    unsigned int tid = threadIdx.x + blockIdx.x * blockDim.x;
    if (tid >= n_cells) return;

    unsigned long k = (unsigned long)fr_cells[tid];
    int bid = fr_obj_id[tid];
    float fx = action_force_x(action, bid);
    float fy = action_force_y(action, bid);
    if (fx == 0.0f && fy == 0.0f) return;

    if (!is_fluid(flag[k])) return;

    // Read physical DDF with correct semantics for the current mode
    float f[NQ];
    unsigned int x, y;
    coordinates(k, x, y);
    unsigned long j[NQ];
    compute_neighbors(x, y, j);

    if (t == 0ul) {
        // Double-buffer: raw load
        for (int i = 0; i < NQ; i++)
            f[i] = load_ddf(fi, index_f(k, (unsigned int)i));
    } else {
        // EsoPull: reconstruct physical f[] from backing slots
        load_physical_node_esopull(k, fi, j, t, f);
    }

#if DIM == 2
    float rho_n, ux, uy;
    compute_rho_u(f, rho_n, ux, uy);

    float rho2 = 0.5f / rho_n;
    ux = fmaf(fx, rho2, ux);
    uy = fmaf(fy, rho2, uy);

    float Fin[NQ];
    compute_guo_forcing(ux, uy, fx, fy, Fin);

    float omega = d_params.omega;
    float c_tau = 1.0f - 0.5f * omega;
    #pragma unroll
    for (int i = 0; i < NQ; i++)
        f[i] += c_tau * Fin[i];
#elif DIM == 3
    float rho_n, ux, uy, uz;
    compute_rho_u(f, rho_n, ux, uy, uz);

    float rho2 = 0.5f / rho_n;
    ux = fmaf(fx, rho2, ux);
    uy = fmaf(fy, rho2, uy);
    float Fin[NQ];
    compute_guo_forcing(ux, uy, uz, fx, fy, 0.0f, Fin);

    float omega = d_params.omega;
    float c_tau = 1.0f - 0.5f * omega;
    #pragma unroll
    for (int i = 0; i < NQ; i++)
        f[i] += c_tau * Fin[i];
#endif

    // Write back with correct semantics
    if (t == 0ul) {
        #pragma unroll
        for (int i = 0; i < NQ; i++)
            store_ddf(fi, index_f(k, (unsigned int)i), f[i]);
    } else {
        // EsoPull: scatter modified physical f[] back to backing slots
        store_physical_node_esopull(k, f, fi, j, t);
    }
}

#endif // CELERIS_STEP_AUX_KERNELS_CU
