import json
import unittest

import numpy as np
from scipy.spatial import cKDTree

import rivers as RV
from tests.test_serve import built_world


def net_and_arrays():
    root = built_world()
    d = dict(np.load(root / "out" / "r2" / "cells.npz"))
    meta = json.loads((root / "out" / "r2" / "cells_meta.json").read_text())
    return RV.RiverNet(d, meta["legends"], meta["radius_km"]), d


class RiverNetTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.net, cls.a = net_and_arrays()
        cls.r = np.where(cls.a["river"].astype(bool))[0]

    def test_levels_never_above_ground_or_below_the_cap(self):
        z, lev = self.a["z_surface_m"][self.r], self.net.level[self.r]
        self.assertTrue(np.all(np.isfinite(lev)))
        self.assertTrue(np.all(lev <= z + 1e-6))
        self.assertTrue(np.all(z - lev <= self.net.cap[self.r] + 1e-6))

    def test_levels_never_rise_downstream(self):
        recv, riv = self.a["recv"], self.a["river"].astype(bool)
        down = recv[self.r]
        on = riv[down]
        self.assertTrue(np.all(self.net.level[down[on]] <= self.net.level[self.r[on]] + 1e-6))

    def test_levels_meet_their_mouth(self):
        recv, ocean, lake = self.a["recv"], self.a["ocean"].astype(bool), self.a["lake"].astype(bool)
        sea = self.r[ocean[recv[self.r]]]
        self.assertTrue(len(sea) > 0)
        k = np.isin(self.net.seg_a, sea)
        np.testing.assert_allclose(self.net.level_b[k], 0.0)
        lk = self.r[lake[recv[self.r]]]
        k = np.isin(self.net.seg_a, lk)
        np.testing.assert_allclose(self.net.level_b[k], self.a["z_filled_m"][recv[self.net.seg_a[k]]])

    def test_drops_below_lakes_are_cataracts_not_cliffs(self):
        recv, riv, lake = self.a["recv"], self.a["river"].astype(bool), self.a["lake"].astype(bool)
        src = np.where(lake & riv[recv] & (recv != np.arange(len(recv))))[0]
        for i in src:
            j = recv[i]
            dist = np.linalg.norm(self.a["g_xyz"][i] - self.a["g_xyz"][j]) * self.net.R
            drop = self.a["z_filled_m"][i] - self.net.level[j]
            self.assertLessEqual(drop, RV.KNICK_M_KM * dist + 1e-6 + max(0.0, self.a["z_filled_m"][i] - self.a["z_surface_m"][j]))

    def test_geometry_per_segment(self):
        n = len(self.net.seg_a)                       # river cells + lake outlets
        self.assertGreaterEqual(n, len(self.r))
        for k in ("a_xyz", "b_xyz", "level_a", "level_b", "half_w", "floor", "wall", "meander", "reach", "seg_len"):
            self.assertEqual(len(getattr(self.net, k)), n, k)
        self.assertTrue(np.all(self.net.half_w >= 0.03) and np.all(self.net.wall > 0))


class StyleTest(unittest.TestCase):
    def test_hard_dry_rock_makes_steeper_narrower_valleys(self):
        g = RV.valley_style("plateau", "—", "granite/gneiss", 400.0)
        s = RV.valley_style("plateau", "—", "sandstone/shale", 1800.0)
        self.assertGreater(g[1], 3 * s[1])       # walls
        self.assertLess(g[0], s[0])              # floor
        self.assertGreater(g[1], 2500)           # the granite gorge: sheer

    def test_floodplains_are_wide_and_meander(self):
        f = RV.valley_style("plain", "floodplain", "sandstone/shale", 900.0)
        m = RV.valley_style("mountains", "—", "granite/gneiss", 900.0)
        self.assertGreater(f[0], 10 * m[0])
        self.assertGreater(f[2], 0)
        self.assertEqual(m[2], 0)

    def test_channel_width(self):
        self.assertAlmostEqual(RV.half_width_km(2.0), 0.004 * np.sqrt(2.0 * 31.7), places=9)


class ValleyTest(unittest.TestCase):
    def test_profile(self):
        d = np.array([0.0, 0.05, 0.2, 1.0, 2.0])
        v = RV.valley_surface(d, 100.0, 0.1, 0.3, 2000.0)
        self.assertEqual(v[0], 100.0)                        # channel: the water level
        self.assertEqual(v[1], 100.0)
        self.assertTrue(101.0 <= v[2] <= 103.0)              # floor
        self.assertAlmostEqual(v[4] - v[3], 2000.0, places=6)   # wall: 2,000 m per km

    def test_points_are_a_global_lattice(self):
        net, _ = net_and_arrays()
        s = int(np.argmax(net.seg_len))
        c = net.a_xyz[s]
        p1, _ = net.segment_points(s, c, 50.0, 5.0)
        p2, _ = net.segment_points(s, c, 400.0, 5.0)
        self.assertGreater(len(p2), len(p1))
        a = {tuple(np.round(v, 12)) for v in p1}
        b = {tuple(np.round(v, 12)) for v in p2}
        self.assertTrue(a <= b, "a wider query returns a superset of the same points")

    def test_antimeridian_segment(self):
        net, _ = net_and_arrays()
        a = np.array([np.cos(np.radians(10)) * np.cos(np.radians(179.9)), np.cos(np.radians(10)) * np.sin(np.radians(179.9)), np.sin(np.radians(10))])
        b = np.array([np.cos(np.radians(10)) * np.cos(np.radians(-179.9)), np.cos(np.radians(10)) * np.sin(np.radians(-179.9)), np.sin(np.radians(10))])
        pts = RV.densify(a, b, 0.0, 1.0, 1.0 / net.R)   # 1 km steps (a chord in planet radii)
        self.assertTrue(np.all(np.abs(np.linalg.norm(pts, axis=1) - 1) < 1e-12))
        self.assertLess(np.max(np.linalg.norm(np.diff(pts, axis=0), axis=1)) * net.R, 1.01)


class RoughWallTest(unittest.TestCase):
    def test_walls_keep_the_terrain_roughness_but_the_channel_stays_flat(self):
        d = np.array([0.0, 0.3, 1.0, 2.0])
        smooth = RV.valley_surface(d, 100.0, 0.1, 0.1, 500.0)
        rough = RV.valley_surface(d, 100.0, 0.1, 0.1, 500.0, rough=np.array([80.0, 80.0, 80.0, -40.0]))
        self.assertEqual(rough[0], smooth[0])                   # the water level
        self.assertAlmostEqual(rough[2] - smooth[2], 80.0)      # a wall well away from the floor: full roughness
        self.assertAlmostEqual(rough[3] - smooth[3], -40.0)
        self.assertTrue(0 < rough[1] - smooth[1] < 80.0)        # fading in above the floor


LEG = {"landform": ["ocean", "plain", "hills"], "ground": ["—", "floodplain"], "lithology": ["granite/gneiss", "sandstone/shale"],
       "age_class": ["oceanic", "craton (3g era)"]}


def synthetic(cells):
    """A tiny world. cells: (lat, lon, z, discharge, landform, ground, receiver index, kind: river|ocean|lake|land)."""
    ll = np.radians([(c[0], c[1]) for c in cells])
    a = {"g_xyz": np.stack([np.cos(ll[:, 0]) * np.cos(ll[:, 1]), np.cos(ll[:, 0]) * np.sin(ll[:, 1]), np.sin(ll[:, 0])], 1),
         "z_surface_m": np.array([c[2] for c in cells], float), "z_filled_m": np.array([c[2] for c in cells], float),
         "discharge_km3_yr": np.array([c[3] for c in cells], float),
         "landform": np.array([LEG["landform"].index(c[4]) for c in cells]), "ground": np.array([LEG["ground"].index(c[5]) for c in cells]),
         "recv": np.array([c[6] for c in cells]), "river": np.array([c[7] == "river" for c in cells]),
         "ocean": np.array([c[7] == "ocean" for c in cells]), "lake": np.array([c[7] == "lake" for c in cells]),
         "lithology": np.ones(len(cells), int), "age_class": np.ones(len(cells), int), "P_ann": np.full(len(cells), 900.0)}
    return RV.RiverNet(a, LEG, 12742.0)


def xyz(lat, lon):
    la, lo = np.radians(lat), np.radians(lon)
    return np.array([[np.cos(la) * np.cos(lo), np.cos(la) * np.sin(lo), np.sin(la)]])


class PerValleyTest(unittest.TestCase):
    # a big floodplain river along the equator (east to the sea) and a small hill stream 22 km north joining it
    net = synthetic([(0, 0.0, 100, 5000, "plain", "floodplain", 1, "river"), (0, 0.3, 90, 5000, "plain", "floodplain", 2, "river"),
                     (0, 0.6, 80, 5000, "plain", "floodplain", 3, "river"), (0, 0.9, 70, 5000, "plain", "floodplain", 4, "river"),
                     (0, 1.2, -50, 0, "ocean", "—", 4, "ocean"),
                     (0.1, 0.3, 400, 2, "hills", "—", 6, "river"), (0.1, 0.6, 380, 2, "hills", "—", 2, "river")])

    def test_a_big_valley_is_not_crowded_out_by_a_nearby_stream(self):
        q = xyz(0.05, 0.45)                                           # 11 km from both centrelines
        V, ch, dr = self.net.valleys(q, 0.5)
        big = float(np.interp(0.45, [0.3, 0.6], [self.net.level[1], self.net.level[2]]))
        self.assertLessEqual(V[0], big + 3.0 + 1e-6, "on the big river's floodplain floor")

    def test_the_answer_does_not_depend_on_the_other_query_points(self):
        q = xyz(0.05, 0.45)
        many = np.vstack([q, xyz(3.0, 20.0), xyz(-0.2, 0.1), xyz(0.08, 0.5)])
        a, b = self.net.valleys(q, 0.5)[0], self.net.valleys(many, 0.5)[0]
        self.assertEqual(a[0], b[0])

    def test_floors_stay_inside_the_reach(self):
        self.assertTrue(np.all(self.net.floor <= RV.FLOOR_MAX_KM + 1e-9))
        self.assertTrue(np.all(self.net.reach >= self.net.half_w + self.net.floor))

    def test_rough_walls_never_dip_below_the_water(self):
        v = RV.valley_surface(np.array([0.6, 1.0]), 100.0, 0.1, 0.3, 500.0, rough=np.array([-500.0, -500.0]))
        self.assertTrue(np.all(v >= 101.0))


class BankTest(unittest.TestCase):
    def test_a_river_doubling_back_below_itself_keeps_its_banks_above_its_water(self):
        # a canyon river runs east, then back west 5.5 km north and 400 m lower: its lower floor must not
        # undercut the upper reach (dry pits below the water beside it)
        net = synthetic([(0, 0.0, 900, 300, "hills", "—", 1, "river"), (0, 0.3, 700, 300, "hills", "—", 2, "river"),
                         (0.05, 0.0, 300, 300, "hills", "—", 3, "river"), (0.05, -0.3, -50, 0, "ocean", "—", 3, "ocean")])
        la, lo = np.meshgrid(np.linspace(-0.03, 0.08, 40), np.linspace(-0.05, 0.35, 80), indexing="ij")
        q = np.vstack([xyz(a, b) for a, b in zip(la.ravel(), lo.ravel())])
        V, ch, _ = net.valleys(q, 0.1)
        best, lev, bank = np.full(len(q), np.inf), np.zeros(len(q)), np.zeros(len(q), bool)
        c = q.mean(0) / np.linalg.norm(q.mean(0))
        for s in range(len(net.seg_a)):                               # the nearest reach, its water, its floor zone
            p, t = net.segment_points(s, c, 60.0, 0.05)
            d, k = cKDTree(p).query(q)
            d *= net.R
            b = d < best
            best[b], bank[b] = d[b], (d <= net.half_w[s] + net.floor[s])[b]
            lev[b] = (net.level_a[s] + (net.level_b[s] - net.level_a[s]) * t[k])[b]
        dry = ~ch & bank
        self.assertGreater(dry.sum(), 100)
        self.assertEqual(int((V[dry] < lev[dry]).sum()), 0, "floor below the water of the reach beside it")


class LakeOutletTest(unittest.TestCase):
    def test_a_lake_outlet_joins_its_river_starting_at_the_lake_surface(self):
        net = synthetic([(0, 0.0, 1100, 0, "hills", "—", 1, "lake"), (0, 0.3, 1100, 300, "plain", "—", 2, "river"),
                         (0, 0.6, 600, 300, "plain", "—", 3, "river"), (0, 0.9, -50, 0, "ocean", "—", 3, "ocean")])
        k = np.where(net.seg_a == 0)[0]
        self.assertEqual(len(k), 1, "the lake outlet has a segment")
        self.assertEqual(net.level_a[k[0]], 1100.0)
        self.assertEqual(net.seg_b[k[0]], 1)


class GroundTest(unittest.TestCase):
    def test_given_levels_reach_their_shoulders_and_sub_pixel_streams_are_not_drawn(self):
        cells = [(0, 0.0, 800, 5000, "plateau", "—", 1, "river"), (0, 0.3, 700, 5000, "plateau", "—", 2, "river"),
                 (0, 0.6, -50, 0, "ocean", "—", 2, "ocean")]
        levels = np.array([400.0, 300.0, 0.0])
        ground = np.array([800.0, 700.0, 0.0])
        LEG2 = dict(LEG, landform=LEG["landform"] + ["plateau"])
        ll = np.radians([(c[0], c[1]) for c in cells])
        arr = {"g_xyz": np.stack([np.cos(ll[:, 0]) * np.cos(ll[:, 1]), np.cos(ll[:, 0]) * np.sin(ll[:, 1]), np.sin(ll[:, 0])], 1),
               "z_surface_m": ground.copy(), "z_filled_m": ground.copy(), "discharge_km3_yr": np.array([5000.0, 5000.0, 0.0]),
               "landform": np.array([3, 3, 0]), "ground": np.zeros(3, int), "recv": np.array([1, 2, 2]),
               "river": np.array([True, True, False]), "ocean": np.array([False, False, True]), "lake": np.zeros(3, bool),
               "lithology": np.zeros(3, int), "age_class": np.ones(3, int), "P_ann": np.full(3, 300.0)}
        net = RV.RiverNet(arr, LEG2, 12742.0, levels=levels, ground=ground)
        k = 0
        self.assertGreaterEqual(net.reach[k], net.half_w[k] + net.floor[k] + (ground[0] - levels[0]) / net.wall[k] - 1e-9)
        small = synthetic([(0, 0.0, 100, 0.2, "plain", "—", 1, "river"), (0, 0.3, 90, 0.2, "plain", "—", 2, "river"),
                           (0, 0.6, -50, 0, "ocean", "—", 2, "ocean")])
        q = xyz(0.0, 0.15)
        self.assertFalse(small.channel(q, 5.0)[1][0], "a 0.2 km³/yr stream is not drawn at 5 km per pixel")
        self.assertTrue(small.channel(q, 0.01)[1][0], "but is at street level")


class ChunkTest(unittest.TestCase):
    def test_a_coarse_tile_splits_into_compact_chunks(self):
        net = synthetic([(0, 0.0, 100, 5, "plain", "—", 1, "river"), (0, 0.3, 90, 5, "plain", "—", 2, "river"),
                         (0, 0.6, -50, 0, "ocean", "—", 2, "ocean")])
        la, lo = np.meshgrid(np.linspace(-52.0, -46.4, 256), np.linspace(-45.0, -33.75, 256), indexing="ij")   # a z5 tile
        a, o = np.radians(la.ravel()), np.radians(lo.ravel())
        chunks = net._chunks(np.stack([np.cos(a) * np.cos(o), np.cos(a) * np.sin(o), np.sin(a)], 1))
        self.assertLess(len(chunks), 300, "not thousands of one-row strips")
        self.assertEqual(sorted(np.concatenate([c[0] for c in chunks]).tolist()), list(range(256 * 256)))


class RiverNetStateTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        import refine as RF
        from tests.test_serve import built_world
        w = RF.WorldCells(built_world() / "out" / "r2")
        cls.a = {k: np.asarray(w.a[k]) for k in RF.NET_KEYS}
        cls.legends, cls.R = w.meta["legends"], float(w.meta["radius_km"])

    def test_state_round_trip_answers_the_same(self):
        import rivers as RV
        net = RV.RiverNet(self.a, self.legends, self.R)
        arrays, meta = net.state()
        self.assertEqual(set(arrays), set(RV.STATE))
        back = RV.RiverNet.from_state(arrays, meta)
        for k in RV.STATE:
            np.testing.assert_array_equal(getattr(back, k), getattr(net, k))
        self.assertEqual(back.max_extent, net.max_extent)
        p = np.asarray(self.a["g_xyz"][:200], dtype=np.float64)
        np.testing.assert_array_equal(back.valleys(p, 5.0)[1], net.valleys(p, 5.0)[1])
        np.testing.assert_array_equal(back.channel(p, 5.0)[1], net.channel(p, 5.0)[1])

    def test_tree_is_built_on_first_use(self):
        import rivers as RV
        net = RV.RiverNet.from_state(*RV.RiverNet(self.a, self.legends, self.R).state())
        self.assertIsNone(net._tree)
        self.assertIsNotNone(net.tree)
        self.assertIs(net.tree, net.tree)