import unittest

import numpy as np

from mapgen.grid import Grid
from tests.helpers import small_grid


class GridTest(unittest.TestCase):
    def test_counts_and_pentagons(self):
        g = small_grid(2)
        self.assertEqual(g.n, 5882)
        self.assertTrue(np.all(np.diff(g.ids.astype(np.int64)) > 0))
        self.assertEqual(int(np.sum(g.counts == 5)), 12)
        self.assertEqual(int(np.sum(g.counts == 6)), g.n - 12)

    def test_neighbours_symmetric(self):
        g = small_grid(2)
        pairs = set(zip(g.src.tolist(), g.dst.tolist()))
        self.assertTrue(all((b, a) in pairs for a, b in pairs))

    def test_area_sums_to_sphere(self):
        g = small_grid(2)
        self.assertAlmostEqual(g.area_km2.sum() / (4 * np.pi * 12742.0**2), 1.0, places=6)

    def test_spacing_and_tangents(self):
        g = small_grid(2)
        self.assertTrue(500 < g.spacing_km < 800)  # res 2 ≈ 632 km
        np.testing.assert_allclose(np.linalg.norm(g.edge_tangents, axis=1), 1.0)
        np.testing.assert_allclose(np.sum(g.edge_tangents * g.xyz[g.src], axis=1), 0.0, atol=1e-12)

    def test_roundtrip_arrays_and_lookup(self):
        g = small_grid(2)
        h = Grid.from_arrays(g.to_arrays(), g.res, g.radius_km)
        np.testing.assert_array_equal(h.nbr_idx, g.nbr_idx)
        i = g.cell_index(10.0, 20.0)
        self.assertLess(abs(g.lat[i] - 10.0), 5.0)


class CellParentsTest(unittest.TestCase):
    def test_same_as_h3_cell_to_parent(self):
        import h3.api.basic_int as h3
        import numpy as np
        from mapgen.grid import cell_parents
        rng = np.random.default_rng(5)
        for res in (1, 3, 6, 9):
            cells = [h3.latlng_to_cell(float(la), float(lo), res + int(k))
                     for la, lo, k in zip(rng.uniform(-90, 90, 300), rng.uniform(-180, 180, 300), rng.integers(0, 4, 300))]
            pent = h3.get_pentagons(max(res - 2, 0))[:2] + [h3.latlng_to_cell(10.0, 20.0, max(res - 2, 0))]
            cells += [c for p in pent for c in h3.cell_to_children(p, res)]                      # incl. pentagons
            for pr in range(0, res + 1):
                want = np.array([h3.cell_to_parent(c, pr) for c in cells], np.uint64)
                self.assertTrue(np.array_equal(cell_parents(cells, pr), want), (res, pr))
        with self.assertRaises(ValueError):
            cell_parents([h3.latlng_to_cell(0.0, 0.0, 2)], 3)


class EdgeBlocksTest(unittest.TestCase):
    def test_blocked_edge_fields_same_as_whole(self):
        from unittest import mock
        from mapgen import grid as GR
        from mapgen import plates as PL
        from mapgen.sphere import tangent_dir
        g = small_grid(3)
        xyz, src, dst = g.xyz, g.src, g.dst
        want_t = tangent_dir(xyz[src], xyz[dst])                      # the whole-array formulas, written out
        want_km = g.radius_km * np.arccos(np.clip(np.sum(xyz[src] * xyz[dst], axis=1), -1.0, 1.0))
        vel = np.random.default_rng(2).normal(size=(g.n, 3))
        rel = vel[dst] - vel[src]
        along = np.sum(rel * want_t, axis=1)
        want_tang = np.linalg.norm(rel - along[:, None] * want_t, axis=1)
        with mock.patch.object(GR, "EDGE_BLOCK", 997):                 # many uneven blocks
            g2 = small_grid(3)
            self.assertTrue(np.array_equal(g2.edge_tangents, want_t))
            self.assertTrue(np.array_equal(g2.edge_km, want_km))
            conv, tang = PL.edge_convergence(g2, vel)
        self.assertTrue(np.array_equal(conv, -along))
        self.assertTrue(np.array_equal(tang, want_tang))
        self.assertGreater(len(GR.edge_blocks(len(dst), 997)), 3)
        with mock.patch.object(GR, "EDGE_BLOCK", 997):
            for idx in (src, dst):
                self.assertTrue(np.array_equal(GR.rowdot_at(vel, idx, want_t), np.sum(vel[idx] * want_t, axis=1)))
                self.assertTrue(np.array_equal(GR.rowdot_at(vel, idx, want_t), np.sum(want_t * vel[idx], axis=1)))