"""Run explicitly on an idle GPU; automatic/manual identical-layout equivalence.

This is not independent flow-accuracy validation. CPU collection never runs it.
PYTHONPATH=src:.:tests/validation/multigrid python -m unittest discover \
    -s tests/validation/multigrid -p test_shear_refinement_example_gpu.py -v
"""
from contextlib import ExitStack
import unittest
import numpy as np
from CelerisLab.cuda.context import CudaContext
from research.shear_refinement.example import (CHILDREN, make, initialize, observe,
                                                topology, objects)


class ShearRefinementGPU(unittest.TestCase):
    def test_measured_growth_matches_explicit_manual_continuation(self):
        with CudaContext(0) as ctx, ExitStack() as owned:
            old = owned.enter_context(make(ctx))
            initialize(old)
            old.step(2)
            before = old.download()
            chosen, report = observe(before[1], topology())
            manual = (CHILDREN[0], (56, 16, 76, 40))
            self.assertEqual(chosen, manual)
            self.assertGreater(report['requested'], 0)
            self.assertEqual(report['corridor_unknown'], 0)
            self.assertEqual(report['corridor_unavailable'], 0)
            auto = owned.enter_context(old.remap_finest(chosen, objects=objects()))
            explicit = owned.enter_context(old.remap_finest(manual, objects=objects()))
            for run in (auto, explicit):
                self.assertEqual(run.hierarchy.pairs[1].bounds[0], CHILDREN[0])
                self.assertEqual(run.step_count, 2)
                run.step(2)
            for a, b, leaf in zip(auto.download(), explicit.download(), topology(manual).leaf_masks):
                np.testing.assert_array_equal(a, b)
                self.assertTrue(np.isfinite(a[:, leaf]).all())
                self.assertTrue(np.isnan(a[:, ~leaf]).all())
            for a, b in zip(old.download(), before):
                np.testing.assert_array_equal(a, b)
            self.assertEqual((old.step_count, auto.step_count, explicit.step_count), (2, 4, 4))


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