import unittest

import numpy as np

from tests.helpers import line_world, make_history, make_race
from worldhistory.migration import drift_and_bud, long_jumps, travel_cost
from worldhistory.config import CONDITIONS
from worldhistory.state import new_state

D_DEPTH = CONDITIONS.index("depth")

MP = make_history().migration


def setup(n, pop, race, K=1e6):
    st = new_state(1, n, np.zeros((1, 6)), seed=0)
    st.P[0, 0] = pop
    return st, np.full((1, n), K), np.ones((1, n), bool)


class MigrationTest(unittest.TestCase):
    def test_travel_cost(self):
        w = line_world(4, ocean=[False, False, False, True], landform=[1, 3, 1, 0], river=[True, False, False, False],
                       holdridge=[20, 20, 23, 0])
        r = make_race(travel={"mountain": 2.0, "desert": 1.0})
        np.testing.assert_allclose(travel_cost(w, r, np.zeros(4)), [-1, 2, 0, 0])   # cell 2: desert +1, coast -1
        np.testing.assert_allclose(travel_cost(w, r, np.full(4, 0.5)), [-1.5, 1.5, -0.5, -0.5])

    def test_budding_needs_founder_group(self):
        w = line_world(3)
        r = make_race(demography={"founder": 50.0, "mobility": 0.1})
        st, K, OK = setup(3, 60.0, r, K=100.0)          # crowded, but below 2 founders
        drift_and_bud(st, w, [r], K, OK, np.zeros((1, 3)), np.zeros((1, 3)), np.ones((1, 3)),
                      np.random.default_rng(0), MP)
        self.assertEqual(float(st.P[0, 1]), 0.0)
        st, K, OK = setup(3, 120.0, r, K=100.0)
        drift_and_bud(st, w, [r], K, OK, np.zeros((1, 3)), np.zeros((1, 3)), np.ones((1, 3)),
                      np.random.default_rng(0), MP)
        self.assertIn(float(st.P[0, 1]), (0.0, 20.0))     # a founder group (capped at P/6 = 20) or nothing
        self.assertAlmostEqual(float(st.P[0].sum()), 120.0)

    def test_river_colonised_before_mountain(self):
        r = make_race(demography={"founder": 5.0, "mobility": 0.1}, travel={"mountain": 2.0})
        hits = {0: 0, 2: 0}
        for seed in range(60):
            st = new_state(1, 3, np.zeros((1, 6)), seed=0)
            st.P[0, 1] = 100.0
            w2 = line_world(3, landform=[1, 1, 3], river=[True, True, False])   # cell 0 river plain, cell 2 mountain
            tc = travel_cost(w2, r, np.zeros(3))[None]
            drift_and_bud(st, w2, [r], np.full((1, 3), 100.0), np.ones((1, 3), bool), tc, np.zeros((1, 3)),
                          np.ones((1, 3)), np.random.default_rng(seed), MP)
            hits[0] += st.P[0, 0] > 0
            hits[2] += st.P[0, 2] > 0
        self.assertGreater(hits[0], hits[2])

    def test_drift_toward_headroom(self):
        w = line_world(3)
        r = make_race(demography={"mobility": 0.2})
        st = new_state(1, 3, np.zeros((1, 6)), seed=0)
        st.P[0] = [100.0, 100.0, 100.0]
        K = np.array([[1000.0, 1000.0, 5000.0]])
        drift_and_bud(st, w, [r], K, np.ones((1, 3), bool), np.zeros((1, 3)), np.zeros((1, 3)),
                      np.ones((1, 3)), np.random.default_rng(0), MP)
        self.assertGreater(float(st.P[0, 2]), float(st.P[0, 0]))  # cell 1 sends more toward the roomier side
        self.assertAlmostEqual(float(st.P[0].sum()), 300.0)

    def test_never_into_uninhabitable(self):
        w = line_world(3)
        r = make_race(demography={"mobility": 0.2, "founder": 1.0})
        st = new_state(1, 3, np.zeros((1, 6)), seed=0)
        st.P[0] = [0.0, 100.0, 100.0]
        drift_and_bud(st, w, [r], np.full((1, 3), 100.0), np.array([[False, True, True]]), np.zeros((1, 3)),
                      np.zeros((1, 3)), np.ones((1, 3)), np.random.default_rng(0), MP)
        self.assertEqual(float(st.P[0, 0]), 0.0)
        self.assertAlmostEqual(float(st.P[0].sum()), 200.0)

    def test_drift_keeps_viable_source_viable(self):
        # uncrowded viable group (22, Allee 20) beside a small one: drift may not take it below the Allee threshold
        w = line_world(2)
        r = make_race(demography={"mobility": 0.5, "allee": 20.0, "founder": 20.0})
        st = new_state(1, 2, np.zeros((1, 6)), seed=0)
        st.P[0] = [22.0, 5.0]
        drift_and_bud(st, w, [r], np.full((1, 2), 1e5), np.ones((1, 2), bool), np.zeros((1, 2)),
                      np.zeros((1, 2)), np.ones((1, 2)), np.random.default_rng(0), MP)
        self.assertGreaterEqual(float(st.P[0, 0]), 20.0)
        self.assertAlmostEqual(float(st.P[0].sum()), 27.0)

    def test_small_group_joins_viable_neighbour(self):
        # below the Allee threshold, a group gathers into its viable neighbour rather than dwindling alone
        w = line_world(3)
        r = make_race(demography={"mobility": 0.1, "allee": 20.0, "founder": 20.0})
        st = new_state(1, 3, np.zeros((1, 6)), seed=0)
        st.P[0] = [8.0, 30.0, 0.0]
        drift_and_bud(st, w, [r], np.full((1, 3), 1e5), np.ones((1, 3), bool), np.zeros((1, 3)),
                      np.zeros((1, 3)), np.ones((1, 3)), np.random.default_rng(0), MP)
        np.testing.assert_allclose(st.P[0], [0.0, 38.0, 0.0])

    def test_push_flees_into_empty_land(self):
        # people driven out (push) settle an empty livable neighbour even when nobody is crowded
        w = line_world(2)
        r = make_race(demography={"mobility": 0.0, "allee": 20.0, "founder": 20.0})
        st = new_state(1, 2, np.zeros((1, 6)), seed=0)
        st.P[0] = [120.0, 0.0]
        st.push[0, 0] = 30.0
        drift_and_bud(st, w, [r], np.full((1, 2), 1e5), np.ones((1, 2), bool), np.zeros((1, 2)),
                      np.zeros((1, 2)), np.ones((1, 2)), np.random.default_rng(0), MP)
        np.testing.assert_allclose(st.P[0], [100.0, 20.0])                  # capped at P/6 per neighbour

    def test_roaming_groups_travel_together(self):
        # roam = chance per step that a whole group (a caravan) moves on, together, to one neighbouring cell
        w = line_world(3)
        r = make_race(demography={"mobility": 0.0, "allee": 2.0, "founder": 2.0, "roam": 1.0})
        ends = set()
        for seed in range(20):
            st = new_state(1, 3, np.zeros((1, 6)), seed=0)
            st.P[0] = [0.0, 120.0, 0.0]
            drift_and_bud(st, w, [r], np.full((1, 3), 1e5), np.ones((1, 3), bool), np.zeros((1, 3)),
                          np.zeros((1, 3)), np.ones((1, 3)), np.random.default_rng(seed), MP)
            self.assertEqual(float(st.P[0, 1]), 0.0)
            self.assertIn(float(st.P[0].max()), (120.0,))
            ends.add(int(st.P[0].argmax()))
        self.assertEqual(ends, {0, 2})                                     # either way, never split

    def test_long_jump_crosses_chasm_on_migrants_adaptation(self):
        # deep-adapted swimmers skip an abyss they cannot live in and land on ground that suits *them*, even though
        # the landing cells' stored optimum (native, 0 m) would call them unfit
        from worldhistory.habitat import environment, suitability
        depth = np.r_[np.full(5, 2000.0), np.full(35, 6000.0), np.full(40, 2000.0)]
        w = line_world(80, spacing_km=10.0, ocean=True, elevation_m=-depth)
        tol = {"depth": {"optimum": 0, "width": 60, "lo": 0, "hi": 7000, "rate": 2.0}}
        r = make_race(demography={"founder": 2.0, "allee": 1.0, "jump": 1.0}, tolerance=tol,
                      habitat={"realm": "sea", "terms": [{"p": "sea", "w": 1.0}]})
        env, suit = environment(w), suitability(w, r)[None]
        landed = set()
        for seed in range(200):
            st = new_state(1, w.n, np.zeros((1, 6)), seed=0)
            st.P[0, 0] = 1e6
            st.O[0, D_DEPTH, :5] = 2000.0
            long_jumps(st, w, [r], np.full((1, w.n), 1e6), np.ones((1, w.n), bool), np.random.default_rng(seed),
                       {**MP, "jump_rate": 1.0}, np.ones((1, w.n)), env=env, suit=suit, f_min=0.1)
            landed |= set(np.flatnonzero(st.P[0, 1:] > 0) + 1)
        self.assertTrue(any(c >= 40 for c in landed))                    # across the chasm
        self.assertFalse(any(5 <= c < 40 for c in landed))               # never into the abyss

    def test_sea_path_jumps_never_cross_land(self):
        # swimmers (jump_path = sea) cannot hop over a strip of land; without it they can
        ocean = np.r_[np.ones(5, bool), np.zeros(5, bool), np.ones(70, bool)]
        w = line_world(80, spacing_km=10.0, ocean=ocean, elevation_m=np.where(ocean, -100.0, 10.0))
        from worldhistory.habitat import environment, suitability
        for path, beyond in (("any", True), ("sea", False)):
            r = make_race(demography={"founder": 2.0, "allee": 1.0, "jump_path": path},
                          habitat={"realm": "sea", "terms": [{"p": "sea", "w": 1.0}]})
            env, suit = environment(w), suitability(w, r)[None]
            landed = set()
            for seed in range(200):
                st = new_state(1, w.n, np.zeros((1, 6)), seed=0)
                st.P[0, 0] = 1e6
                long_jumps(st, w, [r], np.full((1, w.n), 1e6), np.ones((1, w.n), bool), np.random.default_rng(seed),
                           {**MP, "jump_rate": 1.0}, np.ones((1, w.n)), env=env, suit=suit, f_min=0.1)
                landed |= set(np.flatnonzero(st.P[0, 1:] > 0) + 1)
            self.assertEqual(any(c >= 10 for c in landed), beyond, path)

    def test_large_group_splits(self):
        # clans above group_max break up: half leaves as a new group into empty land
        w = line_world(3)
        r = make_race(demography={"mobility": 0.0, "allee": 20.0, "founder": 20.0, "group_max": 500.0})
        st = new_state(1, 3, np.zeros((1, 6)), seed=0)
        st.P[0] = [0.0, 1200.0, 0.0]
        drift_and_bud(st, w, [r], np.full((1, 3), 1e5), np.ones((1, 3), bool), np.zeros((1, 3)),
                      np.zeros((1, 3)), np.ones((1, 3)), np.random.default_rng(0), MP)
        np.testing.assert_allclose(st.P[0], [200.0, 800.0, 200.0])          # 600 leave, capped P/6 per side

    def test_push_leaves(self):
        w = line_world(2)
        r = make_race(demography={"mobility": 0.0})
        st = new_state(1, 2, np.zeros((1, 6)), seed=0)
        st.P[0] = [100.0, 100.0]
        st.push[0, 0] = 10.0
        drift_and_bud(st, w, [r], np.full((1, 2), 1000.0), np.ones((1, 2), bool), np.zeros((1, 2)),
                      np.zeros((1, 2)), np.ones((1, 2)), np.random.default_rng(0), MP)
        np.testing.assert_allclose(st.P[0], [90.0, 110.0])
        self.assertEqual(float(st.push.sum()), 0.0)

    def test_long_jump_distances(self):
        w = line_world(4000, spacing_km=10.0)
        r = make_race(demography={"founder": 2.0, "allee": 1.0})
        st = new_state(1, w.n, np.zeros((1, 6)), seed=0)
        st.P[0, 0] = 1e6
        d = []
        for seed in range(300):
            st.P[0] = 0
            st.P[0, 0] = 1e6
            long_jumps(st, w, [r], np.full((1, w.n), 1e6), np.ones((1, w.n), bool), np.random.default_rng(seed),
                       {**MP, "jump_rate": 1.0}, np.ones((1, w.n)))
            hit = np.flatnonzero(st.P[0, 1:]) + 1
            d += list(np.minimum(hit, w.n - hit) * 10.0)          # the line circles the globe: wrap
        d = np.array(d)
        self.assertGreater(len(d), 100)          # a line world: only jumps along the line land near it
        self.assertTrue((d <= MP["jump_cap_km"] + 10).all())
        self.assertGreater(float(np.median(d)), MP["jump_xm_km"] * 0.8)