// CelerisLab – operators/collision_trt.cuh
// TRT (Two-Relaxation-Time) collision operator.
//
// Symmetric part   uses ω⁺ (= omega, transport viscosity)
// Antisymmetric    uses ω⁻ (linked by magic parameter Λ = 3/16)
//   ω⁻ = 1 / (Λ / (1/ω⁺ - 0.5) + 0.5)
//
// Paired direction ordering makes TRT natural:
//   f_s[i] = 0.5*(f[i] + f[opp(i)])    symmetric
//   f_a[i] = 0.5*(f[i] - f[opp(i)])    antisymmetric
//
// Forcing contract:
//   Fin[i] is the raw Guo source term without the (1 - ω/2) prefactor.
//   This operator applies the prefactor internally, matching collide_srt().
// ============================================================================

#ifndef CELERIS_OPERATORS_COLLISION_TRT_CUH
#define CELERIS_OPERATORS_COLLISION_TRT_CUH

#ifndef TRT_MAGIC_PARAM
#define TRT_MAGIC_PARAM (0.1875f)
#endif

__device__ __forceinline__ float compute_omega_minus(float omega_plus) {
    return 1.0f / (TRT_MAGIC_PARAM / (1.0f / omega_plus - 0.5f) + 0.5f);
}

__device__ __forceinline__ void collide_trt(float* __restrict__ f,
                                            const float* __restrict__ feq,
                                            const float* __restrict__ Fin,
                                            float omega)
{
    const float wp = omega;
    const float wm = compute_omega_minus(wp);
    const float c_tau = 1.0f - 0.5f * wp;

    f[0] = f[0] - wp * (f[0] - feq[0]) + c_tau * Fin[0];

    #pragma unroll
    for (int i = 1; i < NQ; i += 2) {
        const int ib = i + 1;

        const float fi = f[i];
        const float fb = f[ib];

        const float feqi = feq[i];
        const float feqb = feq[ib];

        const float delta_s = 0.5f * ((fi - feqi) + (fb - feqb));
        const float delta_a = 0.5f * ((fi - feqi) - (fb - feqb));

        f[i]  = fi  - wp * delta_s - wm * delta_a + c_tau * Fin[i];
        f[ib] = fb  - wp * delta_s + wm * delta_a + c_tau * Fin[ib];
    }
}

__device__ __forceinline__ void collide_trt_no_force(float* __restrict__ f,
                                                     const float* __restrict__ feq,
                                                     float omega)
{
    const float wp = omega;
    const float wm = compute_omega_minus(wp);

    f[0] = f[0] - wp * (f[0] - feq[0]);

    #pragma unroll
    for (int i = 1; i < NQ; i += 2) {
        const int ib = i + 1;
        const float fi = f[i], fb = f[ib];
        const float feqi = feq[i], feqb = feq[ib];
        const float delta_s = 0.5f * ((fi - feqi) + (fb - feqb));
        const float delta_a = 0.5f * ((fi - feqi) - (fb - feqb));
        f[i]  = fi  - wp * delta_s - wm * delta_a;
        f[ib] = fb  - wp * delta_s + wm * delta_a;
    }
}

#endif // CELERIS_OPERATORS_COLLISION_TRT_CUH
