import unittest

import numpy as np

from mapgen import graph as G
from mapgen.sphere import latlon_to_xyz, east_north
from tests.helpers import small_grid


class GraphTest(unittest.TestCase):
    def setUp(self):
        self.g = small_grid(2)

    def test_mean_max_min_diffuse(self):
        g = self.g
        f = np.zeros(g.n)
        f[0] = 6.0
        m = G.nbr_mean(g, f)
        self.assertAlmostEqual(m[g.nbr_idx[g.nbr_ptr[0]]], 6.0 / g.counts[g.nbr_idx[g.nbr_ptr[0]]])
        self.assertEqual(G.nbr_max(g, f)[g.nbr_idx[g.nbr_ptr[0]]], 6.0)
        d = G.diffuse(g, f, 10)
        self.assertLess(d.max(), 6.0)
        self.assertGreater(np.count_nonzero(d > 1e-6), 20)

    def test_gradient_of_linear_field(self):
        g = self.g
        f = g.xyz[:, 2] * g.radius_km  # height ∝ z → gradient points north near the equator
        gr = G.gradient(g, f)
        eq = np.abs(g.lat) < 20
        _, n = east_north(g.xyz[eq])
        cos_to_north = np.sum(gr[eq] * n, axis=1) / np.linalg.norm(gr[eq], axis=1)
        self.assertGreater(np.median(cos_to_north), 0.99)
        self.assertAlmostEqual(float(np.median(np.linalg.norm(gr[eq], axis=1))), 1.0, delta=0.1)

    def test_distance_matches_great_circle(self):
        g = self.g
        i = g.cell_index(0.0, 0.0)
        d = G.distance_to(g, np.arange(g.n) == i)
        gc = g.radius_km * np.arccos(np.clip(g.xyz @ g.xyz[i], -1, 1))
        far = gc > 3000
        ratio = d[far] / gc[far]
        self.assertTrue(np.all(ratio >= 0.999) and np.median(ratio) < 1.15)
        self.assertTrue(np.all(np.isinf(G.distance_to(g, np.zeros(g.n, bool)))))

    def test_nearest_source_labels(self):
        g = self.g
        a, b = g.cell_index(0.0, -90.0), g.cell_index(0.0, 90.0)
        _, src = G.nearest_source(g, [a, b])
        self.assertEqual(src[g.cell_index(0.0, -60.0)], a)
        self.assertEqual(src[g.cell_index(0.0, 60.0)], b)

    def test_priority_flood_fills_basin_and_drains(self):
        g = self.g
        ocean = g.lat < -30
        z = np.where(ocean, -1000.0, 500.0 + 10 * g.lat)
        pit = g.cell_index(40.0, 0.0)
        z[pit] = -50.0  # land pit below sea level, not connected to ocean
        zf = G.priority_flood(g, z, ocean)
        self.assertGreater(zf[pit], z[pit])
        recv, slope, dist = G.steepest_receivers(g, zf)
        recv[ocean] = np.flatnonzero(ocean)
        levels = G.receiver_levels(recv)
        self.assertEqual(sum(len(l) for l in levels), g.n)
        self.assertTrue(np.all(ocean[levels[0]]))

    def test_priority_flood_needs_sink(self):
        with self.assertRaisesRegex(ValueError, "no sink"):
            G.priority_flood(self.g, np.ones(self.g.n), np.zeros(self.g.n, bool))

    def test_accumulate_conserves(self):
        g = self.g
        ocean = g.lat < -30
        z = np.where(ocean, -1000.0, 1000.0 + 20 * g.lat)
        zf = G.priority_flood(g, z, ocean)
        recv, _, _ = G.steepest_receivers(g, zf)
        recv[ocean] = np.flatnonzero(ocean)
        lv = G.receiver_levels(recv)
        w = np.where(ocean, 0.0, 1.0)
        acc = G.accumulate(recv, lv, w)
        self.assertAlmostEqual(acc[lv[0]].sum(), w.sum())

    def test_cycle_detected(self):
        with self.assertRaisesRegex(ValueError, "cycle"):
            G.receiver_levels(np.array([1, 0, 2]))

    def test_components(self):
        g = self.g
        m = (np.abs(g.lat) < 10) & (np.abs(g.lon) < 20) | (np.abs(g.lat - 50) < 8) & (np.abs(g.lon) < 20)
        lab = G.components(g, m)
        self.assertEqual(len(np.unique(lab[m])), 2)
        self.assertTrue(np.all(lab[~m] == -1))


class SmoothKmTest(unittest.TestCase):
    def test_constant_preserved_and_spike_decays(self):
        g = small_grid(3)
        np.testing.assert_allclose(G.smooth_km(g, np.full(g.n, 5.0), 1000.0), 5.0, rtol=1e-4)
        i = g.cell_index(0.0, 0.0)
        s = G.smooth_km(g, (np.arange(g.n) == i).astype(float), 1000.0)
        d = g.radius_km * np.arccos(np.clip(g.xyz @ g.xyz[i], -1, 1))
        near, far = s[(d > 800) & (d < 1200)].mean(), s[(d > 2800) & (d < 3200)].mean()
        self.assertTrue(near > far > 0)

    def test_resolution_independent(self):
        vals = []
        for res in (2, 3):
            g = small_grid(res)
            d = g.radius_km * np.arccos(np.clip(g.xyz @ g.xyz[g.cell_index(0.0, 0.0)], -1, 1))
            s = G.smooth_km(g, (d < 2000).astype(float), 1500.0)
            vals.append(s[g.cell_index(0.0, 30.0)])      # ~6700 km away
        self.assertAlmostEqual(vals[0], vals[1], delta=0.25 * max(vals))
        self.assertGreater(min(vals), 0.005)


class OceanMaskTest(unittest.TestCase):
    def test_inland_depression_is_not_ocean(self):
        g = small_grid(3)
        z = np.where(g.lat < 0, -3000.0, 500.0)
        c = g.xyz[g.cell_index(40.0, 0.0)]
        d = g.radius_km * np.arccos(np.clip(g.xyz @ c, -1, 1))
        z[d < 600] = -50.0                                   # interior basin below sea level
        ocean = G.ocean_mask(g, z, 1.0e6)
        self.assertTrue(ocean[g.lat < -5].all())
        self.assertFalse(ocean[d < 600].any())


class SmoothKmRobustTest(unittest.TestCase):
    def test_zero_and_tiny_fields(self):
        g = small_grid(3)
        np.testing.assert_array_equal(G.smooth_km(g, np.zeros(g.n), 25.0), 0.0)
        f = np.where(g.lat > 0, 1e-9, 0.0)
        s = G.smooth_km(g, f, 25.0)
        self.assertTrue(np.all(np.isfinite(s)) and s.max() <= 1e-9 * (1 + 1e-6))

    def test_femto_scale_field(self):
        g = small_grid(3)
        rng = np.random.default_rng(0)
        f = np.where(rng.random(g.n) < 0.06, rng.random(g.n) * 8e-14, 0.0)
        s = G.smooth_km(g, f, 25.0)
        self.assertTrue(np.all(np.isfinite(s)))
        self.assertAlmostEqual(float(s.sum() / f.sum()), 1.0, delta=0.05)


def _flood_ref(g, z, sink_mask, eps=0.01):
    """The original pure-Python priority flood (oracle for the compiled one)."""
    import heapq
    has_open = np.bincount(g.src, weights=(~sink_mask)[g.dst].astype(np.float64), minlength=g.n) > 0
    zf, done = np.asarray(z, dtype=np.float64).tolist(), sink_mask.tolist()
    ptr, idx = g.nbr_ptr.tolist(), g.nbr_idx.tolist()
    heap = [(zf[i], i) for i in np.flatnonzero(sink_mask & has_open).tolist()]
    heapq.heapify(heap)
    while heap:
        zc, c = heapq.heappop(heap)
        for k in range(ptr[c], ptr[c + 1]):
            n = idx[k]
            if not done[n]:
                done[n] = True
                zf[n] = max(zf[n], zc + eps)
                heapq.heappush(heap, (zf[n], n))
    return np.array(zf)


def _steepest_ref(g, z):
    """The original lexsort version (oracle)."""
    slope = (z[g.src] - z[g.dst]) / g.edge_km
    first = np.lexsort((-slope, g.src))[g.nbr_ptr[:-1]]
    s = slope[first]
    down = s > 0
    return (np.where(down, g.dst[first], np.arange(g.n)), np.where(down, s, 0.0),
            np.where(down, g.edge_km[first], np.inf))


def _accumulate_ref(recv, levels, w):
    acc = np.asarray(w, dtype=np.float64).copy()
    for lv in reversed(levels[1:]):
        acc += np.bincount(recv[lv], weights=acc[lv], minlength=len(acc))
    return acc


class FastPathsTest(unittest.TestCase):
    """The speed-ups (numba flood, sort-free receivers, per-level accumulate) give bit-identical results."""

    def fields(self, g):
        rng = np.random.default_rng(7)
        rough = rng.normal(0, 300, g.n)
        flat = np.round(rng.normal(0, 2, g.n))            # many exact ties: tie order matters
        pits = np.where(rng.random(g.n) < 0.2, -50.0, rough)
        return {"rough": rough, "flat": flat, "pits": pits}

    def test_flood_matches_reference(self):
        g = small_grid(3)
        for name, z in self.fields(g).items():
            sink = z < np.quantile(z, 0.1)
            for eps in (0.01, 0.0):
                with self.subTest(name=name, eps=eps):
                    want = _flood_ref(g, z, sink, eps)
                    got = G.priority_flood(g, z, sink, eps)
                    self.assertTrue(np.array_equal(got, want))
                    py = G._flood_py(z, sink, g.nbr_ptr, g.nbr_idx,
                                     np.flatnonzero(sink & (np.bincount(g.src, weights=(~sink)[g.dst].astype(float),
                                                                       minlength=g.n) > 0)), eps)
                    self.assertTrue(np.array_equal(py, want))

    def test_steepest_receivers_match_reference(self):
        g = small_grid(3)
        for name, z in self.fields(g).items():
            with self.subTest(name=name):
                for got, want in zip(G.steepest_receivers(g, z), _steepest_ref(g, z)):
                    self.assertTrue(np.array_equal(got, want))

    def test_accumulate_matches_reference(self):
        g = small_grid(3)
        rng = np.random.default_rng(3)
        for name, z in self.fields(g).items():
            with self.subTest(name=name):
                zf = G.priority_flood(g, z, z < np.quantile(z, 0.1))
                recv, _, _ = G.steepest_receivers(g, zf)
                lv = G.receiver_levels(recv)
                w = rng.random(g.n) * 1e3 * np.where(rng.random(g.n) < 0.1, -0.0, 1.0)   # with negative zeros
                got, want = G.accumulate(recv, lv, w), _accumulate_ref(recv, lv, w)
                self.assertTrue(np.array_equal(got, want))
                self.assertTrue(np.array_equal(np.signbit(got), np.signbit(want)))


class SweepTest(unittest.TestCase):
    def test_sweep_matches_full_bincount(self):
        from mapgen import hydrology as HY

        def sweep_ref(recv, levels, water, outlets, cap):
            acc = np.asarray(water, dtype=np.float64).copy()
            loss = np.zeros(len(acc))
            is_out = np.zeros(len(acc), bool)
            is_out[outlets] = True
            cap_cell = np.zeros(len(acc))
            cap_cell[outlets] = cap
            for lv in reversed(levels[1:]):
                push = acc[lv].copy()
                o = is_out[lv]
                if o.any():
                    cells = lv[o]
                    lost = np.minimum(acc[cells], cap_cell[cells])
                    loss[cells] = lost
                    push[o] = acc[cells] - lost
                acc += np.bincount(recv[lv], weights=push, minlength=len(acc))
            return acc, loss

        g = small_grid(3)
        rng = np.random.default_rng(11)
        z = rng.normal(0, 300, g.n)
        zf = G.priority_flood(g, z, z < np.quantile(z, 0.1))
        recv, _, _ = G.steepest_receivers(g, zf)
        lv = G.receiver_levels(recv)
        water = rng.random(g.n) * np.where(rng.random(g.n) < 0.1, -0.0, 1.0)     # with negative zeros
        outlets = rng.choice(g.n, 200, replace=False)
        cap = rng.random(200) * 2
        for got, want in zip(HY._sweep(recv, lv, water, outlets, cap), sweep_ref(recv, lv, water, outlets, cap)):
            self.assertTrue(np.array_equal(got, want))
            self.assertTrue(np.array_equal(np.signbit(got), np.signbit(want)))


class LeavesTest(unittest.TestCase):
    def test_leaves_all_matches_the_walk(self):
        from mapgen import hydrology as HY
        g = small_grid(3)
        rng = np.random.default_rng(4)
        for seed in range(3):
            z = rng.normal(0, 300, g.n)
            ocean = z < np.quantile(z, 0.2)
            zf = G.priority_flood(g, z, ocean)
            lab = G.components(g, ~ocean & (zf - z > 1.0))
            recv, _, _ = G.steepest_receivers(g, zf)
            recv = np.where(ocean, np.arange(g.n), recv)
            lv = G.receiver_levels(recv)
            xs = np.flatnonzero(lab >= 0)
            for limit in (100000, 3):
                want = np.array([HY._leaves(recv, lab, x, limit) for x in xs])
                got = HY._leaves_all(recv, lab, lv, limit)[xs]
                self.assertGreater(want.sum(), 0)
                self.assertTrue(np.array_equal(got, want), (seed, limit))


class BicgstabJacobiTest(unittest.TestCase):
    """The fused solver walks scipy's iterates exactly: same answers bit for bit, same exit codes."""
    def systems(self):
        from scipy import sparse
        rng = np.random.default_rng(0)
        for n in (2000, 9000):
            i = np.repeat(np.arange(n), 6)
            j = (i + rng.integers(-50, 50, len(i))) % n
            L = sparse.csr_matrix((rng.random(len(i)), (i, j)), shape=(n, n))
            S = L + L.T
            yield (sparse.diags(np.asarray(S.sum(1)).ravel()) - S).tocsr() * 40 + sparse.identity(n, format="csr")
            yield (sparse.identity(n, format="csr") * (1 + np.asarray(L.sum(1)).ravel().max() * 0.6) - L).tocsr()

    def setUp(self):
        self.min_n = G.JIT_MIN_N
        G.JIT_MIN_N = 0                                      # the compiled path even on small test systems

    def tearDown(self):
        G.JIT_MIN_N = self.min_n

    def test_matches_scipy_bicgstab(self):
        from scipy.sparse import linalg as splinalg
        rng = np.random.default_rng(1)
        for k, A in enumerate(self.systems()):
            for rtol, maxiter in ((1e-6, 5000), (1e-9, 5000), (1e-12, 7)):
                b = rng.normal(size=A.shape[0])
                inv = 1.0 / A.diagonal()
                M = splinalg.LinearOperator(A.shape, matvec=lambda x: inv * x)
                want = splinalg.bicgstab(A, b, x0=b * 0.5, rtol=rtol, maxiter=maxiter, M=M)
                got = G.bicgstab_jacobi(A, b, b * 0.5, inv, rtol, maxiter)
                with self.subTest(k=k, rtol=rtol, maxiter=maxiter):
                    self.assertEqual(got[1], want[1])
                    self.assertTrue(np.array_equal(got[0], want[0]))

    def test_zero_right_hand_side_and_zero_start(self):
        A = next(self.systems())
        inv = 1.0 / A.diagonal()
        x, info = G.bicgstab_jacobi(A, np.zeros(A.shape[0]), np.zeros(A.shape[0]), inv, 1e-6, 100)
        self.assertEqual(info, 0)
        self.assertFalse(x.any())


class PmapTest(unittest.TestCase):
    def test_order_and_same_floats_as_serial(self):
        import os
        from unittest import mock
        g = small_grid(2)
        fields = [np.random.default_rng(i).normal(size=g.n) for i in range(4)]
        serial = [G.smooth_km(g, f, 900.0) for f in fields]
        with mock.patch.dict(os.environ, {"WORLDGEN_THREADS": "4"}):
            self.assertEqual(G.workers(), 4)
            got = G.pmap(lambda f: G.smooth_km(g, f, 900.0), fields)
        for a, b in zip(serial, got):
            self.assertTrue(np.array_equal(a, b))
        with mock.patch.dict(os.environ, {"WORLDGEN_THREADS": "3"}):
            self.assertEqual(G.pmap(lambda k: k * k, range(9)), [k * k for k in range(9)])
        with mock.patch.dict(os.environ, {"WORLDGEN_THREADS": "x"}):
            self.assertEqual(G.workers(), 3)


class ComponentsTest(unittest.TestCase):
    def test_same_labels_as_scipy(self):
        from scipy import sparse
        from scipy.sparse import csgraph
        g = small_grid(3)
        rng = np.random.default_rng(11)

        def scipy_labels(mask):                                # the previous implementation, as the oracle
            e = mask[g.src] & mask[g.dst]
            m = sparse.csr_matrix((np.ones(int(e.sum())), (g.src[e], g.dst[e])), shape=(g.n, g.n))
            _, lab = csgraph.connected_components(m, directed=False)
            return np.where(mask, lab, -1)
        for p in (0.0, 0.2, 0.45, 0.6, 0.9, 1.0):
            for _ in range(3):
                mask = rng.random(g.n) < p
                want, got = scipy_labels(mask), G.components(g, mask)
                self.assertEqual(got.dtype, want.dtype)
                self.assertTrue(np.array_equal(got, want), p)
        smooth = G.smooth_km(g, rng.normal(size=g.n), 2000.0) > 0      # big blobs, many cells each
        self.assertTrue(np.array_equal(G.components(g, smooth), scipy_labels(smooth)))
        self.assertGreater(len(np.unique(G.components(g, smooth))), 2)