"""CPU-only research-example contracts; no solver/GPU import required."""
import unittest
import numpy as np
from research.shear_refinement.boundary_shear_tagging import tag_cells, strain_rate
from research.shear_refinement.example import (CHILDREN, grow_rectangle, topology,
                                                observe, initial_fields)


class ShearRefinementExample(unittest.TestCase):
    def test_affine_shear_physical_spacing_and_rotation(self):
        y, x = np.indices((12, 20), dtype=float)
        x, y = 7.+.5*(x+.5), 3.+.5*(y+.5)
        valid = np.ones(x.shape, bool)
        np.testing.assert_allclose(strain_rate(.003*y, 0*x, valid, .5), .003, atol=1e-17)
        np.testing.assert_allclose(strain_rate(-.003*y, .003*x, valid, .5), 0., atol=1e-17)

    def test_holes_are_unknown_not_walls_or_zero(self):
        valid = np.zeros((7, 7), bool); valid[3, 3] = True
        u = np.full(valid.shape, np.nan); u[3, 3] = .1
        result = tag_cells(u, u, np.zeros_like(valid), valid=valid, owned=valid,
                           dx_level=.5, origin=(7., 3.), U_ref=.05, L_ref=1.,
                           shear_on=.01, shear_off=0., band_on=0., band_off=0., n_buffer=0)
        self.assertTrue(result.unevaluable[3, 3])
        self.assertTrue(np.isnan(result.strain_star[3, 3]))
        self.assertFalse(result.tagged.any())
        self.assertTrue(np.isinf(result.wall_distance).all())

    def test_masked_hole_never_supplies_a_gradient(self):
        y, x = np.indices((9, 9), dtype=float)
        valid = np.ones(x.shape, bool); valid[:, 4] = False
        ux = .003*(5.+.5*y); uy = np.zeros_like(ux)
        ux[~valid] = np.inf; uy[~valid] = np.nan
        strain = strain_rate(ux, uy, valid, .5)
        np.testing.assert_allclose(strain[valid], .003, atol=1e-17)
        self.assertTrue(np.isnan(strain[~valid]).all())

    def test_predictable_growth_and_input_unchanged(self):
        mask = np.zeros((64, 96), bool); mask[22, 74] = True
        before = mask.copy()
        self.assertEqual(grow_rectangle(CHILDREN[1], mask), (56, 16, 76, 40))
        np.testing.assert_array_equal(mask, before)
        mask[:] = False
        self.assertEqual(grow_rectangle(CHILDREN[1], mask), CHILDREN[1])

    def test_budget_and_geometry_rejection(self):
        mask = np.zeros((64, 96), bool); mask[22, 74] = True
        with self.assertRaisesRegex(ValueError, 'area cap'):
            grow_rectangle(CHILDREN[1], mask, area_cap=479)
        with self.assertRaises(ValueError):
            grow_rectangle((57, 16, 72, 40), mask)
        mask = np.zeros((9, 9), bool); mask[8, 8] = True
        with self.assertRaisesRegex(ValueError, 'frame'):
            grow_rectangle((0, 0, 4, 4), mask)
        with self.assertRaises(ValueError):
            topology((CHILDREN[0], (56, 16, 88, 40)))

    def test_uniform_velocity_observation_does_not_grow(self):
        from CelerisLab.lbm.equilibrium import compute_feq_d2q9
        mesh = topology()
        g = np.broadcast_to(compute_feq_d2q9(1., .01, 0.)[:, None, None],
                            (9, *mesh.leaf_masks[1].shape)).copy()
        g[:, ~mesh.leaf_masks[1]] = np.nan
        rectangles, report = observe(g, mesh)
        self.assertEqual(rectangles, CHILDREN)
        self.assertEqual(report['requested'], 0)

    def test_observation_growth_and_low_strain_no_change(self):
        mesh = topology()
        g = initial_fields(mesh)[1]
        g[:, ~mesh.leaf_masks[1]] = np.nan
        before = g.copy()
        rectangles, report = observe(g, mesh)
        self.assertEqual(rectangles, (CHILDREN[0], (56, 16, 76, 40)))
        self.assertEqual(report['requested'], 64)
        self.assertEqual(report['corridor_unknown'], 0)
        self.assertEqual(report['corridor_unavailable'], 0)
        self.assertGreater(report['out_of_scope'], 0)
        np.testing.assert_array_equal(g, before)
        rectangles, report = observe(g, mesh, shear_on=1.)
        self.assertEqual(rectangles, CHILDREN)
        self.assertEqual(report['requested'], 0)


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