"""CPU-only science, geometry, units, and source tests; no PyCUDA import."""

import sys
from pathlib import Path
import unittest
import numpy as np

ROOT = Path(__file__).resolve().parents[2]
REPO = ROOT
sys.path.insert(0, str(ROOT / "tests/validation"))
sys.path.insert(0, str(REPO / "src"))
from run_jet_circle_cv import (
    CX,
    CY,
    W,
    case,
    state,
    cv,
    aperture_probe,
    measured_aperture_flux,
    verify,
)


def eq(rho, u, v):
    cu = CX[:, None, None] * u + CY[:, None, None] * v
    return W[:, None, None] * rho * (1 + 3 * cu + 4.5 * cu**2 - 1.5 * (u * u + v * v))


class Host(unittest.TestCase):
    def test_descriptor_and_uniform_closed_flux(self):
        from CelerisLab.lbm.descriptors import D2Q9_EX, D2Q9_EY

        np.testing.assert_array_equal(CX, D2Q9_EX)
        np.testing.assert_array_equal(CY, D2Q9_EY)
        rho = np.full((48, 64), 1.01)
        g = eq(rho, 0.023, -0.017)
        inventory, flux = cv(g, np.ones((48, 64), bool), (0, 0), (5, 4, 25, 20), 0.0072)
        np.testing.assert_allclose(
            inventory, 320 * 1.01 * np.array([1, 0.023, -0.017]), atol=1e-12
        )
        np.testing.assert_allclose(flux, 0, atol=1e-14)

    def test_affine_velocity_independent_divergence(self):
        y, x = np.mgrid[:48, :64]
        x = (x + 0.5) * 0.5
        y = (y + 0.5) * 0.5
        a = 0.001
        g = eq(np.ones_like(x), a * x, -a * y)
        _, flux = cv(g, np.ones(x.shape, bool), (0, 0), (5, 4, 25, 20), 0.0072)
        np.testing.assert_allclose(
            flux, [0, a * a * 320 * 15, a * a * 320 * 12], atol=1e-13
        )

    def test_quadratic_viscous_flux_sign_and_units(self):
        y, x = np.mgrid[:48, :64]
        y = (y + 0.5) * 0.5
        a = 0.0001
        nu = 0.0072
        g = eq(np.ones_like(y), a * y * y, np.zeros_like(y))
        _, flux = cv(g, np.ones(y.shape, bool), (0, 0), (5, 4, 25, 20), nu)
        np.testing.assert_allclose(flux, [0, -2 * nu * a * 320, 0], atol=2e-14)

    def test_geometry_adapter_match_and_apertures(self):
        from CelerisLab.config import LBMConfig
        from CelerisLab.body.geometry.circle import CircleGeometry
        from CelerisLab.body.objects import SimObject
        from CelerisLab.body.jets import Jet, Sector, Radial, jet_link_shapes
        from CelerisLab.lbm.refinement import FixedPatch2D, prepare_fine_boundary
        from CelerisLab.lbm.descriptors import SOLID

        c = case()
        center = np.array(c["center"])
        r = c["diameter"] / 2
        cfg = LBMConfig(nx=c["nx"], ny=c["ny"], viscosity=c["nu"])
        cfg.validate()
        patch = FixedPatch2D.from_config(cfg, c["bounds"])
        obj = SimObject(0, CircleGeometry(*center, r), center=tuple(center), radius=r)
        jets = [
            Jet(
                Sector(tuple(center), a, np.pi / 12, r + 0.1),
                Radial(tuple(center), inward=False),
                "uniform",
                s * c["jet_speed"],
            )
            for a, s in [(np.pi / 2, 1), (3 * np.pi / 2, -1)]
        ]
        host = prepare_fine_boundary(cfg, patch, [obj], jets=jets)
        shape, ids = jet_link_shapes(jets, host.curved_hits)
        self.assertTrue((ids == 0).any() and (ids == 1).any() and (ids < 0).any())
        self.assertGreater(host.curved_host[2].min(), 0.010001)
        self.assertFalse(np.any(host.curved_host[-1]))
        links = CircleGeometry(*(2 * center - 0.5), 2 * r).build_cut_links(
            2 * c["nx"], 2 * c["ny"]
        )
        np.testing.assert_array_equal([v.dir for v in links], host.curved_host[1])
        np.testing.assert_allclose(
            [v.q for v in links], host.curved_host[2], atol=2e-7, rtol=0
        )
        np.testing.assert_allclose(
            np.array([[v.hit_x, v.hit_y] for v in links]) * 0.5 + 0.25,
            host.curved_hits,
            atol=2e-12,
        )
        fluid = (host.flags.reshape(patch.fine_shape) & SOLID) == 0
        g = state((9, *patch.fine_shape), c["u"])
        g[:, ~fluid] = np.nan
        for contour in c["contours"]:
            cv(g, fluid, patch.fine_origin, contour, c["nu"])
        probe = aperture_probe(fluid, patch.fine_origin, center, r)
        measured = measured_aperture_flux(g, probe)
        np.testing.assert_allclose(
            measured, 0, atol=1e-14
        )  # x-directed flow has zero flux on symmetric top/bottom arcs
        print(
            "GEOMETRY",
            dict(
                links=len(links),
                aperture_counts=[int((ids == i).sum()) for i in [0, 1]],
                q_min=float(host.curved_host[2].min()),
                q_max=float(host.curved_host[2].max()),
                fallbacks=int(np.count_nonzero(host.curved_host[-1])),
                sample_distances=[p[-1] for p in probe],
            ),
        )

    def test_diffusive_scaling(self):
        a, b = case(1), case(2)
        self.assertEqual(a["re_d"], b["re_d"])
        self.assertEqual(a["nu"], b["nu"])
        self.assertEqual(a["final_t_star"], b["final_t_star"])
        self.assertEqual(a["steps"] * 4, b["steps"])

    def test_source_pin(self):
        verify(REPO)


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