"""Manufactured references for continuation instrumentation, no CUDA."""

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

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "validation"))
from jet_cv_measurements import (
    sample_steps,
    factors,
    ray_geometry,
    reconstruct_aperture,
    contour_terms,
    LINEAR,
    QUADRATIC,
    OFFSETS,
    MeasurementRecorder,
)
from run_jet_circle_cv import (
    case,
    CX,
    CY,
    W,
    cv,
    aperture_probe,
    measured_aperture_flux,
)
from analyze_jet_cv_continuation import integrate, analyze


class Measurements(unittest.TestCase):
    def test_common_physical_schedule(self):
        a = sample_steps(case(1, steps=3200))
        b = sample_steps(case(2, steps=3200))
        np.testing.assert_array_equal(a, b / 4)
        self.assertEqual(len(a), 876)
        self.assertEqual(a[-1], 3200)
        self.assertTrue(all(t in a for t in range(0, 3201, 100)))
        self.assertEqual(len(a[a <= 100]), 101)

    def test_reference_factors(self):
        f = factors(2)
        self.assertEqual(f["mass"], 0.25)
        self.assertEqual(f["impulse"], 0.5)
        self.assertEqual(f["force"], 2)
        self.assertEqual(f["mass_rate"], 1)

    def test_extrapolation_polynomials(self):
        self.assertAlmostEqual(LINEAR @ (2 + 3 * OFFSETS), 2)
        self.assertAlmostEqual(QUADRATIC @ (2 + 3 * OFFSETS + 4 * OFFSETS**2), 2)

    def test_independent_linear_field_volume_and_mass(self):
        h = 0.5
        y, x = np.mgrid[:64, :64]
        x = (x + 0.5) * h
        y = (y + 0.5) * h
        rho = 1 + 0.001 * x
        vel = np.array([0.01 + 0.0001 * x, -0.02 + 0.0002 * y])
        center = np.array([16.0, 16.0])
        r = 5
        fluid = (x - center[0]) ** 2 + (y - center[1]) ** 2 > r * r
        geo = ray_geometry(fluid, (0.0, 0.0), center, r)
        ray, est = reconstruct_aperture(rho, vel, geo)
        points = center + r * geo["normal"]
        rr = 1 + 0.001 * points[..., 0]
        uv = np.stack(
            (0.01 + 0.0001 * points[..., 0], -0.02 + 0.0002 * points[..., 1]), axis=-1
        )
        un = np.sum(uv * geo["normal"], axis=-1)
        exact = np.stack(
            (np.sum(un * geo["ds"], axis=1), np.sum(rr * un * geo["ds"], axis=1)),
            axis=-1,
        )
        np.testing.assert_allclose(est, np.broadcast_to(exact, est.shape), atol=1e-15)
        self.assertGreater(np.max(abs(est[:, :, 0] - est[:, :, 1])), 1e-4)

    def test_split_contour_stress(self):
        y, x = np.mgrid[:48, :64]
        y = (y + 0.5) * 0.5
        rho = np.ones_like(y)
        v = np.array([1e-4 * y * y, np.zeros_like(y)])
        terms = contour_terms(rho, v, (0, 0), (5, 4, 25, 20), 0.0072)
        np.testing.assert_allclose(terms[:2], 0, atol=1e-15)
        np.testing.assert_allclose(terms[2], [-2 * 0.0072 * 1e-4 * 320, 0], atol=1e-15)

    def test_recorder_snapshot_and_reconstructible_arrays(self):
        h = 0.5
        y, x = np.mgrid[:64, :64]
        x = (x + 0.5) * h
        y = (y + 0.5) * h
        fluid = (x - 16) ** 2 + (y - 16) ** 2 > 25
        fine_leaf = fluid.copy()
        fine_leaf[:, :4] = False
        root_leaf = np.zeros((32, 32), bool)
        root_leaf[:, :2] = True

        def equilibrium(ny, nx):
            cu = CX[:, None, None] * 0.01
            return np.broadcast_to(
                W[:, None, None] * (1 + 3 * cu + 4.5 * cu * cu - 1.5 * 0.01**2),
                (9, ny, nx),
            ).copy()

        ug = equilibrium(64, 64)
        levels = [equilibrium(32, 32), ug.copy()]
        c = dict(
            scale=1,
            center=[16, 16],
            diameter=10,
            u=0.01,
            nu=0.0072,
            contours=[[9, 9, 23, 23], [7, 7, 25, 25]],
        )
        data = {
            name: [cv(ug, mask, (0, 0), r, c["nu"]) for r in c["contours"]]
            for name, mask in [("uniform_adapter", fluid), ("amr_api", fine_leaf)]
        }
        with tempfile.TemporaryDirectory() as tmp:
            rec = MeasurementRecorder(
                c, tmp, fluid, [root_leaf, fine_leaf], (0, 0), CX, CY
            )
            rec.capture(0, ug, levels, data, {k: np.zeros(2) for k in data}, 0.0)
            rec.records[-1]["raw_cumulative"] = np.zeros((2, 5, 4))
            rec.flush()
            a = np.load(Path(tmp) / "samples_s1.npz")
            self.assertEqual(a["arc_ray_values"].shape, (1, 2, 3, 2, 32, 3))
            np.testing.assert_allclose(
                a["global_mass_reference"][0, 0], a["global_mass_reference"][0, 1]
            )
            snap = np.load(Path(tmp) / "fields_s1_tref0.npz")
            np.testing.assert_allclose(
                snap["uniform_velocity_reference"][0][fluid], 0.01, rtol=1e-7
            )

    def test_offline_mass_force_normalization(self):
        import json

        c = case(2, steps=3200)
        s = 2
        t = np.arange(0, 3201, 100)
        n = len(t)
        inv = np.zeros((n, 2, 2, 3))
        inv[:, :, :, 0] = 100 * s * s
        inv[:, :, :, 1] = -t[:, None, None] * 0.002 * s
        raw = np.zeros((n, 2, 5, 4))
        raw[:, :, 2, 0] = t[:, None] * 0.002 * s
        with tempfile.TemporaryDirectory() as tmp:
            path = Path(tmp)
            (path / "samples_s2.json").write_text(
                json.dumps(dict(case=c, sample_file="samples_s2.npz"))
            )
            np.savez(
                path / "samples_s2.npz",
                t_reference=t,
                cv_inventory=inv,
                cv_flux=np.zeros_like(inv),
                raw_cumulative=raw,
                arc_estimates=np.zeros((n, 2, 2, 2, 2)),
                global_mass_reference=np.ones((n, 2)) * 100,
                level_stats=np.broadcast_to([1, 1, 0, 0.03, 100], (n, 3, 5)),
                velocity_error=np.zeros(n),
            )
            result = analyze(path)
            self.assertEqual(len(result["windows"]), 32)
            np.testing.assert_allclose(
                result["windows"][0]["normalized_cv_residual"], 0, atol=1e-14
            )
            f0 = 0.5 * 0.03**2 * 24.084
            np.testing.assert_allclose(
                result["windows"][0]["raw_reaction_coefficient"], [[0.002 / f0, 0]] * 2
            )

    def test_nonuniform_trapezoidal_time_units(self):
        t = np.array([0.0, 1.0, 2.0, 4.0, 8.0])
        y = np.stack((3 * t + 2, 2 * t - 7), axis=-1)
        np.testing.assert_allclose(integrate(t, y), [112, 8], atol=1e-15)


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