"""D2Q9 equilibrium helpers — compute_feq_d2q9 and compute_macro_from_ddf correctness.

No GPU required."""

import unittest

import numpy as np

from CelerisLab.lbm.equilibrium import compute_feq_d2q9, compute_macro_from_ddf


class TestEquilibrium(unittest.TestCase):
    """Verify D2Q9 equilibrium and macroscopic helpers."""

    def test_feq_at_rest(self):
        """Equilibrium at rho=1.0, u=0 should give w_i (weights)."""
        feq = compute_feq_d2q9(1.0, 0.0, 0.0)
        w = np.array([4/9, 1/9, 1/9, 1/9, 1/9,
                      1/36, 1/36, 1/36, 1/36], dtype=np.float32)
        np.testing.assert_allclose(feq, w, rtol=1e-6)

    def test_feq_sums_to_rho(self):
        """Sum of feq should equal rho."""
        rho, ux, uy = 1.2, 0.05, -0.02
        feq = compute_feq_d2q9(rho, ux, uy)
        self.assertAlmostEqual(float(np.sum(feq)), rho, places=6)

    def test_macro_preserves_ux_uy(self):
        """compute_macro_from_ddf(feq) should recover rho, ux, uy."""
        rho, ux, uy = 1.0, 0.1, 0.0
        feq = compute_feq_d2q9(rho, ux, uy)
        rho_out, ux_out, uy_out = compute_macro_from_ddf(feq)
        self.assertAlmostEqual(rho_out, rho, places=6)
        self.assertAlmostEqual(ux_out, ux, places=6)
        self.assertAlmostEqual(uy_out, uy, places=6)

    def test_feq_nonzero_vel(self):
        """Equilibrium at non-zero velocity should be asymmetric."""
        feq_x = compute_feq_d2q9(1.0, 0.1, 0.0)
        feq_y = compute_feq_d2q9(1.0, 0.0, 0.1)
        # x-directed flow should have f1 > f2 (right > left)
        self.assertGreater(feq_x[1], feq_x[2],
                           "Right-moving f1 should exceed left-moving f2")
        # y-directed flow should have f3 > f4 (up > down)
        self.assertGreater(feq_y[3], feq_y[4],
                           "Up-moving f3 should exceed down-moving f4")


if __name__ == "__main__":
    unittest.main()