"""Host reduction contract for the optional observation-window exchange ledger."""
import unittest

import numpy as np

from CelerisLab.lbm.raw_exchange import summarize_exchange


def reduce(rows, window, duration=None, h=.5, time_end=10):
    return summarize_exchange(rows, np.array([0, 0, 1]), np.array([False, True, False]),
                              body_id=0, level=1, h=h, time_end=time_end, window=window,
                              reference_center=(0., 0.), duration=duration)


class ExchangeWindowReduction(unittest.TestCase):
    def test_window_duration_scaling(self):
        rows = np.array([[1., 2., 3., 4.], [.5, -1., 0., 1.], [9., 9., 9., 9.]])
        out = reduce(rows, "window", duration=4)
        np.testing.assert_array_equal(out["total"]["impulse"], [1.5, 1.])
        np.testing.assert_array_equal(out["total"]["force"], [1.5 / 4, 1. / 4])
        self.assertEqual(out["total"]["torque"], 3. / 4)
        self.assertEqual(out["total"]["mass_rate"], 5. / 4)
        np.testing.assert_array_equal(out["aperture"]["impulse"], [.5, -1.])
        self.assertEqual((out["time_start"], out["time_end"]), (6., 10))
        self.assertEqual(out["sample_count"], 8)  # duration / h fine substeps
        self.assertEqual(out["window"], "window")

    def test_existing_windows_unchanged(self):
        rows = np.array([[1., 2., 3., 4.], [0., 0., 0., 0.], [0., 0., 0., 0.]])
        base = reduce(rows, "last_base_step")
        self.assertEqual((base["time_start"], base["sample_count"]), (9., 2))
        np.testing.assert_array_equal(base["total"]["force"], [1., 2.])
        sub = reduce(rows, "last_substep")
        self.assertEqual(sub["sample_count"], 1)

    def test_duration_validation(self):
        rows = np.zeros((3, 4))
        for bad in (None, 0, -1, 1.5, True):
            with self.subTest(duration=bad), self.assertRaises(ValueError):
                reduce(rows, "window", duration=bad)
        with self.assertRaises(ValueError):
            reduce(rows, "last_base_step", duration=2)


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