"""Standalone CPU oracle/geometry checks. Set PYTHONPATH=<checkout>/src."""

import unittest
import numpy as np
from oracle import CX, CY, U, exact, centers, initial, rectangles, metrics
from run_case import config


class OracleTests(unittest.TestCase):
    def test_phase_pressure_and_decay(self):
        u, v, p = exact(0.0, 8.0, 0.0, 32.0, 0.01)
        self.assertAlmostEqual(u, -U)
        self.assertAlmostEqual(v, 0.0)
        self.assertAlmostEqual(p, 0.0)
        self.assertAlmostEqual(exact(0.0, 0.0, 0.0, 32.0, 0.01)[2], -U * U / 2)
        a = np.array(exact(3.0, 5.0, 0.0, 32.0, 0.01))
        b = np.array(exact(3.0, 5.0, 20.0, 32.0, 0.01))
        f = np.exp(-2 * 0.01 * (2 * np.pi / 32) ** 2 * 20)
        np.testing.assert_allclose(b, a * np.array([f, f, f * f]), atol=1e-17)

    def test_navier_stokes_residual(self):
        # Independent finite-difference derivatives catch sign/phase mistakes.
        x = np.array([2.7, 8.1, 13.3])
        y = np.array([3.2, 7.4, 21.6])
        t = 10.0
        h = 0.002
        n = 32.0
        nu = 0.01
        f = lambda x, y, t: np.array(exact(x, y, t, n, nu))
        z = f(x, y, t)
        dx = (f(x + h, y, t) - f(x - h, y, t)) / (2 * h)
        dy = (f(x, y + h, t) - f(x, y - h, t)) / (2 * h)
        dt = (f(x, y, t + h) - f(x, y, t - h)) / (2 * h)
        lap = (
            f(x + h, y, t) + f(x - h, y, t) + f(x, y + h, t) + f(x, y - h, t) - 4 * z
        ) / h**2
        self.assertLess(np.max(np.abs(dx[0] + dy[1])), 1e-11)
        residual = (
            dt[:2]
            + z[0] * dx[:2]
            + z[1] * dy[:2]
            + np.array([dx[2], dy[2]])
            - nu * lap[:2]
        )
        self.assertLess(np.max(np.abs(residual)), 1e-10)

    def test_equilibrium_moments_and_metrics(self):
        n = 32
        nu = 0.01
        g = initial((n, n), 0, n, nu)
        rho = g.sum(axis=0)
        x, y = centers((n, n))
        u, v, p = exact(x, y, 0, n, nu)
        np.testing.assert_allclose(rho, 1 + 3 * p, atol=1e-15)
        np.testing.assert_allclose(np.einsum("q,qyx->yx", CX, g) / rho, u, atol=1e-15)
        np.testing.assert_allclose(np.einsum("q,qyx->yx", CY, g) / rho, v, atol=1e-15)
        m = metrics((g,), [np.ones((n, n), bool)], 0, n, nu)
        for k in ("velocity_l2", "pressure_l2", "energy_relative"):
            self.assertLess(m[k], 1e-11)
        with self.assertRaises(FloatingPointError):
            g[:, 2, 2] = np.nan
            metrics((g,), [np.ones((n, n), bool)], 0, n, nu)

    def test_geometry_cell_centers_and_scaling(self):
        from CelerisLab.lbm.hierarchy import build_hierarchy, leaf_partition
        from CelerisLab.lbm.refinement import ACTIVE

        for n in (32, 64):
            for re in (100, 1000):
                for depth in (2, 3):
                    cfg = config(n, re)
                    h = build_hierarchy(cfg, rectangles(n, depth))
                    np.testing.assert_array_equal(leaf_partition(h), np.ones((n, n)))
                    for level, cfg_l in enumerate(h.configs):
                        self.assertAlmostEqual(
                            cfg_l.viscosity, cfg.viscosity * 2**level
                        )
                        self.assertLessEqual(cfg_l.omega, cfg_l.omega_max)
                        x, y = centers(h.kinds[level].shape, level)
                        self.assertEqual(x[0, 0], 0.5 / 2**level)
                        self.assertEqual(x[0, -1], n - 0.5 / 2**level)
                    masks = [k == ACTIVE for k in h.kinds]
                    g = [
                        initial(k.shape, l, n, cfg.viscosity)
                        for l, k in enumerate(h.kinds)
                    ]
                    result = metrics(g, masks, 0, n, cfg.viscosity)
                    self.assertLess(result["velocity_l2"], 1e-12)
                    self.assertLess(result["pressure_l2"], 1e-10)
                    self.assertLess(result["energy_relative"], 0.002)

    def test_complete_assessment_and_corruptions(self):
        from copy import deepcopy
        from assess import assess
        from oracle import TIMES

        records = []
        for re in (100, 1000):
            cases = [(n, 1) for n in (32, 64, 128, 256)]
            cases += [(n, d) for n in (32, 64) for d in (2, 3)]
            for n, d in cases:
                rows = [
                    dict(
                        tstar=t,
                        velocity_l2=0.0,
                        pressure_l2=0.0,
                        energy_relative=0.0,
                        mass_relative=0.0,
                    )
                    for t in TIMES
                ]
                records.append(
                    dict(
                        re_L=re,
                        n=n,
                        depth=d,
                        U=U,
                        viscosity=U * n / re,
                        source=dict(revision="synthetic", dirty=False),
                        completed=True,
                        rows=rows,
                    )
                )
        self.assertTrue(assess(records)["passed"])
        bad = deepcopy(records)
        bad[-1]["rows"][-1]["pressure_l2"] = 0.5
        failed = assess(bad)
        self.assertFalse(failed["passed"])
        self.assertEqual(len(failed["comparisons"]), 24)
        partial = assess(records[0:1] + records[4:5])
        self.assertFalse(partial["passed"])
        self.assertEqual(len(partial["comparisons"]), 3)
        self.assertIsNone(partial["comparisons"][0]["uniform_finest"])

        bad = deepcopy(records)
        bad[-1]["rows"][-1]["pressure_l2"] = float("nan")
        self.assertFalse(assess(bad)["passed"])
        bad = deepcopy(records)
        bad[-1]["source"]["dirty"] = True
        self.assertFalse(assess(bad)["passed"])
        bad = deepcopy(records)
        bad[-1]["rows"][-1]["tstar"] = 0.49
        self.assertFalse(assess(bad)["passed"])
        with self.assertRaises(ValueError):
            assess(records + [records[0]])

    def test_workload(self):
        from run_case import workload
        from CelerisLab.lbm.hierarchy import build_hierarchy
        from CelerisLab.lbm.refinement import ACTIVE

        cases = [(n, 1) for n in (32, 64, 128, 256)]
        cases += [(n, d) for n in (32, 64) for d in (2, 3)]
        costs = []
        for n, d in cases:
            h = build_hierarchy(config(n, 100), [] if d == 1 else rectangles(n, d))
            costs.append(workload([k == ACTIVE for k in h.kinds], 16 * n))
        self.assertEqual(sum(c["active_collision_updates"] for c in costs), 384270336)
        self.assertEqual(sum(c["root_intervals"] for c in costs), 10752)
        self.assertEqual(costs[-1]["active_collision_updates"], 50987008)

    def test_incomplete_assessment_fails(self):
        from assess import assess

        self.assertFalse(assess([])["passed"])
        self.assertEqual(len(assess([])["missing"]), 16)


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