import unittest

import numpy as np

from mapgen import check as CK
from tests.helpers import small_grid


class CheckTest(unittest.TestCase):
    def setUp(self):
        self.g = small_grid(2)

    def test_land_fraction(self):
        a = np.ones(100)
        z = np.where(np.arange(100) < 22, 100.0, -100.0)
        self.assertIsNone(CK.check_land_fraction(z, a, 0.22))
        self.assertIn("land fraction", CK.check_land_fraction(z, a, 0.40))

    def test_hypsometry(self):
        rng = np.random.default_rng(0)
        z = np.concatenate([rng.normal(-4000, 600, 7000), rng.normal(500, 400, 3000)])
        self.assertIsNone(CK.check_hypsometry(z, np.ones(len(z))))
        self.assertIn("bimodal", CK.check_hypsometry(rng.normal(0, 3000, 10000), np.ones(10000)))

    def test_max_elevation(self):
        self.assertIsNone(CK.check_max_elevation(np.array([11000.0]), np.array([1.0])))
        self.assertIsNone(CK.check_max_elevation(np.array([20000.0]), np.array([0.5])))
        self.assertIn("12000", CK.check_max_elevation(np.array([13000.0]), np.array([1.0])))

    def test_drainage_and_uphill(self):
        recv = np.array([0, 0, 1, 3])
        z = np.array([-10.0, 5.0, 8.0, 20.0])
        endo = np.array([False, False, False, False])
        self.assertIn("sink", CK.check_drainage(recv, z, endo))
        endo[3] = True
        self.assertIsNone(CK.check_drainage(recv, z, endo))
        self.assertIsNone(CK.check_no_uphill(z, recv, z, endo))
        self.assertIn("uphill", CK.check_no_uphill(np.array([-10.0, 9.0, 8.0, 20.0]), recv, z, endo))

    def test_temperature_and_rain_bands(self):
        g = self.g
        self.assertIsNone(CK.check_temperature(g, 30 - 0.6 * np.abs(g.lat)))
        self.assertIn("poleward", CK.check_temperature(g, 0.3 * np.abs(g.lat)))
        wet = 2000 * np.exp(-(g.lat / 8) ** 2) + 900 * np.exp(-((np.abs(g.lat) - 40) / 8) ** 2) + 100
        self.assertIsNone(CK.check_rain_bands(g, wet, 20.0))
        self.assertIn("subtropic", CK.check_rain_bands(g, np.full(g.n, 800.0), 20.0))

    def test_drainage_terminals_must_hold_water_or_salt(self):
        recv = np.array([0, 0, 1, 3])
        z = np.array([-10.0, 5.0, 8.0, 20.0])
        endo = np.array([False, False, False, True])
        dry = np.zeros(4, bool)
        self.assertIn("lake", CK.check_drainage(recv, z, endo, terminal_ok=dry))
        wet = dry.copy()
        wet[3] = True
        self.assertIsNone(CK.check_drainage(recv, z, endo, terminal_ok=wet))

    def test_lake_surface_routing_is_exempt(self):
        recv = np.array([0, 0, 1, 2])
        zf = np.array([-10.0, 5.0, 5.02, 5.01])      # cell 3 → 2 is ε-uphill across a flat lake
        z = np.array([-10.0, 5.0, 4.0, 4.5])
        endo = np.zeros(4, bool)
        lake = np.array([False, False, True, True])
        self.assertIn("uphill", CK.check_no_uphill(zf, recv, z, endo))
        self.assertIsNone(CK.check_no_uphill(zf, recv, z, endo, lake=lake))


class RevisionCheckTest(unittest.TestCase):
    P1 = {"name": "east-flank", "center": [-39.9, -18.0], "area_km2": 3.0e6, "elongation": 1.9, "azimuth_deg": 30.0,
          "top_m": [1500.0, 3000.0]}
    P2 = {"name": "plateau-02", "center": [-0.4, 176.6], "area_km2": 1.0e6, "top_m": [1000.0, 1500.0],
          "islands": True}

    def world(self):
        from mapgen import plateaus as PL
        g = small_grid(4)
        ids = PL.cell_ids(g.xyz, [self.P1, self.P2], 7, g.radius_km)
        z = PL.apply(g, np.full(g.n, -6000.0), [self.P1, self.P2], ids, 7)
        return g, {"elevation_eroded_m": z, "ocean": z <= 0, "plateau_id": ids}

    def test_plateaus_pass_and_fail(self):
        g, d = self.world()
        self.assertEqual(CK.check_plateaus(g, d, [self.P1, self.P2], 7), [])
        z = d["elevation_eroded_m"].copy()
        z[np.flatnonzero(d["plateau_id"] == 0)[0]] = -100.0
        z[z > 0] = -700.0
        bad = CK.check_plateaus(g, {**d, "elevation_eroded_m": z, "ocean": z <= 0}, [self.P1, self.P2], 7)
        self.assertTrue(any("east-flank" in m and "−500" in m for m in bad), bad)
        self.assertTrue(any("plateau-02" in m and "island" in m for m in bad), bad)

    def test_zones_gradual_rule(self):
        from mapgen import zones as ZN
        g = small_grid(4)
        m = ZN.contribution(g.xyz, {"name": "R1", "field": "o2", "center": [0.0, 0.0], "radius_km": 3200.0,
                                    "v": 1.0}, 7, g.radius_km)
        smooth = {"o2_fraction": 0.21 * (1 + 0.5 * m), "gravity_g": np.full(g.n, 1.05)}
        self.assertEqual(CK.check_zones(g, smooth), [])
        step = {"o2_fraction": np.where(g.lon > 0, 0.35, 0.21), "gravity_g": np.full(g.n, 1.05)}
        self.assertTrue(any("o2_fraction" in m for m in CK.check_zones(g, step)))

    def test_compare_skips_o2_fields_inside_o2_zones(self):
        g = small_grid(3)
        tect = {"plateau": [], "land_patch": [],
                "zone": [{"name": "r", "field": "o2", "center": [0.0, 0.0], "radius_km": 2000.0, "v": 1.0}]}
        o2 = CK.o2_areas(g.xyz, tect, g.radius_km)
        self.assertTrue(o2[g.cell_index(0.0, 0.0)])
        self.assertFalse(o2[g.cell_index(0.0, 90.0)])
        ex = CK.edited_areas(g.xyz, tect, g.radius_km)
        self.assertFalse(ex.any(), "O₂ zones don't change relief")
        old = {"po2_bar": np.zeros(g.n), "T_mean": np.zeros(g.n)}
        new = {"po2_bar": np.where(o2, 0.1, 0.0), "T_mean": np.where(o2, 1.0, 0.0)}
        lines = CK.compare(old, new, ex, {k: o2 for k in CK.O2_FIELDS})
        self.assertIn("po2_bar: 0.0% changed, max 0, p99 0", lines)
        self.assertTrue(any(s.startswith("T_mean:") and not s.startswith("T_mean: 0.0%") for s in lines), lines)

    def test_compare_command_needs_a_world_cells_file(self):
        import importlib.util
        import tempfile
        from pathlib import Path
        from mapgen.testing import built_world
        spec = importlib.util.spec_from_file_location("mapgen_cli", Path(CK.__file__).parents[1] / "mapgen.py")
        cli = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(cli)
        old = Path(tempfile.mkdtemp()) / "other.npz"
        np.savez(old, a=np.zeros(3))
        self.assertEqual(cli.main(["compare", "--res", "2", "--old", str(old)], root=built_world()), 2)

    def test_compare_ignores_edited_areas(self):
        g = small_grid(3)
        tect = {"plateau": [self.P1], "land_patch": [], "zone": []}
        ex = CK.edited_areas(g.xyz, tect, g.radius_km)
        self.assertTrue(ex[g.cell_index(-39.9, -18.0)])
        self.assertFalse(ex[g.cell_index(40.0, 100.0)])
        old = {"a": np.zeros(g.n), "b": np.zeros(g.n), "g_ids": g.ids}
        new = {"a": np.where(ex, 5.0, 0.0), "b": np.where(ex, 0.0, 1.0), "g_ids": g.ids}
        lines = CK.compare(old, new, ex)
        self.assertIn("a: 0.0% changed, max 0, p99 0", lines)
        self.assertTrue(any(line.startswith("b: 100.0% changed, max 1") for line in lines), lines)

    def test_site_report_warns_about_land_nearby(self):
        from mapgen import plateaus as PL
        g, d = self.world()
        land = np.abs(g.lat - (-39.9)) < 2.0
        near = land & (np.abs(g.lon - (-18.0)) < 30.0) & (d["plateau_id"] < 0)
        data = {**d, "ocean": ~near, "bnd_type": np.zeros(g.n, np.int8)}
        msgs = PL.site_report(g, data, [self.P1, self.P2])
        self.assertTrue(any("east-flank" in m and "from land" in m for m in msgs), msgs)

    def test_compare_command(self):
        import importlib.util
        import shutil
        import tempfile
        from pathlib import Path
        from mapgen.testing import built_world
        spec = importlib.util.spec_from_file_location("mapgen_cli", Path(CK.__file__).parents[1] / "mapgen.py")
        cli = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(cli)
        old = Path(tempfile.mkdtemp()) / "cells.npz"
        shutil.copy(built_world() / "out" / "r2" / "cells.npz", old)
        self.assertEqual(cli.main(["compare", "--res", "2", "--old", str(old)], root=built_world()), 0)

    def test_compare_equal_infinities_are_unchanged(self):
        ex = np.zeros(3, bool)
        lines = CK.compare({"d": np.array([np.inf, 1.0, 2.0])}, {"d": np.array([np.inf, 1.0, 3.0])}, ex)
        self.assertIn("d: 33.3% changed, max 1, p99 0.98", lines)