import unittest

import numpy as np

from mapgen import crust, elevation as EL, plates
from tests.helpers import make_ctx


def scenario(tect, land_fn, land_fraction, res=3, gravity=None, lock=None):
    ctx = make_ctx(res, tect=tect, cfg={"build": {"land_fraction": land_fraction}})
    g = ctx.grid
    z0 = np.zeros(g.n)
    ctx.data.update({"sk_land": land_fn(g).astype(float), "m_land_hint": z0, "sk_mountains": z0,
                     "m_mountain_hint": z0, "m_gravity_zones": z0 if gravity is None else gravity(g),
                     "m_lock": z0 if lock is None else lock(g)})
    ctx.data.update(plates.run(ctx))
    ctx.data.update(crust.run(ctx))
    return ctx


SUBDUCT = {"plate": [
    {"id": "c", "seed": [0.0, -30.0], "kind": "continental", "motion": [90.0, 3.0]},
    {"id": "o", "seed": [0.0, 10.0], "kind": "oceanic", "motion": [270.0, 5.0]}]}
COLLIDE = {"plate": [
    {"id": "a", "seed": [0.0, -30.0], "kind": "continental", "motion": [90.0, 4.0]},
    {"id": "b", "seed": [0.0, 30.0], "kind": "continental", "motion": [270.0, 4.0]}]}


class ElevationTest(unittest.TestCase):
    def test_sea_level_hits_target(self):
        rng = np.random.default_rng(1)
        z = rng.normal(size=5000) * 1000
        area = rng.uniform(1, 2, size=5000)
        zs = EL.solve_sea_level(z, area, 0.3)
        self.assertAlmostEqual(area[zs > 0].sum() / area.sum(), 0.3, delta=0.005)

    def test_subduction_trench_and_coastal_range(self):
        ctx = scenario(SUBDUCT, lambda g: (np.abs(g.lat) < 40) & (g.lon > -70) & (g.lon < 0), 0.12)
        out = EL.run(ctx)
        z, g = out["elevation_m"], ctx.grid
        band = np.abs(g.lat) < 30
        ocean_side = band & (ctx.data["plate"] == 1) & (out["d_sub_km"] < 300)
        cont_side = band & (ctx.data["plate"] == 0) & (out["d_over_km"] < 600)
        self.assertLess(z[ocean_side].min(), -7000)
        self.assertGreater(z[cont_side].max(), 2500)

    def test_collision_builds_high_range(self):
        ctx = scenario(COLLIDE, lambda g: (np.abs(g.lat) < 30) & (np.abs(g.lon) < 60), 0.15)
        out = EL.run(ctx)
        near = out["d_coll_km"] < 300
        self.assertGreater(out["elevation_m"][near].max(), 6000)
        self.assertLessEqual(out["elevation_m"].max(), 12000)

    def test_low_gravity_zone_raises_relief(self):
        land = lambda g: (np.abs(g.lat) < 30) & (np.abs(g.lon) < 60)
        zone = lambda g: np.where((np.abs(g.lat) < 20) & (g.lon > 10) & (g.lon < 50), -1.0, 0.0)
        base = EL.run(scenario(COLLIDE, land, 0.15))["elevation_m"]
        ctx = scenario(COLLIDE, land, 0.15, gravity=zone)
        low = EL.run(ctx)["elevation_m"]
        inside = zone(ctx.grid) < 0
        self.assertGreater(low[inside & (base > 0)].mean(), base[inside & (base > 0)].mean() * 1.5)

    def test_lock_forces_coast(self):
        land = lambda g: (np.abs(g.lat) < 30) & (np.abs(g.lon) < 60)
        lock = lambda g: np.ones(g.n)
        ctx = scenario(COLLIDE, land, 0.15, lock=lock)
        z = EL.run(ctx)["elevation_m"]
        self.assertTrue(np.all((z > 0) == (ctx.data["sk_land"] > 0.5)))


class SeaLevelRegressionTest(unittest.TestCase):
    LAND = staticmethod(lambda g: (g.lat > -25) & (g.lat < 45) & (np.abs(g.lon) < 62))   # ≈ 0.195 of the sphere

    def test_continents_not_lifted(self):
        ctx = scenario(COLLIDE, self.LAND, 0.19)
        out = EL.run(ctx)
        z, g = out["elevation_m"], ctx.grid
        self.assertLess(np.median(z[self.LAND(g) & (out["d_coll_km"] > 1500)]), 1500)
        self.assertLess(np.median(z[~ctx.data["continental"]]), -3000)

    def test_target_above_continental_area_errors(self):
        from mapgen.pipeline import StageError
        with self.assertRaisesRegex(StageError, "land_fraction"):
            EL.run(scenario(COLLIDE, self.LAND, 0.30))


class LockThresholdTest(unittest.TestCase):
    def test_near_neutral_lock_does_nothing(self):
        land = lambda g: (np.abs(g.lat) < 30) & (np.abs(g.lon) < 60)
        base = EL.run(scenario(COLLIDE, land, 0.15))["elevation_m"]
        tiny = EL.run(scenario(COLLIDE, land, 0.15, lock=lambda g: np.full(g.n, 1 / 255)))["elevation_m"]
        np.testing.assert_array_equal(base, tiny)


class MarginProfileTest(unittest.TestCase):
    @staticmethod
    def _profile(extra):
        from mapgen.graph import distance_to, ocean_mask
        land = lambda g: (np.abs(g.lat) < 30) & (np.abs(g.lon) < 60)
        g = make_ctx(4).grid
        frac = g.area_km2[land(g)].sum() / g.area_km2.sum()
        ctx = make_ctx(4, tect=COLLIDE, cfg={"build": {"land_fraction": round(frac - 0.003, 4)},
                                             "crust": {"edge_noise": 0.0}, "elevation": extra})
        z0 = np.zeros(g.n)
        ctx.data.update({"sk_land": land(g).astype(float), "m_land_hint": z0, "sk_mountains": z0,
                         "m_mountain_hint": z0, "m_gravity_zones": z0, "m_lock": z0})
        ctx.data.update(plates.run(ctx))
        ctx.data.update(crust.run(ctx))
        z = EL.run(ctx)["elevation_m"].astype(float)
        sea = ocean_mask(g, z, 5.0e6)
        d_coast_land = distance_to(g, sea)
        d_coast_sea = distance_to(g, ~sea)
        far = (~sea) & (d_coast_land > 1200) & (np.abs(g.lon) > 20)       # interior, away from the collision belt
        coastal = (~sea) & (d_coast_land < 250) & (np.abs(g.lon) > 20)
        shelf = sea & (d_coast_sea < 200)
        return np.median(z[coastal]), np.median(z[far]), np.median(z[shelf])

    def test_coastal_lowlands_and_shelves(self):
        step = self._profile({"coast_noise_m": 0.0, "margin_km": 1.0, "slope_km": 1.0})
        coastal, interior, shelf = self._profile({})
        self.assertLess(coastal, interior - 300)              # coastal plains sit well below the interior
        self.assertGreater(shelf, -800)                        # a shallow shelf fringes the coast
        self.assertGreaterEqual(step[0], step[1] - 300)        # the old step margin had no coastal lowlands


class ConnectedSeaLevelTest(unittest.TestCase):
    def test_interior_pits_do_not_count_as_sea(self):
        from mapgen.graph import ocean_mask
        from tests.helpers import small_grid
        g = small_grid(3)
        land = (np.abs(g.lat) < 40) & (np.abs(g.lon) < 70)
        z = np.where(land, 300.0 + 40.0 * (70.0 - np.abs(g.lon)), -4000.0)     # rises inland (continuous)
        pits = land & ((g.lat % 10) < 3) & ((g.lon % 10) < 3) & (np.abs(g.lon) < 55)
        z[pits] = -500.0                                                    # many deep interior pits
        target = 0.8 * g.area_km2[land].sum() / g.area_km2.sum()
        zs = EL.solve_sea_level_connected(g, z, target, 5.0e6)
        sea = ocean_mask(g, zs, 5.0e6)
        self.assertAlmostEqual(g.area_km2[~sea].sum() / g.area_km2.sum(), target, delta=0.01)


class RevisionElevationTest(unittest.TestCase):
    LAND = staticmethod(lambda g: np.abs(g.lon + 30) < 25)

    def test_without_patches_the_relief_is_solved_once(self):
        from unittest import mock
        ctx = scenario(SUBDUCT, self.LAND, 0.1)
        with mock.patch.object(EL, "_relief", wraps=EL._relief) as rel:
            EL.run(ctx)
        self.assertEqual(rel.call_count, 1)

    def test_a_land_patch_keeps_the_sea_level_of_the_world_without_it(self):
        from mapgen.crust import center_dist
        za = EL.run(scenario(SUBDUCT, self.LAND, 0.1))["elevation_m"]
        tect = {**SUBDUCT, "land_patch": [{"name": "fill", "center": [0.0, 120.0], "radius_km": 1500.0}]}
        b = scenario(tect, self.LAND, 0.1)
        zb = EL.run(b)["elevation_m"]
        g = b.grid
        far = (za > 0) & (center_dist(g, [0.0, 120.0]) > 5000.0)
        self.assertGreater(far.sum(), 20)
        np.testing.assert_allclose(zb[far], za[far], atol=1e-3)     # the other coasts don't move
        i = g.cell_index(0.0, 120.0)
        self.assertGreater(zb[i] - za[i], 1000.0, "the patch is continental ground now")

    def test_a_land_patch_is_land_though_the_sea_level_drowns_bare_crust(self):
        from mapgen.crust import center_dist
        tect = {**SUBDUCT, "land_patch": [{"name": "fill", "center": [0.0, 120.0], "radius_km": 1500.0}]}
        b = scenario(tect, self.LAND, 0.05)
        z = EL.run(b)["elevation_m"]
        inner = center_dist(b.grid, [0.0, 120.0]) < 1000.0
        self.assertGreater(float(np.mean(z[inner] > 0)), 0.9)

    def test_plateaus_get_their_surface_after_the_sea_level_solve(self):
        from mapgen import plateaus as PL
        from mapgen.crust import center_dist
        plats = [{"name": "p", "center": [0.0, 120.0], "area_km2": 3.0e6, "top_m": [1500.0, 3000.0]},
                 {"name": "q", "center": [40.0, 150.0], "area_km2": 1.0e6, "top_m": [1000.0, 1500.0],
                  "islands": True}]
        za = EL.run(scenario(SUBDUCT, self.LAND, 0.1))["elevation_m"]
        b = scenario({**SUBDUCT, "plateau": plats}, self.LAND, 0.1)
        zb = EL.run(b)["elevation_m"]
        g, ids = b.grid, b.data["plateau_id"]
        away = (center_dist(g, [0.0, 120.0]) > 2500.0) & (center_dist(g, [40.0, 150.0]) > 1500.0)
        np.testing.assert_allclose(zb[away], za[away], atol=1e-3)
        _, short = PL.semi_axes(plats[0])
        core = np.flatnonzero(ids == 0)
        core = core[(1.0 - PL.rho(g.xyz[core], plats[0], b.seed, g.radius_km)) * short > PL.MARGIN_KM]
        self.assertGreater(len(core), 5)
        self.assertTrue(1500.0 - PL.RELIEF_M <= float(np.median(-zb[core])) <= 3000.0 + PL.RELIEF_M)
        self.assertLessEqual(zb[ids == 0].max(), PL.HIDDEN_MAX_M)
        near_q = center_dist(g, [40.0, 150.0]) < 1500.0
        self.assertGreater(zb[near_q].max(), 0.0, "the island plateau breaks the surface")

    def test_land_added_marks_what_the_patch_makes_land(self):
        from mapgen.crust import center_dist
        from mapgen.graph import ocean_mask
        a = EL.run(scenario(SUBDUCT, self.LAND, 0.05))
        tect = {**SUBDUCT, "land_patch": [{"name": "fill", "center": [0.0, 120.0], "radius_km": 1500.0}]}
        b = scenario(tect, self.LAND, 0.05)
        ob = EL.run(b)
        g = b.grid
        land_a = ~ocean_mask(g, a["elevation_m"])
        land_b = ~ocean_mask(g, ob["elevation_m"])
        added = ob["land_added"]
        self.assertGreater(float(added[center_dist(g, [0.0, 120.0]) < 1000.0].mean()), 0.9)
        self.assertFalse((added & land_a).any(), "land without the patch is not added land")
        self.assertFalse((added & ~land_b).any(), "added land is land")
        self.assertFalse(a["land_added"].any(), "no patch, no plateau: nothing added")