"""CPU-only analytic checks, runnable with stdlib unittest (or pytest)."""
import unittest
import numpy as np
from research.fig2_materials.metrics import (
    common_fluid_roi_mask, field_error, motion_summary, roi_slices,
)


class MetricsTests(unittest.TestCase):
    def test_u0_normalization_not_target_energy(self):
        zero = np.zeros((2, 2)); support = np.ones((2, 2), bool)
        self.assertEqual(field_error(zero + 3, zero + 4, zero, zero, 2, support), 2.5)
        self.assertEqual(field_error(zero + 13, zero + 4, zero + 10, zero, 2, support), 2.5)
        for u0 in (0, -1, np.nan):
            with self.assertRaises(ValueError):
                field_error(zero, zero, zero, zero, u0, support)

    def test_common_mask_roi_finite_and_fluid(self):
        a = np.zeros((3, 4)); b = a.copy(); b[1, 2] = np.nan
        m = np.ones(a.shape, bool); n = m.copy(); n[1, 1] = False
        support = common_fluid_roi_mask([m, n], [a, b], (slice(1, 3), slice(1, 4)))
        self.assertEqual(int(support.sum()), 4)
        self.assertFalse(support[1, 1]); self.assertFalse(support[1, 2])
        self.assertTrue(support[2, 1])  # zero velocity is fluid
        self.assertEqual(field_error(a, a, b, a, 1, support), 0)
        with self.assertRaises(ValueError):
            common_fluid_roi_mask([m & False], [a], (slice(None), slice(None)))
        with self.assertRaises(ValueError):
            field_error(a, a, b, a, 1, m)  # fixed support must be finite

    def test_mean_field_differs_from_mean_instantaneous_error(self):
        snapshots = np.array([[[2.]], [[-2.]]]); z = np.zeros((1, 1)); mask = np.ones((1, 1), bool)
        self.assertEqual(field_error(snapshots.mean(0), z, z, z, 1, mask), 0)
        self.assertEqual(np.mean([field_error(s, z, z, z, 1, mask) for s in snapshots]), 2)

    def test_actual_geometry_and_motion_bias(self):
        roi = roi_slices((1022, 2400), D=40, xs=1600, y0=510.5, downstream_solid_extent=1350)
        self.assertEqual(roi, (slice(311, 711), slice(1360, 2161)))
        m = motion_summary([8, 10, 12])
        self.assertEqual(m['mean'], 10)
        self.assertAlmostEqual(m['centered_rms'], np.sqrt(8 / 3))
        self.assertEqual(m['half_range'], 2)
        for a in ([], [np.nan]):
            with self.assertRaises(ValueError):
                motion_summary(a)


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