import shutil
import importlib.util
import tempfile
import unittest
from pathlib import Path

import numpy as np

from mapgen import eras as ER, events as EV, pipeline as P
from mapgen.graph import components, distance_to
from mapgen.sphere import gc_dist_km, latlon_to_xyz, xyz_to_latlon
from tests.test_render import small_world

R = 12742.0
ERAS = """
[[event]]
name = "cut"
kind = "disintegrate"
center = {cut}
radius_km = 3000.0
depth_m = 3000.0
[[event]]
name = "aura"
kind = "zone"
shape = "landmass"
seed = {seed}
reach_km = [500.0, 1000.0]
edge_km = 5.0
fields = {{ gravity_g = 0.35, pressure_bar = 2.0, o2_fraction = 0.35, fire_reactivity = {fire} }}
[[event]]
name = "rise"
kind = "volcano"
center = {vol}
radius_km = 900.0
peak_m = 1500.0
peak_mode = "absolute"
[[event]]
name = "ghost"
kind = "zone"
shape = "landmass"
seed = {seed}
reach_km = [100.0, 200.0]
edge_km = 5.0
fields = {{ fire_reactivity = 0.2 }}
[eras]
order = ["before", "glow", "after", "risen", "haunt"]
default = "after"
[eras.before]
label = "Before"
events = []
[eras.glow]
label = "Glow"
events = ["aura"]
[eras.after]
label = "After"
events = ["cut"]
[eras.risen]
label = "Risen"
events = ["rise"]
years = 8000
[eras.haunt]
label = "Haunt"
events = ["ghost"]
"""


def cli():
    spec = importlib.util.spec_from_file_location("mapgen_cli", Path(ER.__file__).parents[1] / "mapgen.py")
    m = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(m)                          # the script (mapgen/ is the package)
    return m


def latlon(p):
    lat, lon = xyz_to_latlon(np.asarray(p, dtype=np.float64))
    return [round(float(lat), 3), round(float(lon), 3)]


class ErasTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.root = Path(tempfile.mkdtemp())
        small_world(cls.root)
        cls.base = P.build(cls.root, 2, log=lambda m: None)
        g, d = cls.base.grid, cls.base.data
        land = ~np.asarray(d["ocean"])
        lab = components(g, land)
        big = lab == np.bincount(lab[land]).argmax()
        mid = g.xyz[big].mean(axis=0)
        cls.cut = latlon(g.xyz[big][np.argmax(g.xyz[big] @ (mid / np.linalg.norm(mid)))])
        far = gc_dist_km(g.xyz, latlon_to_xyz(*cls.cut), R)
        cls.seed = latlon(g.xyz[np.argmax(np.where(big, far, -1.0))])     # stays land after the cut
        cls.vol_i = int(np.argmax(distance_to(g, land)))                   # the loneliest sea
        cls.vol = latlon(g.xyz[cls.vol_i])
        cls.g0 = float(cls.base.cfg["planet"]["gravity_g"])
        cls.write_eras()
        for n in ("glow", "after", "risen", "haunt"):
            ER.build_era(cls.root, 2, n, log=lambda m: None, base=cls.base)

    def test_low_memory_era_writes_identical_files(self):
        # oracle: the normal-mode era built in setUpClass
        tmp = Path(tempfile.mkdtemp())
        try:
            shutil.copytree(self.root, tmp, dirs_exist_ok=True)
            shutil.rmtree(ER.era_dir(tmp, 2, "glow"))
            ER.build_era(tmp, 2, "glow", log=lambda m: None, low_memory=True)
            a, b = ER.era_dir(self.root, 2, "glow"), ER.era_dir(tmp, 2, "glow")
            fa = sorted(p.relative_to(a) for p in a.rglob("*") if p.is_file())
            self.assertEqual(fa, sorted(p.relative_to(b) for p in b.rglob("*") if p.is_file()))
            for f in fa:
                self.assertEqual((a / f).read_bytes(), (b / f).read_bytes(), str(f))
        finally:
            shutil.rmtree(tmp)

    @classmethod
    def write_eras(cls, fire=0.5):
        (cls.root / "config" / "eras.toml").write_text(ERAS.format(cut=cls.cut, seed=cls.seed, vol=cls.vol, fire=fire))

    def load(self, name=None):
        d = self.root / "out" / "r2" if name is None else ER.era_dir(self.root, 2, name)
        with np.load(d / "cells.npz") as z:
            return {k: z[k] for k in z.files}

    def test_outside_the_mask_is_the_base(self):
        b = self.load()
        for n in ("glow", "after", "risen"):
            e = self.load(n)
            m = e["era_mask"]
            self.assertTrue(m.any() and not m.all(), n)
            for k, v in b.items():
                if v.shape[:1] == m.shape:
                    np.testing.assert_array_equal(e[k][~m], v[~m], err_msg=f"{n}: {k}")

    def test_deposits_are_the_base_geology(self):
        b = self.load()                                      # an event moves ground; it neither makes nor unmakes ore
        for n in ("glow", "after", "risen"):
            e = self.load(n)
            for k in ("deposits", "deposit_main", "iron_potential"):
                np.testing.assert_array_equal(e[k], b[k], err_msg=f"{n}: {k}")

    def test_disintegration_only_lowers_and_reaches_depth(self):
        b, e = self.load(), self.load("after")
        zb, ze = b["elevation_eroded_m"].astype(np.float64), e["elevation_eroded_m"].astype(np.float64)
        self.assertTrue(np.all(ze <= zb + 1e-3))
        d = gc_dist_km(b["g_xyz"], latlon_to_xyz(*self.cut), R)
        _, z_ref = EV.disintegrate(b["g_xyz"], zb, ~b.get("open_water", b["ocean"]), {"center": self.cut, "radius_km": 3000.0,
                                                                "depth_m": 3000.0}, R)
        inner = d < 1500.0
        self.assertTrue(inner.any())
        self.assertTrue(np.all(ze[inner] <= z_ref - 3000.0 * np.sqrt(0.75) + 1e-3))
        self.assertTrue((e["ocean"] | e["lake"])[inner].all(), "the bowl fills with water")

    def test_zone_fields_hold_inside_the_landmass(self):
        b, e = self.load(), self.load("glow")
        land, w = e["zone_aura_land"], e["zone_aura"]
        self.assertTrue(land.any())
        self.assertTrue(np.all(w[land] == 1.0))
        np.testing.assert_allclose(e["gravity_g"][land], 0.35, rtol=1e-6)
        np.testing.assert_allclose(e["fire_reactivity"][land], 0.5)
        np.testing.assert_allclose(e["po2_bar"], e["o2_fraction"] * e["pressure_bar"], rtol=1e-5)
        np.testing.assert_array_equal(e["gravity_g"][w == 0], b["gravity_g"][w == 0])

    def test_zone_follows_the_landmass_before_the_era(self):
        b, e = self.load(), self.load("after")
        d = gc_dist_km(b["g_xyz"], latlon_to_xyz(*self.cut), R)
        lost = ~b["open_water"] & e["open_water"] & (d < 1500.0)
        self.assertTrue(lost.any())
        self.assertTrue(e["zone_aura_land"][lost].all(), "the old coast, not the bitten one")
        np.testing.assert_allclose(e["gravity_g"][lost], 0.35, rtol=1e-6)

    def test_plants_in_a_zone(self):
        from mapgen import environment as EN
        e = self.load("glow")
        land = e["zone_aura_land"] & ~e["ocean"]
        zone, _ = EN.holdridge(e["biotemp"], e["P_ann"] * np.sqrt(2.0), e["T_min"])
        np.testing.assert_array_equal(e["holdridge"][land], zone[land])
        np.testing.assert_allclose(e["plant_height_x"], self.g0 / e["gravity_g"], rtol=1e-5)

    def test_era_years_set_its_erosion(self):
        e = self.load("risen")
        self.assertGreater(float(e["elevation_eroded_m"][self.vol_i]), 1400.0, "8,000 years barely wear a cone")
        self.assertNotEqual(ER.fingerprint("k", [("a", [], 8000)]), ER.fingerprint("k", [("a", [], 9000)]))

    def test_steps_list_each_eras_own_events_and_years(self):
        from mapgen import config as C
        st = ER.steps(C.load(self.root)[1], "risen")
        self.assertEqual([(n, [e["name"] for e in ev], y) for n, ev, y in st],
                         [("before", [], None), ("glow", ["aura"], None), ("after", ["cut"], None),
                          ("risen", ["rise"], 8000)])

    def test_a_later_zone_follows_the_land_of_its_own_era(self):
        b, e = self.load(), self.load("haunt")
        d = gc_dist_km(b["g_xyz"], latlon_to_xyz(*self.cut), R)
        lost = ~b["open_water"] & self.load("after")["open_water"] & (d < 1500.0)
        self.assertTrue(lost.any())
        self.assertFalse(e["zone_ghost_land"][lost].any(), "ghost came after the cut: the bitten coast")
        self.assertTrue(e["zone_aura_land"][lost].all(), "aura still follows the land before the cut")

    def test_a_new_label_needs_no_rebuild(self):
        import json
        f = ER.era_dir(self.root, 2, "glow") / "cells.npz"
        before = f.stat().st_mtime_ns
        p = self.root / "config" / "eras.toml"
        p.write_text(p.read_text().replace('label = "Glow"', 'label = "The Glow"'))
        try:
            logs = []
            ER.build_era(self.root, 2, "glow", log=logs.append)
            self.assertIn("era glow: cached", logs)
            meta = json.loads((ER.era_dir(self.root, 2, "glow") / "cells_meta.json").read_text())
            self.assertEqual(meta["era"]["label"], "The Glow")
            self.assertEqual(f.stat().st_mtime_ns, before)
        finally:
            self.write_eras()
            ER.build_era(self.root, 2, "glow", log=lambda m: None)

    def test_a_renamed_era_whose_folder_moved_needs_no_rebuild(self):
        import json
        old, new = ER.era_dir(self.root, 2, "glow"), ER.era_dir(self.root, 2, "shine")
        before = (old / "cells.npz").stat().st_mtime_ns
        p = self.root / "config" / "eras.toml"
        p.write_text(p.read_text().replace('"glow"', '"shine"').replace("[eras.glow]", "[eras.shine]"))
        old.rename(new)
        try:
            logs = []
            ER.build_era(self.root, 2, "shine", log=logs.append)
            self.assertIn("era shine: cached", logs)
            self.assertEqual(json.loads((new / "cells_meta.json").read_text())["era"]["name"], "shine")
            self.assertEqual((new / "cells.npz").stat().st_mtime_ns, before)
        finally:
            new.rename(old)
            self.write_eras()
            ER.build_era(self.root, 2, "glow", log=lambda m: None)

    def test_one_lake_id_never_names_two_lakes(self):
        b, e = self.load(), self.load("after")
        m, ids = e["era_mask"], e["lake_id"]
        for k in np.unique(ids[m & (ids >= 0)]).tolist():
            if (ids[~m] == k).any():
                np.testing.assert_array_equal(ids == k, b["lake_id"] == k, err_msg=f"lake {k}")

    def test_zone_events_never_change_height(self):
        b, e = self.load(), self.load("glow")
        for k in ("elevation_m", "elevation_eroded_m", "z_surface_m"):
            np.testing.assert_array_equal(e[k], b[k], err_msg=k)

    def test_a_volcano_raises_new_land(self):
        e = self.load("risen")
        self.assertFalse(e["ocean"][self.vol_i])
        self.assertGreater(float(e["elevation_eroded_m"][self.vol_i]), 0.0)

    def test_era_is_cached_until_its_events_change(self):
        logs = []
        ER.build_era(self.root, 2, "glow", log=logs.append)
        self.assertIn("era glow: cached", logs)
        first = self.load("glow")
        self.write_eras(fire=0.6)
        try:
            logs, base_logs = [], []
            ER.build_era(self.root, 2, "glow", log=logs.append)
            self.assertNotIn("era glow: cached", logs)
            P.build(self.root, 2, log=base_logs.append)
            self.assertTrue(base_logs and all(line.endswith("cached") for line in base_logs),
                            "tuning an event never rebuilds the base")
        finally:
            self.write_eras()
            ER.build_era(self.root, 2, "glow", log=lambda m: None)
        again = self.load("glow")                             # repeatable: the same inputs, the same arrays
        for k in first:
            np.testing.assert_array_equal(again[k], first[k], err_msg=k)

    def test_build_command_builds_eras_and_removes_stale_ones(self):
        stale = ER.era_dir(self.root, 2, "gone")
        stale.mkdir(parents=True, exist_ok=True)
        self.assertEqual(cli().main(["build", "--res", "2"], root=self.root), 0)
        self.assertFalse(stale.exists())
        for n in ("glow", "after", "risen"):
            self.assertTrue((ER.era_dir(self.root, 2, n) / "cells.npz").exists(), n)
        self.assertFalse(ER.era_dir(self.root, 2, "before").exists(), "the base era is the base build")
        self.assertEqual(cli().main(["era", "glow", "--res", "2"], root=self.root), 0)
        self.assertEqual(cli().main(["era", "nowhen", "--res", "2"], root=self.root), 2)


class MergeLakeIdsTest(unittest.TestCase):
    def test_kept_lakes_keep_ids_changed_and_new_lakes_get_fresh_ones(self):
        base = np.array([0, 0, 1, 1, -1, 2, -1], np.int32)
        new = np.array([5, 5, 0, 0, 0, -1, 1], np.int32)   # 5 = base 0 as it was; 0 = base 1 grown; 1 new; 2 gone
        out = ER.merge_lake_ids(base, new, np.ones(7, bool))
        np.testing.assert_array_equal(out, [0, 0, 3, 3, 3, -1, 4])
        self.assertEqual(out.dtype, np.int32)

    def test_outside_the_mask_stays_the_base(self):
        base = np.array([0, 0, -1, 1])
        new = np.array([3, 3, 7, 7])
        out = ER.merge_lake_ids(base, new, np.array([False, False, True, True]))
        np.testing.assert_array_equal(out, [0, 0, 2, 2])


class SameTest(unittest.TestCase):
    def test_nan_equals_nan_in_float_arrays_only(self):
        self.assertTrue(ER._same(np.array([np.nan, 1.0]), np.array([np.nan, 1.0])))
        self.assertFalse(ER._same(np.array([np.nan, 1.0]), np.array([np.nan, 2.0])))
        self.assertTrue(ER._same(np.array([1, 2]), np.array([1, 2])))
        self.assertTrue(ER._same(np.array([True]), np.array([True])))