"""Direct-current-header regression; historical candidate manifests stay frozen."""

import importlib.util
from pathlib import Path
import shutil
import sys
from tempfile import TemporaryDirectory
import unittest

ROOT = Path(__file__).resolve().parents[2]
HELPERS = ROOT / "tests/validation/les"
sys.path.insert(0, str(HELPERS))
spec = importlib.util.spec_from_file_location(
    "smagorinsky_oracle", HELPERS / "test_candidate.py"
)
oracle = importlib.util.module_from_spec(spec)
spec.loader.exec_module(oracle)
HEADER = Path("src/CelerisLab/lbm/kernels/operators/turbulence_smag.cuh")
DESCRIPTORS = Path("src/CelerisLab/lbm/kernels/core/descriptors.cuh")


class SmagorinskyNormalization(unittest.TestCase):
    def test_current_header_matches_independent_tensor_oracle(self):
        with TemporaryDirectory() as directory:
            build_dir = Path(directory)
            for dim in (2, 3):
                for cs in (0.0, 0.16):
                    # Compile directly against this checkout, not saved candidate headers.
                    closed = oracle.build(
                        ROOT.parent, ROOT.name, dim, cs, 1, 3, build_dir
                    )
                    converged = oracle.build(
                        ROOT.parent, ROOT.name, dim, cs, 0, 200, build_dir
                    )
                    for rho in (0.7, 1.0, 1.3):
                        for tau0 in (0.56, 0.8):
                            for kind, magnitude in (
                                ("equilibrium", 0.0),
                                ("shear", -0.02),
                                ("shear", 0.02),
                                ("extension", -0.02),
                                ("extension", 0.02),
                            ):
                                stress = oracle.tensor(
                                    dim, kind, magnitude * rho, 0, dim - 1
                                )
                                f, eq = oracle.state(dim, rho, stress)
                                expected = 1 / oracle.standard_tau(
                                    tau0, cs, oracle.norm(stress), rho
                                )
                                with self.subTest(
                                    dim=dim,
                                    cs=cs,
                                    rho=rho,
                                    tau=tau0,
                                    kind=kind,
                                    sign=magnitude,
                                ):
                                    self.assertAlmostEqual(
                                        closed(f, eq, rho, tau0), expected, delta=3e-6
                                    )
                                    self.assertAlmostEqual(
                                        converged(f, eq, rho, tau0),
                                        expected,
                                        delta=3e-6,
                                    )
                    if cs == 0:
                        f, eq = oracle.state(
                            dim, 1.0, oracle.tensor(dim, "equilibrium", 0.0)
                        )
                        self.assertAlmostEqual(
                            closed(f, eq, 1.0, 0.50001), 1.99, delta=3e-6
                        )
                        self.assertAlmostEqual(
                            closed(f, eq, 1.0, 200.0), 0.01, delta=3e-6
                        )

    def test_removing_sqrt_two_is_detected(self):
        with TemporaryDirectory() as directory:
            root = Path(directory)
            for relative in (HEADER, DESCRIPTORS):
                target = root / "old_convention" / relative
                target.parent.mkdir(parents=True, exist_ok=True)
                shutil.copyfile(ROOT / relative, target)
            target = root / "old_convention" / HEADER
            text = target.read_text()
            factor = "18.0f * sqrtf(2.0f) * smag_const_sq"
            self.assertEqual(text.count(factor), 2)
            target.write_text(text.replace(factor, "18.0f * smag_const_sq"))
            for dim in (2, 3):
                old = oracle.build(root, "old_convention", dim, 0.16, 1, 3, root)
                stress = oracle.tensor(dim, "shear", 0.02)
                f, eq = oracle.state(dim, 1.0, stress)
                expected = 1 / oracle.standard_tau(0.8, 0.16, oracle.norm(stress), 1.0)
                self.assertGreater(abs(old(f, eq, 1.0, 0.8) - expected), 1e-3)


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