// CelerisLab – boundary/inlet_outlet.cuh
// Inlet and outlet boundary conditions  (D2Q9).
//
// Parabolic inlet (non-equilibrium extrapolation, Zou-He style):
//   Left wall (x=0): reconstruct cx>0 populations (i=1,5,7)
//
// Pressure outlet (non-equilibrium extrapolation):
//   Right wall (x=NX-1): reconstruct cx<0 populations (i=2,6,8)
//
// New paired D2Q9 ordering:
//   cx = {0, 1,-1, 0, 0, 1,-1, 1,-1}
//   cy = {0, 0, 0, 1,-1, 1,-1,-1, 1}
// ============================================================================

#ifndef CELERIS_BOUNDARY_INLET_OUTLET_CUH
#define CELERIS_BOUNDARY_INLET_OUTLET_CUH

#ifndef INLET_PROFILE
#define INLET_PROFILE 1
#endif

#ifndef OUTLET_MODE
#define OUTLET_MODE 0
#endif

#ifndef OUTLET_BACKFLOW_CLAMP
#define OUTLET_BACKFLOW_CLAMP 1
#endif

#ifndef OUTLET_BLEND_ALPHA
#define OUTLET_BLEND_ALPHA 0.70f
#endif

// OUTLET_SRT_NEQ_DAMP and INLET_TRT_NEQ_DAMP are injected by config_method.h.
// These fallback defaults are only active if building outside the normal
// Python compile pipeline (e.g. standalone nvcc tests).
#ifndef OUTLET_SRT_NEQ_DAMP
#define OUTLET_SRT_NEQ_DAMP 0.50f
#endif

#ifndef INLET_TRT_NEQ_DAMP
#define INLET_TRT_NEQ_DAMP 0.50f
#endif

__device__ __forceinline__ float inlet_target_u(float y_coord) {
#if INLET_PROFILE == 0
    // Uniform profile: U0 is the imposed streamwise velocity everywhere on the
    // inlet fluid band.
    return U0;
#else
    // Parabolic profile on the fluid-node band y in [1, NY-2].
    //
    // U0 is treated here as the mean streamwise inlet velocity. The returned
    // peak centerline velocity is 1.5 * U0, matching the discrete Poiseuille
    // profile used throughout initialization and boundary reconstruction.
    // Keep this convention aligned with case setup and validation scripts.
    const float y_clamped = fminf((float)(NY - 2), fmaxf(1.0f, y_coord));
    const float H = fmaxf((float)(NY - 2), 1.0f);
    const float eta = (y_clamped - 0.5f) / H;  // first and last fluid rows map near 0 and 1
    const float shape = fmaxf(0.0f, 4.0f * eta * (1.0f - eta));
    return U0 * 1.5f * shape;
#endif
}

#if LATTICE_MODEL == LATTICE_D2Q9

// ---------------------------------------------------------------------------
// Parabolic inlet  (x = 0, non-equilibrium extrapolation)
//
// f, f_neb are local DDF arrays:
//   f      = populations at the boundary node (x=0)
//   f_neb  = populations at the interior neighbor (x=1)
//   y      = y-coordinate of the boundary node
//
// Reconstructs f[1], f[5], f[7]  (cx > 0 directions in new ordering).
//
// Velocity convention:
//   uniform   -> U0 is the imposed inlet velocity
//   parabolic -> U0 is the mean inlet velocity, so inlet_target_u() returns a
//                centerline peak of 1.5 * U0
//
// Reconstruction keeps the Zou-He mass closure but does not copy all three
// unknown-direction NEQ parts from the donor.
//
// Why this split form:
// - In narrow high-blockage channels, the donor diagonals f_neb[5], f_neb[7]
//   are strongly contaminated by near-wall shear and tend to inject a spurious
//   negative shift into the first interior column. Empirically this is not just
//   a corner-node artifact: the whole x=1 profile can be biased low when the
//   channel becomes very narrow.
// - Removing NEQ entirely hurts stability, so the streamwise unknown f[1] keeps
//   donor NEQ information.
// - The diagonal unknowns are instead reconstructed from local density and
//   transverse-velocity constraints. A positivity limiter is applied only if
//   the local constraints are mutually inconsistent, preferring exact rho and
//   streamwise flux over exact v_target at that node.
// ---------------------------------------------------------------------------
__device__ inline void apply_parabolic_inlet(float* __restrict__ f,
                                             const float* __restrict__ f_neb,
                                             float y_coord)
{
    // Donor macros from the first interior fluid column.
    float rho_neb, u_neb, v_neb;
    compute_rho_u(f_neb, rho_neb, u_neb, v_neb);

    // Target inlet velocity.
    const float u_target = inlet_target_u(y_coord);
    const float v_target = 0.0f;

    // Zou-He mass closure at the west boundary.
    // Known (after pull and any wall pre-repair): f[0],f[2],f[3],f[4],f[6],f[8]
    // Unknown to reconstruct:                         f[1],f[5],f[7]
    const float rho_in = (f[0] + f[3] + f[4] + 2.0f * (f[2] + f[6] + f[8]))
                       / (1.0f - u_target);

    float feq_tar[9], feq_neb[9];
    compute_feq(rho_in, u_target, v_target, feq_tar);
    compute_feq(rho_neb, u_neb, v_neb, feq_neb);

#if COLLISION_MODEL == 1
    const float beta_n = INLET_TRT_NEQ_DAMP;
#else
    const float beta_n = 1.0f;
#endif

    // Keep donor NEQ only for the streamwise incoming population. This retains
    // the stabilizing normal-flux information without feeding both diagonal
    // donor modes back into the inlet every step.
    const float f1_try = feq_tar[1] + beta_n * (f_neb[1] - feq_neb[1]);

    // Known-part density contribution.
    const float known_sum = f[0] + f[2] + f[3] + f[4] + f[6] + f[8];

    // From uy = (f3 - f4 + f5 - f6 - f7 + f8) / rho, so with v_target = 0:
    //   f5 - f7 = rho*v_target + (f4 - f3) + (f6 - f8)
    float pair_diff = rho_in * v_target + (f[4] - f[3]) + (f[6] - f[8]);

    // Density fixes f5 + f7 once f1 is chosen:
    //   f5 + f7 = rho - known_sum - f1
    // To keep both diagonals non-negative we need pair_sum >= |pair_diff|,
    // hence f1 <= rho - known_sum - |pair_diff|.
    const float f1_hi = fmaxf(0.0f, rho_in - known_sum - fabsf(pair_diff));
    const float f1 = fminf(fmaxf(f1_try, 0.0f), f1_hi);

    float pair_sum = rho_in - known_sum - f1;

    // If the local constraints are still inconsistent because of roundoff or an
    // extremely distorted incoming state, clip the transverse difference rather
    // than emitting negative diagonal populations.
    if (fabsf(pair_diff) > pair_sum) {
        pair_diff = copysignf(pair_sum, pair_diff);
    }

    f[1] = f1;
    f[5] = 0.5f * (pair_sum + pair_diff);
    f[7] = 0.5f * (pair_sum - pair_diff);
}

// ---------------------------------------------------------------------------
// Pressure outlet  (x = NX-1, non-equilibrium extrapolation)
//
// Reconstructs f[2], f[6], f[8]  (cx < 0 directions in new ordering)
// p_out = 0 (gauge pressure), uses velocity from neighbor.
// ---------------------------------------------------------------------------
__device__ inline void apply_pressure_outlet(float* __restrict__ f,
                                             const float* __restrict__ f_neb,
                                             float y_coord)
{
    (void)y_coord;

#if OUTLET_MODE == 1
    // Simple zero-gradient copy for unknown incoming directions at outlet.
    f[2] = f_neb[2];
    f[8] = f_neb[8];
    f[6] = f_neb[6];
#else

    // Prescribed-pressure outlet: keep neighbor velocity and impose reference
    // outlet density for NEQ reconstruction.
    float rho_neb, u_neb, v_neb;
    compute_rho_u(f_neb, rho_neb, u_neb, v_neb);
#if OUTLET_BACKFLOW_CLAMP
    u_neb = fmaxf(u_neb, 0.0f);
#endif
    float rho_out = RHO;

    float feq_tar[9], feq_neb[9];
    compute_feq(rho_out, u_neb, v_neb, feq_tar);
    compute_feq(rho_neb, u_neb, v_neb, feq_neb);

#if COLLISION_MODEL == 0 || COLLISION_MODEL == 1
    // SRT and TRT path: use full-population damped NEQ reconstruction at
    // outlet to suppress checkerboard and boundary-source noise.
    const float beta = OUTLET_SRT_NEQ_DAMP;
    #pragma unroll
    for (int i = 0; i < 9; i++) {
        const float fneq = f_neb[i] - feq_neb[i];
        f[i] = feq_tar[i] + beta * fneq;
    }
#elif OUTLET_MODE == 2
    const float a = OUTLET_BLEND_ALPHA;
    f[2] = a * (f_neb[2] - feq_neb[2] + feq_tar[2]) + (1.0f - a) * f_neb[2];
    f[8] = a * (f_neb[8] - feq_neb[8] + feq_tar[8]) + (1.0f - a) * f_neb[8];
    f[6] = a * (f_neb[6] - feq_neb[6] + feq_tar[6]) + (1.0f - a) * f_neb[6];
#else
    f[2] = f_neb[2] - feq_neb[2] + feq_tar[2];
    f[8] = f_neb[8] - feq_neb[8] + feq_tar[8];
    f[6] = f_neb[6] - feq_neb[6] + feq_tar[6];
#endif
#endif
}

#endif // LATTICE_D2Q9

// ============================================================================
// D3Q19 inlet / outlet  (non-equilibrium extrapolation)
//
// Parabolic inlet (x=0):  reconstruct cx>0 populations i=1,7,9,13,15
// Pressure outlet (x=NX-1): reconstruct cx<0 populations i=2,8,10,14,16
//
// Uses generic feq computation from macro.cuh to avoid hand-expanded formulas.
// ============================================================================
#if LATTICE_MODEL == LATTICE_D3Q19

__device__ inline void apply_parabolic_inlet_3d(float* __restrict__ f,
                                                const float* __restrict__ f_neb,
                                                float y_coord)
{
    // Donor macros from the first interior fluid column.
    float rho_neb, un, vn, wn;
    compute_rho_u(f_neb, rho_neb, un, vn, wn);

    // Target velocity: parabolic in y, uniform in z.
    const float u_tar = inlet_target_u(y_coord);
    const float v_tar = 0.0f;
    const float w_tar = 0.0f;

    // Zou-He mass balance at the west boundary.
    // Unknown cx>0 populations are i = 1, 7, 9, 13, 15.
    const float rho_in = (f[0] + f[3] + f[4] + f[5] + f[6] + f[11] + f[12] + f[17] + f[18]
                        + 2.0f * (f[2] + f[8] + f[10] + f[14] + f[16]))
                       / (1.0f - u_tar);

    float feq_tar[19], feq_neb[19];
    compute_feq(rho_in, u_tar, v_tar, w_tar, feq_tar);
    compute_feq(rho_neb, un, vn, wn, feq_neb);

#if COLLISION_MODEL == 1
    const float beta_n = INLET_TRT_NEQ_DAMP;
#else
    const float beta_n = 1.0f;
#endif

    // D3Q19 counterpart of the D2Q9 split strategy:
    // - keep donor NEQ on the pure streamwise incoming population f[1]
    // - retain the x-z pair only through its total incoming mass, because the
    //   present channel setup has y-walls but no z-wall BC path
    // - reconstruct the y-coupled diagonals from local rho/uy constraints to
    //   avoid feeding wall-shear contamination back into the inlet every step
    const float f1_try = feq_tar[1] + beta_n * (f_neb[1] - feq_neb[1]);
    const float zsum_try = (feq_tar[9] + feq_tar[15])
                         + beta_n * ((f_neb[9] - feq_neb[9])
                                   + (f_neb[15] - feq_neb[15]));

    const float known_sum = f[0] + f[2] + f[3] + f[4] + f[5] + f[6]
                          + f[8] + f[10] + f[11] + f[12] + f[14]
                          + f[16] + f[17] + f[18];
    const float rem_total = rho_in - known_sum;

    // uy constraint:
    // uy = (f3 - f4 + f7 - f8 + f11 - f12 + f14 - f13 + f17 - f18) / rho
    float y_diff = rho_in * v_tar
                 - (f[3] - f[4] - f[8] + f[11] - f[12] + f[14] + f[17] - f[18]);

    // uz constraint:
    // uz = (f5 - f6 + f9 - f10 + f11 - f12 + f16 - f15 + f18 - f17) / rho
    float z_diff = rho_in * w_tar
                 - (f[5] - f[6] - f[10] + f[11] - f[12] + f[16] + f[18] - f[17]);

    // Reserve enough total mass for the y-diagonal pair to satisfy positivity.
    const float f1_hi = fmaxf(0.0f, rem_total - fabsf(y_diff));
    const float f1 = fminf(fmaxf(f1_try, 0.0f), f1_hi);

    const float zsum_hi = fmaxf(0.0f, rem_total - f1 - fabsf(y_diff));
    const float z_sum = fminf(fmaxf(zsum_try, 0.0f), zsum_hi);
    if (fabsf(z_diff) > z_sum) {
        z_diff = copysignf(z_sum, z_diff);
    }

    float y_sum = rem_total - f1 - z_sum;
    y_sum = fmaxf(y_sum, 0.0f);
    if (fabsf(y_diff) > y_sum) {
        y_diff = copysignf(y_sum, y_diff);
    }

    f[1]  = f1;
    f[9]  = 0.5f * (z_sum + z_diff);
    f[15] = 0.5f * (z_sum - z_diff);
    f[7]  = 0.5f * (y_sum + y_diff);
    f[13] = 0.5f * (y_sum - y_diff);
}

__device__ inline void apply_pressure_outlet_3d(float* __restrict__ f,
                                                const float* __restrict__ f_neb,
                                                float y_coord)
{
    (void)y_coord;

#if OUTLET_MODE == 1
    // Simple zero-gradient copy for unknown incoming directions at outlet.
    f[2]  = f_neb[2];
    f[8]  = f_neb[8];
    f[10] = f_neb[10];
    f[14] = f_neb[14];
    f[16] = f_neb[16];
#else

    // Neighbor macros
    float rho_neb, un, vn, wn;
    compute_rho_u(f_neb, rho_neb, un, vn, wn);
#if OUTLET_BACKFLOW_CLAMP
    un = fmaxf(un, 0.0f);
#endif

    // Prescribed-pressure outlet: keep neighbor velocity and impose reference
    // outlet density for NEQ reconstruction.
    float rho_out = RHO;
    float feq_tar[19], feq_neb[19];
    compute_feq(rho_out, un, vn, wn, feq_tar);
    compute_feq(rho_neb, un, vn, wn, feq_neb);

    // Reconstruct cx<0 directions: i = 2, 8, 10, 14, 16
#if COLLISION_MODEL == 0 || COLLISION_MODEL == 1
    const float beta = OUTLET_SRT_NEQ_DAMP;
    #pragma unroll
    for (int i = 0; i < 19; i++) {
        const float fneq = f_neb[i] - feq_neb[i];
        f[i] = feq_tar[i] + beta * fneq;
    }
#elif OUTLET_MODE == 2
    const float a = OUTLET_BLEND_ALPHA;
    f[2]  = a * (f_neb[2]  - feq_neb[2]  + feq_tar[2])  + (1.0f - a) * f_neb[2];
    f[8]  = a * (f_neb[8]  - feq_neb[8]  + feq_tar[8])  + (1.0f - a) * f_neb[8];
    f[10] = a * (f_neb[10] - feq_neb[10] + feq_tar[10]) + (1.0f - a) * f_neb[10];
    f[14] = a * (f_neb[14] - feq_neb[14] + feq_tar[14]) + (1.0f - a) * f_neb[14];
    f[16] = a * (f_neb[16] - feq_neb[16] + feq_tar[16]) + (1.0f - a) * f_neb[16];
#else
    f[2]  = f_neb[2]  - feq_neb[2]  + feq_tar[2];
    f[8]  = f_neb[8]  - feq_neb[8]  + feq_tar[8];
    f[10] = f_neb[10] - feq_neb[10] + feq_tar[10];
    f[14] = f_neb[14] - feq_neb[14] + feq_tar[14];
    f[16] = f_neb[16] - feq_neb[16] + feq_tar[16];
#endif
#endif
}

#endif // LATTICE_D3Q19

#endif // CELERIS_BOUNDARY_INLET_OUTLET_CUH