// CelerisLab – operators/turbulence_smag.cuh
// Smagorinsky LES: compute effective omega from local non-equilibrium stress.
//
// Model:
//   nu_t = (C_s * Delta)^2 * |S|,  Delta = 1 (lattice unit)
//   omega_eff = 1 / (3*(nu + nu_t) + 0.5)
//
// |S| is estimated from non-equilibrium second-order moments using tau0.
// ============================================================================

#ifndef CELERIS_OPERATORS_TURBULENCE_SMAG_CUH
#define CELERIS_OPERATORS_TURBULENCE_SMAG_CUH

#ifndef USE_LES
#define USE_LES 0
#endif

#ifndef LES_CS
#define LES_CS 0.16f
#endif

#ifndef LES_CLOSED_FORM
#define LES_CLOSED_FORM 1
#endif

#ifndef LES_FP_ITERS
#define LES_FP_ITERS 3
#endif

#ifndef LES_FP_RELAX
#define LES_FP_RELAX 0.70f
#endif

#ifndef LES_NUT_MAX_RATIO
#define LES_NUT_MAX_RATIO 20.0f
#endif


__device__ __forceinline__ float clamp_omega(float w) {
    return fminf(OMEGA_COLLISION_MAX, fmaxf(OMEGA_COLLISION_MIN, w));
}

#if LATTICE_MODEL == LATTICE_D2Q9

__device__ __forceinline__ float compute_omega_smag(const float* __restrict__ f,
                                                    const float* __restrict__ feq,
                                                    float rho,
                                                    float omega0)
{
    const float rho_safe = fmaxf(rho, 1.0e-12f);
    const float tau0 = fmaxf(1.0f / fmaxf(omega0, 1.0e-6f), 0.500001f);

    // Πneq = Σ (ci ci)(fi-feqi)
    float pixx = 0.0f, piyy = 0.0f, pixy = 0.0f;

    #pragma unroll
    for (int i = 0; i < 9; i++) {
        const float fneq = f[i] - feq[i];
        const float cx = (float)d_cx[i];
        const float cy = (float)d_cy[i];
        pixx += fneq * cx * cx;
        piyy += fneq * cy * cy;
        pixy += fneq * cx * cy;
    }

#if LES_CLOSED_FORM
    // Closed-form Smagorinsky (no iteration):
    //   ω = 2 / (τ₀ + √(τ₀² + 18·Cs²·√Q / ρ))
    // where Q = Πneq_ij · Πneq_ij (Frobenius norm squared)
    const float Q = pixx*pixx + piyy*piyy + 2.0f*pixy*pixy;
    const float sqrtQ = sqrtf(fmaxf(Q, 0.0f));
    const float smag_const_sq = LES_CS * LES_CS;
    const float discriminant = tau0*tau0 + 18.0f * smag_const_sq * sqrtQ / rho_safe;
    const float tau_eff = 0.5f * (tau0 + sqrtf(fmaxf(discriminant, tau0*tau0)));
    return clamp_omega(1.0f / fmaxf(tau_eff, 0.500001f));
#else
    // Iterative fixed-point (original)
    const float nu0 = (tau0 - 0.5f) * (1.0f / 3.0f);
    float tau_eff = tau0;
    const float tau_max = 1.0f / fmaxf(OMEGA_COLLISION_MIN, 1.0e-6f);

    #pragma unroll
    for (int it = 0; it < LES_FP_ITERS; ++it) {
        const float denom = 2.0f * rho_safe * CS2 * fmaxf(tau_eff, 0.500001f);
        const float sxx = -pixx / denom;
        const float syy = -piyy / denom;
        const float sxy = -pixy / denom;

        // Use deviatoric strain to reduce SGS sensitivity to bulk/pressure-wave content.
        const float tr = 0.5f * (sxx + syy);
        const float sxx_dev = sxx - tr;
        const float syy_dev = syy - tr;
        const float sxy_dev = sxy;

        // |S_dev| = sqrt(2 S_dev,ij S_dev,ij)
        const float s_mag = sqrtf(2.0f * (sxx_dev*sxx_dev + syy_dev*syy_dev + 2.0f*sxy_dev*sxy_dev));

        float nut = (LES_CS * LES_CS) * s_mag;
        const float nut_cap = fmaxf(0.0f, LES_NUT_MAX_RATIO * nu0);
        nut = fminf(fmaxf(nut, 0.0f), nut_cap);

        const float tau_new = fminf(tau_max, fmaxf(0.500001f, 0.5f + 3.0f * (nu0 + nut)));
        tau_eff = LES_FP_RELAX * tau_new + (1.0f - LES_FP_RELAX) * tau_eff;
    }

    return clamp_omega(1.0f / tau_eff);
#endif // LES_CLOSED_FORM
}

#elif LATTICE_MODEL == LATTICE_D3Q19

__device__ __forceinline__ float compute_omega_smag(const float* __restrict__ f,
                                                    const float* __restrict__ feq,
                                                    float rho,
                                                    float omega0)
{
    const float rho_safe = fmaxf(rho, 1.0e-12f);
    const float tau0 = fmaxf(1.0f / fmaxf(omega0, 1.0e-6f), 0.500001f);

    // Πneq = Σ (ci ci)(fi-feqi)
    float pixx = 0.0f, piyy = 0.0f, pizz = 0.0f;
    float pixy = 0.0f, pixz = 0.0f, piyz = 0.0f;

    #pragma unroll
    for (int i = 0; i < 19; i++) {
        const float fneq = f[i] - feq[i];
        const float cx = (float)d_cx[i];
        const float cy = (float)d_cy[i];
        const float cz = (float)d_cz[i];
        pixx += fneq * cx * cx;
        piyy += fneq * cy * cy;
        pizz += fneq * cz * cz;
        pixy += fneq * cx * cy;
        pixz += fneq * cx * cz;
        piyz += fneq * cy * cz;
    }

#if LES_CLOSED_FORM
    const float Q = pixx*pixx + piyy*piyy + pizz*pizz
                  + 2.0f*(pixy*pixy + pixz*pixz + piyz*piyz);
    const float sqrtQ = sqrtf(fmaxf(Q, 0.0f));
    const float smag_const_sq = LES_CS * LES_CS;
    const float discriminant = tau0*tau0 + 18.0f * smag_const_sq * sqrtQ / rho_safe;
    const float tau_eff = 0.5f * (tau0 + sqrtf(fmaxf(discriminant, tau0*tau0)));
    return clamp_omega(1.0f / fmaxf(tau_eff, 0.500001f));
#else
    // Self-consistent tau iteration for LES/TRT coupling.
    const float nu0 = (tau0 - 0.5f) * (1.0f / 3.0f);
    float tau_eff = tau0;
    const float tau_max = 1.0f / fmaxf(OMEGA_COLLISION_MIN, 1.0e-6f);

    #pragma unroll
    for (int it = 0; it < LES_FP_ITERS; ++it) {
        const float denom = 2.0f * rho_safe * CS2 * fmaxf(tau_eff, 0.500001f);
        const float sxx = -pixx / denom;
        const float syy = -piyy / denom;
        const float szz = -pizz / denom;
        const float sxy = -pixy / denom;
        const float sxz = -pixz / denom;
        const float syz = -piyz / denom;

        // Use deviatoric strain to reduce SGS sensitivity to bulk/pressure-wave content.
        const float tr = (sxx + syy + szz) / 3.0f;
        const float sxx_dev = sxx - tr;
        const float syy_dev = syy - tr;
        const float szz_dev = szz - tr;
        const float sxy_dev = sxy;
        const float sxz_dev = sxz;
        const float syz_dev = syz;

        // |S_dev| = sqrt(2 S_dev,ij S_dev,ij)
        const float s_mag = sqrtf(2.0f * (sxx_dev*sxx_dev + syy_dev*syy_dev + szz_dev*szz_dev
                           + 2.0f*(sxy_dev*sxy_dev + sxz_dev*sxz_dev + syz_dev*syz_dev)));

        float nut = (LES_CS * LES_CS) * s_mag;
        const float nut_cap = fmaxf(0.0f, LES_NUT_MAX_RATIO * nu0);
        nut = fminf(fmaxf(nut, 0.0f), nut_cap);

        const float tau_new = fminf(tau_max, fmaxf(0.500001f, 0.5f + 3.0f * (nu0 + nut)));
        tau_eff = LES_FP_RELAX * tau_new + (1.0f - LES_FP_RELAX) * tau_eff;
    }

    return clamp_omega(1.0f / tau_eff);
#endif // LES_CLOSED_FORM
}

#endif // NQ

#endif // CELERIS_OPERATORS_TURBULENCE_SMAG_CUH
