// CelerisLab – kernel_v2.cu
// ============================================================================
// Modular compilation entry point  (Stage 1 Architecture)
//
// This file includes all modular headers and step kernels,
// and exports extern "C" functions for PyCUDA.
//
// Controlled by macros.h:
//   NQ              = 9 (D2Q9) or 19 (D3Q19)
//   COLLISION_MODEL = 0 (SRT) / 1 (TRT) / 2 (MRT)
//   STREAMING_MODEL = 0 (double-buffer) / 1 (Esoteric-Pull)
//   STORE_PRECISION = 0 (FP32) / 1 (FP16S) / 2 (FP16C)
//   USE_DDF_SHIFTING= 0 / 1
// ============================================================================

#include <stdio.h>
#include <stdint.h>
#include <cuda.h>

// ---------------------------------------------------------------------------
// Layer 0: Configuration (compile-time)
// ---------------------------------------------------------------------------
#include "macros.h"

// ---------------------------------------------------------------------------
// Layer 1: Core primitives
// ---------------------------------------------------------------------------
#include "core/precision.cuh"
#include "core/flags.cuh"
#include "core/descriptors.cuh"
#include "core/layout.cuh"
#include "core/params.cuh"

// ---------------------------------------------------------------------------
// Layer 2: Operators
// ---------------------------------------------------------------------------
#include "operators/macro.cuh"
#include "operators/collision_srt.cuh"
#include "operators/collision_trt.cuh"
#include "operators/collision_mrt.cuh"
#include "operators/forcing_guo.cuh"

// ---------------------------------------------------------------------------
// Layer 3: Streaming
// ---------------------------------------------------------------------------
#include "streaming/pull_double_buffer.cuh"
#include "streaming/esopull_single_buffer.cuh"

// ---------------------------------------------------------------------------
// Layer 4: Boundary conditions
// ---------------------------------------------------------------------------
#include "boundary/bounce_back.cuh"
#include "boundary/curved_boundary.cuh"
#include "boundary/inlet_outlet.cuh"

// ---------------------------------------------------------------------------
// Layer 5: Step kernels (selected by STREAMING_MODEL)
// ---------------------------------------------------------------------------
#include "step/one_step_double.cu"
#include "step/one_step_esopull.cu"

// ---------------------------------------------------------------------------
// Layer 6: IBM kernels (always available, launched separately)
// ---------------------------------------------------------------------------
#include "boundary/ibm_kernels.cuh"

// ============================================================================
// Extern "C" wrappers  (for PyCUDA / ctypes)
//
// Naming:
//   OneStep            – backward-compatible main step (double-buffer)
//   InitTubeFlow_v2    – channel initialization (new ordering)
//   StreamCollide      – new API main step (double or esopull)
//   CurvedBoundary     – post-step curved boundary correction
//   UpdateFields       – recompute rho/u from DDF
//   InitEsoPull        – Esoteric-Pull initialization
//   IBM_UpdateLag      – IBM Lagrangian point update
//   IBM_Eul2Lag        – Euler → Lagrangian interpolation
//   IBM_Lag2Eul        – Lagrangian → Euler spreading
// ============================================================================
extern "C"
{

// ----- Backward-compatible init (uses new direction ordering) -----
__global__ void InitTubeFlow_v2(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;
        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);
        }
    }
#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)));

    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
}

// ----- Main step (double-buffer) -----
// Signature compatible with driver.py: flag, fi_in, fi_out, indx, delta, action, obs
__global__ void OneStep(
    uint8_t*  flag,
    fpxx*     fi_in,
    fpxx*     fi_out,
    int32_t*  indx,
    float*    delta,
    float*    action,
    float*    obs)
{
    // Redirect to the modular StreamCollideDouble kernel body
    // (We inline here to keep a single extern "C" entry point.)

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

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

    float f[NQ];

#if STREAMING_MODEL == 0
    stream_pull_load(k, f, fi_in, j);
#else
    // Esoteric-Pull for OneStep compat would need timestep — not available
    // in legacy signature. Use double-buffer mode only for OneStep.
    stream_pull_load(k, f, fi_in, j);
#endif

    float rho_n, ux, uy;
#if NQ == 9
    compute_rho_u(f, rho_n, ux, uy);

    // Inlet / outlet / wall bounce-back
    if (fl & LEGACY_SOLID) {
        bool interior_y = (y > 0u) && (y < (unsigned int)(NY - 1));
        if (x == 0 && interior_y) {
            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 == (unsigned int)(NX - 1) && interior_y) {
            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);
        }
        else {
            // Wall / corner: full-way bounce-back (swap all pairs)
            #pragma unroll
            for (int i = 1; i < NQ; i += 2) {
                float t = f[i]; f[i] = f[i+1]; f[i+1] = t;
            }
        }
    }

    // Collision
    if (fl & LEGACY_FLUID) {
        float feq[NQ], Fin[NQ];
        compute_feq(rho_n, ux, uy, feq);
        zero_forcing(Fin);

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

    stream_pull_store(k, f, fi_out);

    // Curved boundary processing
    if ((fl & LEGACY_SOLID) && (fl & LEGACY_INTERFACE)) {
        int id_off = indx[k];
        int id_obj = *reinterpret_cast<int*>(&delta[id_off]);
        float Uw = action[id_obj] * delta[id_off + NQ];
        float Vw = action[id_obj] * delta[id_off + NQ + 1];
        float* obs_fx = &obs[DIM * id_obj];
        float* obs_fy = &obs[DIM * id_obj + 1];
        apply_curved_boundary(k, x, y, fi_out, delta, id_off, Uw, Vw, obs_fx, obs_fy);
    }

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

#elif NQ == 19
    // ---- Macroscopic (computed for all nodes; only used for fluid collision) ----
    float uz;
    compute_rho_u(f, rho_n, ux, uy, uz);

    // ---- Boundary conditions ----
    if (fl & LEGACY_SOLID) {
        // Interior of inlet/outlet planes (skip corners at y=0/NY-1)
        bool interior_y = (y > 0u) && (y < (unsigned int)(NY - 1));
        if (x == 0 && interior_y) {
            // Parabolic velocity inlet
            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 == (unsigned int)(NX - 1) && interior_y) {
            // Pressure outlet
            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);
        }
        else {
            // Wall nodes and corners: full-way bounce-back.
            #pragma unroll
            for (int i = 1; i < NQ; i += 2) {
                float t = f[i]; f[i] = f[i+1]; f[i+1] = t;
            }
        }
    }
    // Obstacle nodes: full-way bounce-back (same swap)
    if (fl & LEGACY_OBSTACLE) {
        #pragma unroll
        for (int i = 1; i < NQ; i += 2) {
            float t = f[i]; f[i] = f[i+1]; f[i+1] = t;
        }
    }

    // ---- Collision (fluid only) ----
    if (fl & LEGACY_FLUID) {
        float feq[NQ], Fin[NQ];
        compute_feq(rho_n, ux, uy, uz, feq);
        zero_forcing(Fin);
        collide_srt(f, feq, Fin, d_params.omega);
    }

    stream_pull_store(k, f, fi_out);
#endif
}

}  // extern "C"
