// CelerisLab – operators/collision_mrt.cuh
// MRT (Multiple-Relaxation-Time) collision operator.
//
// D2Q9 Lallemand-Luo (Phys. Rev. E 61, 2000) with NEW paired ordering:
//   cx = {0, 1,-1, 0, 0, 1,-1, 1,-1}
//   cy = {0, 0, 0, 1,-1, 1,-1,-1, 1}
//
// Moment space:
//   m[0] = ρ     (conserved,  s₀ = 0 or 1.0 overwrite)
//   m[1] = e     (energy,     s₁ = s_e)
//   m[2] = ε     (energy²,   s₂ = s_eps)
//   m[3] = jx    (conserved,  s₃ = 0 or 1.0 overwrite)
//   m[4] = qx    (energy flux x, s₄ = s_q)
//   m[5] = jy    (conserved,  s₅ = 0 or 1.0 overwrite)
//   m[6] = qy    (energy flux y, s₆ = s_q)
//   m[7] = pxx   (stress,     s₇ = s_nu = ω = 1/(3ν + 0.5))
//   m[8] = pxy   (stress,     s₈ = s_nu)
//
// 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_MRT_CUH
#define CELERIS_OPERATORS_COLLISION_MRT_CUH

#if LATTICE_MODEL == LATTICE_D2Q9

__device__ __forceinline__ void collide_mrt(float* __restrict__ g,
                                            float rho, float ux, float uy,
                                            const float* __restrict__ Fin,
                                            float omega)
{
    const float s_rho = 0.0f;
    const float s_e   = 1.2f;
    const float s_eps = 1.2f;
    const float s_jx  = 0.0f;
    const float s_q   = 1.2f;
    const float s_jy  = 0.0f;
    const float s_nu  = omega;
    const float c_tau = 1.0f - 0.5f * omega;

    const float p = rho * (1.0f / 3.0f);

    float m[9];
    m[0] =    g[0] +   g[1] +   g[2] +   g[3] +   g[4] +   g[5] +   g[6] +   g[7] +   g[8];
    m[1] = -4*g[0] -   g[1] -   g[2] -   g[3] -   g[4] + 2*g[5] + 2*g[6] + 2*g[7] + 2*g[8];
    m[2] =  4*g[0] - 2*g[1] - 2*g[2] - 2*g[3] - 2*g[4] +   g[5] +   g[6] +   g[7] +   g[8];
    m[3] =            g[1] -   g[2]                      +   g[5] -   g[6] +   g[7] -   g[8];
    m[4] =         -2*g[1] + 2*g[2]                      +   g[5] -   g[6] +   g[7] -   g[8];
    m[5] =                            g[3] -   g[4]      +   g[5] -   g[6] -   g[7] +   g[8];
    m[6] =                         -2*g[3] + 2*g[4]      +   g[5] -   g[6] -   g[7] +   g[8];
    m[7] =            g[1] +   g[2] -  g[3] -   g[4];
    m[8] =                                                    g[5] +   g[6] -   g[7] -   g[8];

    const float u2 = ux * ux + uy * uy;
    float meq[9];
    meq[0] = 3.0f * p;
    meq[1] = -6.0f * p + 3.0f * rho * u2;
    meq[2] =  3.0f * p - 3.0f * rho * u2;
    meq[3] = rho * ux;
    meq[4] = -rho * ux;
    meq[5] = rho * uy;
    meq[6] = -rho * uy;
    meq[7] = rho * (ux * ux - uy * uy);
    meq[8] = rho * ux * uy;

    float dm[9];
    dm[0] = s_rho * (meq[0] - m[0]);
    dm[1] = s_e   * (meq[1] - m[1]);
    dm[2] = s_eps * (meq[2] - m[2]);
    dm[3] = s_jx  * (meq[3] - m[3]);
    dm[4] = s_q   * (meq[4] - m[4]);
    dm[5] = s_jy  * (meq[5] - m[5]);
    dm[6] = s_q   * (meq[6] - m[6]);
    dm[7] = s_nu  * (meq[7] - m[7]);
    dm[8] = s_nu  * (meq[8] - m[8]);

    g[0] += (    dm[0] -    dm[1] +    dm[2]                                              ) / 9.0f;
    g[1] += (4 * dm[0] -    dm[1] - 2* dm[2] + 6* dm[3] - 6* dm[4]               + 9*dm[7]) / 36.0f;
    g[2] += (4 * dm[0] -    dm[1] - 2* dm[2] - 6* dm[3] + 6* dm[4]               + 9*dm[7]) / 36.0f;
    g[3] += (4 * dm[0] -    dm[1] - 2* dm[2]                        + 6*dm[5] - 6*dm[6] - 9*dm[7]) / 36.0f;
    g[4] += (4 * dm[0] -    dm[1] - 2* dm[2]                        - 6*dm[5] + 6*dm[6] - 9*dm[7]) / 36.0f;
    g[5] += (4 * dm[0] + 2* dm[1] +    dm[2] + 6* dm[3] + 3* dm[4] + 6*dm[5] + 3*dm[6] + 9*dm[8]) / 36.0f;
    g[6] += (4 * dm[0] + 2* dm[1] +    dm[2] - 6* dm[3] - 3* dm[4] - 6*dm[5] - 3*dm[6] + 9*dm[8]) / 36.0f;
    g[7] += (4 * dm[0] + 2* dm[1] +    dm[2] + 6* dm[3] + 3* dm[4] - 6*dm[5] - 3*dm[6] - 9*dm[8]) / 36.0f;
    g[8] += (4 * dm[0] + 2* dm[1] +    dm[2] - 6* dm[3] - 3* dm[4] + 6*dm[5] + 3*dm[6] - 9*dm[8]) / 36.0f;

    #pragma unroll
    for (int i = 0; i < 9; i++) {
        g[i] += c_tau * Fin[i];
    }
}

__device__ __forceinline__ void collide_mrt_no_force(float* __restrict__ g,
                                                     float rho, float ux, float uy,
                                                     float omega)
{
    float Fin[9] = {0};
    collide_mrt(g, rho, ux, uy, Fin, omega);
}

#elif LATTICE_MODEL == LATTICE_D3Q19

__device__ __forceinline__ void collide_mrt(float* __restrict__ g,
                                            float rho, float ux, float uy, float uz,
                                            const float* __restrict__ Fin,
                                            float omega)
{
    float feq[19];
    compute_feq(rho, ux, uy, uz, feq);

    const float s_nu   = omega;
    const float s_bulk = 1.2f;
    const float s_high = 1.4f;
    const float c_tau  = 1.0f - 0.5f * omega;

    const float one_minus_s_nu   = 1.0f - s_nu;
    const float one_minus_s_bulk = 1.0f - s_bulk;
    const float one_minus_s_high = 1.0f - s_high;

    float neq[19];
    #pragma unroll
    for (int i = 0; i < 19; i++) {
        neq[i] = g[i] - feq[i];
    }

    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 ci_x = (float)d_cx[i];
        const float ci_y = (float)d_cy[i];
        const float ci_z = (float)d_cz[i];
        const float fneq = neq[i];
        pixx += fneq * ci_x * ci_x;
        piyy += fneq * ci_y * ci_y;
        pizz += fneq * ci_z * ci_z;
        pixy += fneq * ci_x * ci_y;
        pixz += fneq * ci_x * ci_z;
        piyz += fneq * ci_y * ci_z;
    }

    const float tr_pi      = pixx + piyy + pizz;
    const float tr_pi_thrd = tr_pi * (1.0f / 3.0f);

    const float dev_xx = pixx - tr_pi_thrd;
    const float dev_yy = piyy - tr_pi_thrd;
    const float dev_zz = pizz - tr_pi_thrd;
    const float dev_xy = pixy;
    const float dev_xz = pixz;
    const float dev_yz = piyz;

    const float proj_pref = 4.5f;

    #pragma unroll
    for (int i = 0; i < 19; i++) {
        const float ci_x = (float)d_cx[i];
        const float ci_y = (float)d_cy[i];
        const float ci_z = (float)d_cz[i];

        const float q_xx = ci_x * ci_x - CS2;
        const float q_yy = ci_y * ci_y - CS2;
        const float q_zz = ci_z * ci_z - CS2;
        const float q_xy = ci_x * ci_y;
        const float q_xz = ci_x * ci_z;
        const float q_yz = ci_y * ci_z;

        const float q_dot_dev =
            q_xx * dev_xx + q_yy * dev_yy + q_zz * dev_zz
            + 2.0f * (q_xy * dev_xy + q_xz * dev_xz + q_yz * dev_yz);

        const float q_trace = q_xx + q_yy + q_zz;
        const float q_dot_iso = q_trace * tr_pi_thrd;

        const float neq_dev = d_w[i] * proj_pref * q_dot_dev;
        const float neq_iso = d_w[i] * proj_pref * q_dot_iso;
        const float neq_h   = neq[i] - neq_dev - neq_iso;

        g[i] = feq[i]
             + one_minus_s_nu   * neq_dev
             + one_minus_s_bulk * neq_iso
             + one_minus_s_high * neq_h
             + c_tau * Fin[i];
    }
}

#endif  // NQ

#endif  // CELERIS_OPERATORS_COLLISION_MRT_CUH
