import unittest

import numpy as np

from mapgen import environment as EN
from tests.helpers import make_ctx


def name(bio, p, tmin):
    z, _ = EN.holdridge(np.array([bio]), np.array([p]), np.array([tmin]))
    return EN.HOLDRIDGE_NAMES[z[0]]


class HoldridgeTest(unittest.TestCase):
    def test_count(self):
        self.assertEqual(len(EN.HOLDRIDGE_NAMES), 38)

    def test_table(self):
        self.assertEqual(name(26, 3000, 20), "tropical moist forest")
        self.assertEqual(name(26, 9000, 20), "tropical rain forest")
        self.assertEqual(name(26, 100, 20), "tropical desert")
        self.assertEqual(name(0.5, 100, -30), "polar desert")
        self.assertEqual(name(8, 300, -10), "cool temperate steppe")
        self.assertEqual(name(15, 1500, -5), "warm temperate moist forest")
        self.assertEqual(name(15, 1500, 5), "subtropical moist forest")
        self.assertEqual(name(4, 700, -20), "boreal wet forest")
        self.assertEqual(name(2, 200, -25), "subpolar moist tundra")

    def test_polar_desert_with_summer_monsoon(self):
        tag = EN.seasonality(np.array([180.0, 20.0]), np.array([20.0, 180.0]), np.array([80.0, 80.0]),
                             np.array([-30.0, -30.0]), np.array([100.0, 100.0]), EN.DEFAULTS)
        self.assertEqual(tag[0], EN.SEAS_W)   # NH: wet June (summer), dry December
        self.assertEqual(tag[1], EN.SEAS_S)   # NH: dry summer
        sh = EN.seasonality(np.array([20.0]), np.array([180.0]), np.array([-80.0]), np.array([-30.0]),
                            np.array([100.0]), EN.DEFAULTS)
        self.assertEqual(sh[0], EN.SEAS_W)    # SH summer is December

    def test_tropical_monsoon(self):
        tag = EN.seasonality(np.array([3000.0]), np.array([200.0]), np.array([15.0]), np.array([20.0]),
                             np.array([1800.0]), EN.DEFAULTS)
        self.assertEqual(tag[0], EN.SEAS_M)


class StageTest(unittest.TestCase):
    def test_run(self):
        ctx = make_ctx(3)
        g = ctx.grid
        n = g.n
        land = (np.abs(g.lat) < 40) & (np.abs(g.lon) < 60)
        z = np.where(land, 300.0, -4000.0)
        z = np.where(land & (np.abs(g.lon) < 4), 4000.0, z)
        zeros = np.zeros(n)
        ctx.data.update({
            "elevation_eroded_m": z, "continental": land, "age_class": np.where(land, 1, 0).astype(np.int8),
            "d_over_km": np.full(n, np.inf), "T_mean": np.where(np.abs(g.lat) > 30, -5.0, 20.0),
            "T_min": np.full(n, -10.0), "T_range": np.full(n, 20.0), "P_ann": np.full(n, 800.0),
            "P_jun": np.full(n, 800.0), "P_dec": np.full(n, 800.0), "biotemp": np.full(n, 10.0),
            "dist_ocean_km": np.where(land, 1000.0, 0.0), "river": np.zeros(n, bool),
            "strahler": np.zeros(n, np.int8), "discharge_km3_yr": zeros, "lake": np.zeros(n, bool),
            "salt_flat": np.zeros(n, bool)})
        out = EN.run(ctx)
        self.assertEqual(out["landform"][g.cell_index(0.0, 0.0)], EN.LF_MOUNTAINS)
        self.assertEqual(out["landform"][g.cell_index(0.0, 150.0)], EN.LF_OCEAN)
        self.assertEqual(out["ground"][g.cell_index(35.0, 40.0)], EN.GR_PERMAFROST)
        for k in ("holdridge", "hold_region", "seasonality", "landform", "relief_m", "lithology", "ground",
                  "coal_potential", "iron_potential"):
            self.assertEqual(len(out[k]), n)


class OceanMaskEnvTest(unittest.TestCase):
    def test_inland_basin_not_ocean_landform(self):
        ctx = make_ctx(3)
        g = ctx.grid
        n = g.n
        land = (np.abs(g.lat) < 40) & (np.abs(g.lon) < 60)
        z = np.where(land, 300.0, -4000.0)
        basin = (np.abs(g.lat) < 6) & (np.abs(g.lon) < 6)
        z[basin] = -30.0
        zeros = np.zeros(n)
        ctx.data.update({
            "elevation_eroded_m": z, "ocean": ~land, "continental": land, "age_class": np.where(land, 1, 0).astype(np.int8),
            "d_over_km": np.full(n, np.inf), "T_mean": np.full(n, 20.0), "T_min": np.full(n, 10.0),
            "T_range": np.full(n, 10.0), "P_ann": np.full(n, 800.0), "P_jun": np.full(n, 800.0),
            "P_dec": np.full(n, 800.0), "biotemp": np.full(n, 20.0), "dist_ocean_km": np.where(land, 1000.0, 0.0),
            "river": np.zeros(n, bool), "strahler": np.zeros(n, np.int8), "discharge_km3_yr": zeros,
            "lake": np.zeros(n, bool), "salt_flat": np.zeros(n, bool)})
        out = EN.run(ctx)
        self.assertFalse(np.any(out["landform"][basin] == EN.LF_OCEAN))


class ReliefResolutionTest(unittest.TestCase):
    def test_relief_window_is_in_km(self):
        from tests.helpers import small_grid
        med = []
        for res in (2, 3):
            g = small_grid(res)
            z = 1000.0 * np.sin(10.0 * g.xyz[:, 0]) + 800.0 * np.cos(9.0 * g.xyz[:, 1])
            _, relief = EN.landform(g, z, np.ones(g.n, np.int8), np.ones(g.n, np.int8), np.full(g.n, np.inf),
                                    np.full(g.n, 800.0), np.zeros(g.n, bool), relief_km=700.0)
            med.append(np.median(relief))
        self.assertAlmostEqual(med[0] / med[1], 1.0, delta=0.3)


class WetlandRuleTest(unittest.TestCase):
    def test_wetlands_need_real_rivers_on_real_flats(self):
        from types import SimpleNamespace
        from mapgen import graph as G
        g = small_grid_env()
        n = g.n
        z = np.where(g.lat > -30, 100.0 + 0.0001 * g.lat, -3000.0)          # a very flat humid plain
        common = dict(lit=np.full(n, EN.LI_GRANITE, np.int8), t_mean=np.full(n, 15.0), t_min=np.full(n, 5.0),
                      p_ann=np.full(n, 1200.0), dist_ocean=np.full(n, 2000.0), river=np.zeros(n, bool),
                      strahler=np.zeros(n, np.int8), lake=np.zeros(n, bool), salt_flat=np.zeros(n, bool))
        def wet(q):
            gr = EN.ground(g, z, common["lit"], common["t_mean"], common["t_min"], common["p_ann"], common["dist_ocean"],
                           common["river"], common["strahler"], np.full(n, q), common["lake"], common["salt_flat"],
                           EN.DEFAULTS, g.lat <= -30)
            return np.mean(gr[g.lat > 0] == EN.GR_WETLAND)
        self.assertLess(wet(3.0), 0.05)          # small streams don't make a wetland
        self.assertGreater(wet(8.0), 0.9)        # a big river across a flat does


def small_grid_env():
    from tests.helpers import small_grid
    return small_grid(3)


class DepositsTest(unittest.TestCase):
    def test_shares_rules_and_main(self):
        from mapgen import minerals as MN
        from mapgen import environment as EN
        from mapgen.crust import CRATON, POST_OROGEN
        g = make_ctx(4).grid
        n = g.n
        land = np.abs(g.lat) < 60
        mountain = land & (g.lon > 0) & (g.lon < 60)
        f = {"land": land, "z": np.where(mountain, 3500.0, 200.0), "lit": np.where(mountain, EN.LI_METAMORPHIC, EN.LI_SANDSTONE).astype(np.int8),
             "age": np.where(mountain, POST_OROGEN, CRATON), "lf": np.where(mountain, EN.LF_MOUNTAINS, EN.LF_PLAIN),
             "relief": np.where(mountain, 2500.0, 50.0), "t_mean": np.full(n, 15.0), "p_ann": np.full(n, 900.0),
             "d_coll": np.full(n, np.inf), "salt": np.zeros(n, bool), "endo": np.zeros(n, bool),
             "dist_ocean": np.full(n, 500.0), "coal": land & ~mountain}
        bits, main = MN.place(g, f, 7)
        self.assertEqual(len(MN.DEPOSITS), 32)
        self.assertFalse(bits[~land].any())
        for name in ("ironstone", "oil_gas", "iron_high", "tungsten"):
            has = (bits & np.uint32(MN.BIT[name])) > 0
            share = dict((d[0], d[2]) for d in MN.DEPOSITS)[name]
            self.assertLessEqual(has.sum(), round(share * land.sum()) + 1, name)
            self.assertTrue(has.any(), name)
        for name in MN.DEEP[:2]:                          # the good stuff: in the mountains only here
            has = (bits & np.uint32(MN.BIT[name])) > 0
            self.assertTrue(np.all(mountain[has]), name)
        some = bits > 0
        self.assertTrue(np.all(main[some] > 0) and np.all(main[~some] == 0))
        k = np.flatnonzero(some)[0]
        self.assertTrue(bits[k] & np.uint32(1 << (int(main[k]) - 1)))

    def test_deposits_are_sprinkled_not_one_hotspot(self):
        from mapgen import minerals as MN
        g = make_ctx(4).grid
        land = np.abs(g.lat) < 60
        hot = land & (g.lat > 0) & (g.lat < 30) & (g.lon > 0) & (g.lon < 40)      # a world-class district
        score = np.where(hot, 3.0, np.where(land, 1.0, 0.0))
        prov = MN.provinces(g, 7)
        sel = MN._rank_select(score, 0.02, land, prov)
        self.assertEqual(sel.sum(), round(0.02 * land.sum()))
        self.assertFalse(sel[~land].any())
        self.assertGreater((sel & hot).sum(), 0.3 * sel.sum())                 # the district keeps the most
        self.assertGreater((sel & ~hot).sum(), 0.3 * sel.sum())                # the rest is sprinkled
        self.assertGreater(len(np.unique(prov[sel & ~hot])), 0.3 * len(np.unique(prov[land & ~hot])))
        self.assertFalse(MN._rank_select(np.where(hot, 1.0, 0.0), 0.02, land, prov)[~hot].any())   # rules still rule
        poor = MN._rank_select(np.where(hot, 10.0, np.where(land, 1.0, 0.0)), 0.02, land, prov)
        self.assertFalse(poor[~hot].any())                                     # far poorer ground: no local mines