import unittest

import numpy as np

import refine as RF
import seafloor as SF
from mapgen.graph import priority_flood
from mapgen.sphere import gc_dist_km, latlon_to_xyz
from tests.test_refine import square

R = 12742.0


def grid(lat=0.0, lon=0.0, d=4.0):
    (cells,) = RF.area_cells([square(lat, lon, d)], 5)
    return RF.subgrid(cells, R)


def shelf():
    """A res-5 square across a coast: land west of 3° W, a shelf, a slope with gullies, an abyssal plain with a deep
    hollow (≈ 300 m, at 0°, 2.5° E) and a shallow one (≈ 30 m, at 2.5° S, 1.5° E)."""
    g, halo = grid()
    x = g.lon
    z = np.where(x < -3.0, 200.0,
                 np.where(x < -1.5, -100.0 - 100.0 * (x + 3.0) / 1.5,
                          np.where(x < 0.0, -200.0 - 3800.0 * (x + 1.5) / 1.5, -4000.0 - 75.0 * x)))
    z = z + np.where((x >= -1.5) & (x < 0.0), 40.0 * np.sin(np.radians(g.lat) * 400.0), 0.0)   # gullies
    for la, lo, r_km, depth in ((0.0, 2.5, 60.0, 300.0), (-2.5, 1.5, 40.0, 30.0)):
        d = gc_dist_km(g.xyz, latlon_to_xyz(la, lo), R)
        z = z - depth * np.clip(1.0 - (d / r_km) ** 2, 0.0, 1.0)
    return g, halo, z, z <= 0


class CanyonTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.g, cls.halo, cls.z, cls.sea = shelf()
        cls.zc, cls.cut, cls.fan = SF.canyons(cls.g, cls.z, cls.sea, cls.halo)

    def test_canyons_cut_the_slope(self):
        self.assertGreater(self.cut.max(), 50.0)
        self.assertLessEqual(self.cut.max(), SF.CANYON["cap_m"])
        slope = (self.g.lon > -1.5) & (self.g.lon < 0.0)
        self.assertGreater(self.cut[slope].max(), self.cut[~slope].max(initial=0.0))

    def test_only_lower_in_canyons_and_raise_at_most_100_m_on_fans(self):
        dz = self.zc - self.z
        self.assertGreaterEqual(dz.min(), -SF.CANYON["cap_m"] - 1e-9)
        self.assertLessEqual(dz.max(), SF.FAN["max_m"] + 1e-9)
        np.testing.assert_array_equal(dz < 0, self.cut > 0)
        self.assertTrue(np.all(self.fan[dz > 0] > 0))
        self.assertGreater(self.fan.max(), 0.0, "fans where the flows stop")

    def test_land_and_halo_stay(self):
        keep = ~self.sea | self.halo
        np.testing.assert_array_equal(self.zc[keep], self.z[keep])
        self.assertTrue(np.all(self.zc[self.sea] <= np.maximum(self.z[self.sea], SF.SEA_TOP_M)))

    def test_deterministic(self):
        again = SF.canyons(self.g, self.z, self.sea, self.halo)
        for a, b in zip(again, (self.zc, self.cut, self.fan)):
            np.testing.assert_array_equal(a, b)


class PondTest(unittest.TestCase):
    def setUp(self):
        self.g, self.halo, self.z, self.sea = shelf()
        self.sed = np.full(self.g.n, 500.0)

    def test_deep_hollows_fill_flat_below_their_spill(self):
        zp, fill = SF.ponds(self.g, self.z, self.sea, self.halo, self.sed)
        zf = priority_flood(self.g, np.where(self.sea, self.z, SF.WALL_M), self.sea & self.halo, eps=0.0)
        self.assertTrue(np.all(zp <= np.maximum(zf, self.z) + 1e-9), "never above the spill level")
        self.assertTrue(np.all(zp >= self.z))
        np.testing.assert_allclose(fill, zp - self.z)
        deep = gc_dist_km(self.g.xyz, latlon_to_xyz(0.0, 2.5), R) < 30.0
        self.assertTrue(np.all(fill[deep] > 0))
        self.assertLess(float(np.ptp(zp[deep])), 1e-6, "a flat floor")
        self.assertAlmostEqual(float(zp[deep][0]), float(self.z[deep].min()) + 0.3 * 500.0, delta=1.0)
        shallow = gc_dist_km(self.g.xyz, latlon_to_xyz(-2.5, 1.5), R) < 20.0
        self.assertTrue(np.all(fill[shallow] == 0), "hollows under 50 m stay as they are")

    def test_no_sea_is_a_no_op(self):
        z = np.full(self.g.n, 300.0)
        sea = np.zeros(self.g.n, bool)
        zc, cut, fan = SF.canyons(self.g, z, sea, self.halo)
        zp, fill = SF.ponds(self.g, z, sea, self.halo, self.sed)
        np.testing.assert_array_equal(zc, z)
        np.testing.assert_array_equal(zp, z)
        self.assertFalse(cut.any() or fan.any() or fill.any())

    def test_sea_without_halo_sea_still_ponds(self):
        d = gc_dist_km(self.g.xyz, latlon_to_xyz(0.0, 0.0), R)
        z = np.where(d < 200.0, -1000.0 - 300.0 * np.clip(1.0 - (d / 100.0) ** 2, 0.0, 1.0), 500.0)
        sea = z <= 0
        self.assertFalse((sea & self.halo).any())
        zp, fill = SF.ponds(self.g, z, sea, self.halo, self.sed)
        zc, _, _ = SF.canyons(self.g, z, sea, self.halo)
        self.assertTrue(np.all(np.isfinite(zp)) and np.all(np.isfinite(zc)))


from mapgen import plateaus as PL, seabed as SB

PLAT = {"name": "p", "center": [0.0, 0.0], "area_km2": 2.0e5, "top_m": [1500.0, 2500.0]}


def world_values(n=1):
    return {"vent_potential": np.full(n, 0.8), "seabed_type": np.full(n, SB.SB_CONTINENTAL, np.int8),
            "seabed_mineral": np.zeros(n, np.int8), "bottom_temp_c": np.full(n, 3.0),
            "sediment_m": np.full(n, 400.0), "d_div_km": np.full(n, 900.0), "T_mean": np.full(n, 10.0)}


class CellDrawTest(unittest.TestCase):
    def test_cell_draws_are_uniform_and_stable(self):
        ids = np.arange(10 ** 5, dtype=np.uint64) * np.uint64(7919) + np.uint64(612345678901)
        u = SF._u(ids, 3)
        self.assertTrue(0.0 <= u.min() and u.max() < 1.0)
        self.assertAlmostEqual(float(u.mean()), 0.5, delta=0.01)
        np.testing.assert_array_equal(SF._u(ids, 3), u)
        self.assertFalse(np.array_equal(SF._u(ids, 4), u))


class VolcanicTest(unittest.TestCase):
    def test_fine_features_on_the_plateau_only(self):
        g, halo = grid()
        pid = PL.cell_ids(g.xyz, [PLAT], 7, R)
        add, near = SF.volcanic(g, np.ones(g.n, bool), halo, pid, [PLAT], 7)
        self.assertGreater(add.max(), 150.0)
        self.assertTrue(np.all(add[pid < 0] == 0) and np.all(near[pid < 0] == 0))
        self.assertTrue(np.all(add[halo] == 0))

    def test_features_and_vents_do_not_depend_on_the_area(self):
        outs = []
        for g, h in (grid(0.0, 0.0, 4.0), grid(0.5, 0.5, 3.0)):
            pid = PL.cell_ids(g.xyz, [PLAT], 7, R)
            sea = np.ones(g.n, bool)
            add, near = SF.volcanic(g, sea, h, pid, [PLAT], 7)
            _, v = SF.vents(g, np.full(g.n, -2000.0) + add, sea, h, np.full(g.n, 0.8), near, np.full(g.n, 500.0),
                            np.full(g.n, 100.0), np.full(g.n, 2.0), 7)
            outs.append((g, h, add, v))
        (ga, ha, aa, va), (gb, hb, ab, vb) = outs
        ca, cb = np.flatnonzero(~ha), np.flatnonzero(~hb)
        common, ia, ib = np.intersect1d(ga.ids[ca], gb.ids[cb], return_indices=True)
        self.assertGreater(len(common), 100)
        np.testing.assert_allclose(aa[ca[ia]], ab[cb[ib]], atol=1e-9)
        np.testing.assert_array_equal(np.intersect1d(va["cell"], common), np.intersect1d(vb["cell"], common))

    def test_hidden_plateaus_stay_500_m_down_and_no_sea_rises_to_land(self):
        g, halo = grid()
        z = np.full(g.n, -600.0)
        sea = np.ones(g.n, bool)
        zn, f, v = SF.run(g, z, sea, halo, np.zeros(g.n, np.int64), world_values(), [PLAT], 7)
        on = (f["plateau_id"] >= 0) & ~halo
        self.assertTrue(on.any())
        self.assertLessEqual(float(zn[on].max()), -500.0)
        isl = {**PLAT, "islands": True}
        zi, fi, _ = SF.run(g, z, sea, halo, np.zeros(g.n, np.int64), world_values(), [isl], 7)
        self.assertLessEqual(float(zi.max()), SF.SEA_TOP_M)
        np.testing.assert_array_equal(zn[halo], z[halo])


class VentTest(unittest.TestCase):
    def test_vents_only_in_the_sea_and_smokers_on_the_ridge(self):
        g, halo = grid()
        z = np.where(g.lon < -3.0, 100.0, -2500.0)
        sea = z <= 0
        d_div = np.abs(g.lon - 2.0) * 111.0                    # a ridge axis along 2° E
        pot = np.where(sea, np.clip(1.0 - d_div / 300.0, 0.05, 1.0), 0.0)
        st, v = SF.vents(g, z, sea, halo, pot, np.zeros(g.n), d_div, np.full(g.n, 1500.0), np.full(g.n, 2.0), 7)
        cells = np.searchsorted(g.ids, v["cell"])
        self.assertGreater(len(cells), 0)
        self.assertTrue(np.all(sea[cells] & ~halo[cells] & (z[cells] < -SF.VENT["min_depth_m"])))
        self.assertTrue(np.all(v["type"][d_div[cells] < SF.VENT["ridge_km"]] == 0), "black smokers on the ridge")
        t, temp = v["type"], v["temp_c"]
        self.assertTrue(np.all((temp[t == 0] >= 300) & (temp[t == 0] <= 400)))
        self.assertTrue(np.all(st[~sea] == 0) and st.max() <= 1.0)
        np.testing.assert_array_equal(v["mineral"], v["type"])


class FieldsTest(unittest.TestCase):
    def test_fields_zero_on_land_and_features_mark_the_floor(self):
        pv = {"seabed_type": np.array([SB.SB_CLAY, SB.SB_NONE, SB.SB_CLAY, SB.SB_CLAY, SB.SB_NONE], np.int8),
              "seabed_mineral": np.zeros(5, np.int8), "bottom_temp_c": np.array([2.0, 0.0, 2.0, 2.0, 0.0]),
              "sediment_m": np.array([400.0, 0.0, 400.0, 400.0, 0.0]), "T_mean": np.array([10.0, 12.0, 10.0, 10.0, 12.0])}
        sea = np.array([True, True, True, True, False])
        z0 = np.zeros(5)
        f = SF.fields(pv, sea, np.zeros(5, bool), np.array([0.0, 0, 100, 0, 0]), np.array([0.0, 0, 0, 5, 0]), z0,
                      np.array([0.0, 0, 0, 0, 0]), np.array([0.9, 0, 0, 0, 0]))
        np.testing.assert_array_equal(f["seabed_type"], [SB.SB_VENTS, SB.SB_TERRIGENOUS, SB.SB_TERRIGENOUS,
                                                         SB.SB_TERRIGENOUS, SB.SB_NONE])
        self.assertEqual(int(f["seabed_mineral"][0]), SB.MI_SULFIDES)
        np.testing.assert_allclose(f["sediment_m"], [0.0, 300.0, 120.0, 405.0, 0.0])
        np.testing.assert_allclose(f["bottom_temp_c"], [20.0, 12.0, 2.0, 2.0, 0.0])
        np.testing.assert_allclose(f["vent"], [0.9, 0, 0, 0, 0], rtol=1e-6)
        np.testing.assert_allclose(f["canyon_m"], [0, 0, 100, 0, 0])