import unittest

import numpy as np

from mapgen import erosion as ER
from mapgen.crust import CRATON
from tests.helpers import make_ctx


def mountain_ctx(k=None):
    cfg = {"build": {"land_fraction": 0.05}}
    if k is not None:
        cfg["erosion"] = {"k": k}
    ctx = make_ctx(3, cfg=cfg)
    g = ctx.grid
    d = g.radius_km * np.arccos(np.clip(g.xyz @ g.xyz[g.cell_index(0.0, 0.0)], -1, 1))
    z = np.where(d < 6000, 6000 * np.clip(1 - d / 6000, 0, None) ** 1.2 + 50, -4000.0)
    ctx.data.update({"elevation_m": z.astype(np.float32), "age_class": np.full(g.n, CRATON, np.int8)})
    return ctx, z


class ErosionTest(unittest.TestCase):
    def test_erode_lowers_land(self):
        ctx, z0 = mountain_ctx()
        kmult = np.full(ctx.grid.n, 1.5)
        z1 = ER.erode(ctx.grid, z0, kmult, ER.DEFAULTS)
        land = z0 > 0
        self.assertLess(z1[land].mean(), z0[land].mean())
        self.assertLessEqual(z1.max(), z0.max() + 1e-6)

    def test_stage_keeps_land_fraction(self):
        ctx, _ = mountain_ctx()
        z = ER.run(ctx)["elevation_eroded_m"]
        g = ctx.grid
        self.assertTrue(np.all(np.isfinite(z)))
        self.assertAlmostEqual(g.area_km2[z > 0].sum() / g.area_km2.sum(), 0.05, delta=0.01)

    def test_stable_with_huge_k(self):
        ctx, _ = mountain_ctx(k=1e3)
        z = ER.run(ctx)["elevation_eroded_m"]
        self.assertTrue(np.all(np.isfinite(z)))
        self.assertGreaterEqual(z.min(), -11000)

    def test_stage_outputs_connected_ocean(self):
        ctx, z0 = mountain_ctx()
        g = ctx.grid
        c = g.xyz[g.cell_index(0.0, 0.0)]
        d = g.radius_km * np.arccos(np.clip(g.xyz @ c, -1, 1))
        z0 = np.where(d < 700, -500.0, z0)                     # deep interior pit in the mountain block
        ctx.data["elevation_m"] = z0.astype(np.float32)
        out = ER.run(ctx)
        self.assertIn("ocean", out)
        self.assertFalse(out["ocean"][d < 300].any())
        self.assertTrue(out["ocean"][d > 7000].all())

    def test_stage_adds_the_elevation_stages_added_land(self):
        ctx, z0 = mountain_ctx()
        ctx.cfg["build"]["land_fraction"] = 0.02
        g = ctx.grid
        mtn = z0 > 0
        ctx.data.update({"continental": mtn, "continental_base": mtn})   # lifted shelf: no new crust at all
        added = mtn & (g.lon >= 0)
        ctx.data["land_added"] = added
        extra = g.area_km2[added].sum() / g.area_km2.sum()
        z = ER.run(ctx)["elevation_eroded_m"]
        self.assertAlmostEqual(g.area_km2[z > 0].sum() / g.area_km2.sum(), 0.02 + extra, delta=0.005)

class BaseLevelTest(unittest.TestCase):
    def test_erosion_keeps_coastline(self):
        ctx = make_ctx(3)
        g = ctx.grid
        block = (np.abs(g.lat) < 30) & (np.abs(g.lon) < 50)
        z0 = np.where(block, 800.0 + 2500.0 * np.exp(-((g.lon + 45) / 6.0) ** 2), -6000.0)   # coastal range by a deep sea
        z1 = ER.erode(g, z0, np.full(g.n, 1.5), ER.DEFAULTS)
        np.testing.assert_array_equal(z1 > 0, z0 > 0)
        self.assertTrue(np.all(z1[~block] == z0[~block]))


class HillslopeResolutionTest(unittest.TestCase):
    def test_hillslope_smoothing_is_resolution_independent(self):
        from tests.helpers import small_grid
        drops = []
        for res in (3, 4):
            g = small_grid(res)
            c = g.xyz[g.cell_index(20.0, 0.0)]
            d = g.radius_km * np.arccos(np.clip(g.xyz @ c, -1, 1))
            z0 = np.where(g.lat > -40, 500.0 + 3000.0 * np.exp(-(d / 600.0) ** 2), -4000.0)
            P = dict(ER.DEFAULTS, k=0.0)                       # isolate hillslope smoothing
            z1 = ER.erode(g, z0, np.ones(g.n), P)
            drops.append(z0.max() - z1[d < 150].max())
        self.assertAlmostEqual(drops[0] / max(drops[1], 1e-9), 1.0, delta=0.35)


class SedimentFillTest(unittest.TestCase):
    def test_shallow_basins_fill_deep_ones_remain(self):
        g = make_ctx(3).grid
        z0 = np.where(g.lat > -20, 200.0 + 2.0 * (g.lat + 20), -3000.0)
        depth = {}
        for name, (lat, dz) in {"shallow": (20.0, 60.0), "deep": (50.0, 2500.0)}.items():
            c = g.xyz[g.cell_index(lat, 0.0)]
            d = g.radius_km * np.arccos(np.clip(g.xyz @ c, -1, 1))
            z0 = np.where(d < 700, z0 - dz * (1 - d / 700), z0)
            depth[name] = d < 250
        from mapgen.graph import priority_flood
        def fill_after(frac):
            z1 = ER.erode(g, z0, np.full(g.n, 1.0), dict(ER.DEFAULTS, k=0.0, sediment_fill=frac))
            return priority_flood(g, z1, z1 <= -1000) - z1
        self.assertGreater(fill_after(0.0)[depth["shallow"]].max(), 20.0)     # without sediment it stays a basin
        fill = fill_after(ER.DEFAULTS["sediment_fill"])
        self.assertLess(fill[depth["shallow"]].max(), 0.2 * 60.0)             # noise-scale basins mostly filled
        self.assertGreater(fill[depth["deep"]].max(), 20.0)