import unittest

import numpy as np

from tests.helpers import globe_world, line_world, make_history, make_race
from worldhistory.hazards import make_hazard
from worldhistory.regions import Regions
from worldhistory.state import new_state


def hist(**hz):
    return make_history(regions={"north": {"kind": "box", "lat": [30, 90], "lon": [-180, 180]}}, hazard=[hz])


class HazardTest(unittest.TestCase):
    def setUp(self):
        self.w = globe_world(1)
        self.st = new_state(1, self.w.n, np.zeros((1, 6)), seed=0)
        self.st.P[:] = 100.0

    def test_static_in_window_and_region(self):
        h = hist(kind="static", region="north", mortality=0.1, years=[10, 30])
        hz = make_hazard(h.hazards[0])
        reg = Regions(self.w, h.regions)
        north = reg("north")
        hz.step(self.st, self.w, [make_race()], reg, 0, np.random.default_rng(0))
        self.assertTrue((self.st.P == 100).all())
        hz.step(self.st, self.w, [make_race()], reg, 10, np.random.default_rng(0))
        np.testing.assert_allclose(self.st.P[0, north], 90.0)
        np.testing.assert_allclose(self.st.P[0, ~north], 100.0)

    def test_static_push_drives_people_off_land(self):
        # megafauna: on land in the region, a share flees each step (not killed); drift is told to avoid those cells
        w = line_world(4, ocean=[False, False, True, True])
        h = make_history(regions={"coast": {"kind": "box", "lat": [-1, 1], "lon": [-1, 99]}},
                         hazard=[{"kind": "static", "name": "beasts", "region": "coast", "mortality": 0.0,
                                  "years": [0, 100], "push": 0.2}])
        st = new_state(1, 4, np.zeros((1, 6)), seed=0)
        st.P[:] = 100.0
        hz = make_hazard(h.hazards[0])
        hz.step(st, w, [make_race()], Regions(w, h.regions), 0, np.random.default_rng(0))
        np.testing.assert_allclose(st.P[0], 100.0)
        np.testing.assert_allclose(st.push[0], [20, 20, 0, 0])
        np.testing.assert_allclose(hz.danger, [1, 1, 0, 0])
        hz.step(st, w, [make_race()], Regions(w, h.regions), 200, np.random.default_rng(0))
        np.testing.assert_allclose(hz.danger, 0)                          # outside its years: no danger

    def test_roaming_local_stay_home_and_raid(self):
        h = hist(kind="roaming", start=20, home_region="north", home_near="a", local=10, wanderers=0,
                 radius_km=600, raid_chance=0.5, growth=0.0, infight=0.0, raid_death=0.0)
        hz = make_hazard(h.hazards[0])
        reg = Regions(self.w, h.regions)
        rng = np.random.default_rng(3)
        hz.step(self.st, self.w, [make_race()], reg, 10, rng)
        self.assertTrue((self.st.P == 100).all())                    # not started yet
        stats = [hz.step(self.st, self.w, [make_race()], reg, t, rng) for t in range(20, 200, 10)]
        self.assertTrue(reg("north")[hz.home].all())
        self.assertEqual(len(hz.home), 10)
        self.assertGreater(sum(s["raids"] for s in stats), 0)
        self.assertTrue((self.st.P[0, reg("north")] < 100).any())
        np.testing.assert_array_equal(self.st.P[0, self.w.lat < 20], 100.0)   # north (lat ≥ 30) + 600 km only

    def test_roaming_wanderer_leaves_home_region(self):
        h = hist(kind="roaming", start=0, home_region="north", local=0, wanderers=1, wander_stop=1.0,
                 wander_step_km=1500.0, growth=0.0, infight=0.0)
        hz = make_hazard(h.hazards[0])
        reg = Regions(self.w, h.regions)
        rng = np.random.default_rng(0)
        for t in range(0, 400, 10):
            hz.step(self.st, self.w, [make_race()], reg, t, rng)
        self.assertTrue((self.st.P[0, ~reg("north")] < 100).any())

    def test_roaming_growth_and_deaths(self):
        h = hist(kind="roaming", start=0, home_region="north", local=100, wanderers=0, growth=0.1,
                 clutch=[2, 4], infight=0.05, raid_chance=0.0)
        hz = make_hazard(h.hazards[0])
        reg = Regions(self.w, h.regions)
        rng = np.random.default_rng(0)
        births = deaths = 0
        for t in range(0, 500, 10):
            s = hz.step(self.st, self.w, [make_race()], reg, t, rng)
            births += s["births"]
            deaths += s["deaths"]
        self.assertGreater(births, 0)
        self.assertGreater(deaths, 0)
        self.assertEqual(len(hz.home), 100 + births - deaths)

    def test_roaming_growth_levels_off_at_cap(self):
        # growth 0.1/step uncapped would give 20·1.1^100 ≈ 275k units; logistic cap holds it near 60
        h = hist(kind="roaming", start=0, home_region="north", local=20, wanderers=0, growth=0.1,
                 clutch=[2, 4], infight=0.0, raid_chance=0.0, cap=60)
        hz = make_hazard(h.hazards[0])
        reg = Regions(self.w, h.regions)
        rng = np.random.default_rng(0)
        for t in range(0, 1000, 10):
            hz.step(self.st, self.w, [make_race()], reg, t, rng)
        self.assertGreater(len(hz.home), 40)
        self.assertLessEqual(len(hz.home), 60 + 4)

    def test_roaming_fallback_home(self):
        self.st.P[:] = 0.0
        h = hist(kind="roaming", start=0, home_region="north", home_near="a", local=5, wanderers=0)
        hz = make_hazard(h.hazards[0])
        reg = Regions(self.w, h.regions)
        hz.step(self.st, self.w, [make_race()], reg, 0, np.random.default_rng(0))
        self.assertTrue(reg("north")[hz.home].all())