import unittest

import numpy as np

from mapgen import events as EV
from mapgen.pipeline import StageError
from mapgen.sphere import east_north, gc_dist_km, great_circle_point, latlon_to_xyz
from tests.helpers import small_grid

R = 12742.0
CUT = {"name": "cut", "kind": "disintegrate", "center": [0.0, 0.0], "radius_km": 3000.0, "depth_m": 3000.0}


class DisintegrateTest(unittest.TestCase):
    def test_bowl_only_lowers_and_reaches_its_depth(self):
        g = small_grid(3)
        d = gc_dist_km(g.xyz, latlon_to_xyz(0.0, 0.0), R)
        z = np.where(d < 4000, 800.0, -4000.0)
        z2, z_ref = EV.disintegrate(g.xyz, z, z > 0, CUT, R)
        self.assertAlmostEqual(z_ref, 800.0)
        self.assertTrue(np.all(z2 <= z))
        np.testing.assert_array_equal(z2[d >= 3000], z[d >= 3000])
        c = int(np.argmin(d))
        self.assertAlmostEqual(z2[c], 800.0 - 3000.0 * np.sqrt(1 - (d[c] / 3000.0) ** 2), places=6)


    def test_disintegration_over_open_sea_has_no_rim_height(self):
        g = small_grid(3)
        z2, z_ref = EV.disintegrate(g.xyz, np.full(g.n, -1000.0), np.zeros(g.n, bool), CUT, R)
        self.assertIsNone(z_ref)
        self.assertLess(float(z2[g.cell_index(0.0, 0.0)]), -2990.0, "the bowl hangs from 0 m")

    def test_heights_warn_when_a_cut_has_no_rim_land(self):
        from mapgen import eras as ER
        g = small_grid(3)
        logs = []
        ER._heights(g, np.full(g.n, -1000.0), [dict(CUT)], 5.0e6, logs.append, "x")
        self.assertTrue(any("no land in its rim" in m for m in logs), logs)


class VolcanoTest(unittest.TestCase):
    def test_shapes_and_modes(self):
        g = small_grid(3)
        d = gc_dist_km(g.xyz, latlon_to_xyz(0.0, 90.0), R)
        c = int(np.argmin(d))
        z = np.full(g.n, -3000.0)
        above = EV.volcano(g.xyz, z, {"center": [0.0, 90.0], "radius_km": 900.0, "peak_m": 2000.0}, R)
        self.assertTrue(np.all(above >= z))
        self.assertAlmostEqual(above[c], -3000.0 + 2000.0 * (1 - d[c] / 900.0) ** 1.5)
        absolute = EV.volcano(g.xyz, z, {"center": [0.0, 90.0], "radius_km": 900.0, "peak_m": 1500.0,
                                         "peak_mode": "absolute"}, R)
        self.assertGreater(absolute[c], 0.0, "a volcano above 0 m makes land")
        np.testing.assert_array_equal(absolute[d >= 900.0], z[d >= 900.0])
        cal = EV.volcano(g.xyz, np.zeros(g.n), {"center": [0.0, 90.0], "radius_km": 2000.0, "peak_m": 2000.0,
                                                "shape": "caldera"}, R)
        ring = (d > 300.0) & (d < 700.0)
        self.assertLess(cal[c], cal[ring].max(), "a caldera's middle sits below its rim")


class ZoneWeightTest(unittest.TestCase):
    def test_sharp_circle_edge_ramps_over_edge_km(self):
        ev = {"name": "z", "shape": "circle", "center": [10.0, 20.0], "radius_km": 800.0, "edge_km": 5.0}
        p0 = latlon_to_xyz(10.0, 20.0)
        e, n = east_north(p0[None])
        s = np.arange(780.0, 810.0, 0.5)
        pts = np.array([great_circle_point(p0, e[0], si, R) for si in s])
        w = EV.zone_weight(pts, ev, R, 1296)
        self.assertTrue(np.all(w[s <= 794.5] == 1.0))
        self.assertTrue(np.all(w[s >= 800.5] == 0.0))
        ramp = s[(w > 0) & (w < 1)]
        self.assertLessEqual(ramp.max() - ramp.min(), 5.0)

    def test_landmass_reach_varies_between_its_bounds(self):
        g = small_grid(3)
        land = (np.abs(g.lat) < 20) & (np.abs(g.lon) < 30)
        ev = {"name": "aura", "shape": "landmass", "seed": [0.0, 0.0], "reach_km": [500.0, 1000.0], "edge_km": 5.0}
        lm = EV.landmass(g, land, ev["seed"], ev["name"])
        np.testing.assert_array_equal(lm, land)
        w = EV.zone_weight(g.xyz, ev, R, 1296, g.xyz[lm])
        self.assertTrue(np.all(w[land] == 1.0))
        from scipy.spatial import cKDTree
        chord, _ = cKDTree(g.xyz[lm]).query(g.xyz)
        d = 2 * np.arcsin(chord / 2) * R
        self.assertTrue(np.all(w[d > 1000.0] == 0.0))
        self.assertTrue(np.all(w[(d > 0) & (d < 495.0)] == 1.0))

    def test_landmass_seed_in_the_sea_names_the_event(self):
        g = small_grid(3)
        land = (np.abs(g.lat) < 20) & (np.abs(g.lon) < 30)
        with self.assertRaises(StageError) as cm:
            EV.landmass(g, land, [0.0, 150.0], "impact-aftermath")
        self.assertIn("impact-aftermath", str(cm.exception))
        self.assertIn("nearest land", str(cm.exception))


class ApplyZoneTest(unittest.TestCase):
    EV = {"name": "aura", "fields": {"gravity_g": 0.35, "pressure_bar": 2.0, "o2_fraction": 0.35,
                                     "fire_reactivity": 0.5}}

    def base(self, n=4):
        return {"gravity_g": np.full(n, 1.05, np.float32), "o2_fraction": np.array([0.21, 0.3, 0.1, 0.21]),
                "pressure_bar": np.array([1.0, 1.0, np.exp(-1), 1.0]), "fire_reactivity": np.ones(n),
                "po2_bar": np.array([0.21, 0.3, 0.1 * np.exp(-1), 0.21])}

    def test_zone_values_are_absolute_over_any_base(self):
        f = self.base()
        w = np.array([1.0, 1.0, 1.0, 0.0])
        z = np.array([0.0, -50.0, 8000.0, 0.0])
        out = EV.apply_zone(f, w, self.EV, z, 8000.0)
        np.testing.assert_allclose(out["o2_fraction"][:3], 0.35)
        np.testing.assert_allclose(out["gravity_g"][:3], 0.35, rtol=1e-6)
        np.testing.assert_allclose(out["pressure_bar"][:3], [2.0, 2.0, 2.0 * np.exp(-1)])
        np.testing.assert_allclose(out["po2_bar"], out["o2_fraction"] * out["pressure_bar"])
        self.assertEqual(out["gravity_g"].dtype, np.float32)
        for k in out:
            self.assertEqual(out[k][3], f[k][3], f"{k}: untouched outside the zone")

    def test_half_weight_blends(self):
        out = EV.apply_zone(self.base(), np.full(4, 0.5), self.EV, np.zeros(4), 8000.0)
        self.assertAlmostEqual(float(out["fire_reactivity"][0]), 0.75)
        self.assertAlmostEqual(float(out["pressure_bar"][0]), 1.5)