import unittest

import numpy as np

from tests.helpers import line_world, make_race
from worldhistory.adaptation import adapt
from worldhistory.config import CONDITIONS
from worldhistory.habitat import environment, fitness, native_optima
from worldhistory.state import new_state

D = CONDITIONS.index("depth")
TOL = {"depth": {"optimum": 0, "width": 50, "lo": 0, "hi": 600, "rate": 2.0}}


def world_and_state(race, pops):
    w = line_world(4, ocean=[False, True, True, True], elevation_m=[10, -200, -600, -1000])
    st = new_state(1, 4, native_optima([race]), seed=0)
    st.P[0] = pops
    return w, st


class AdaptationTest(unittest.TestCase):
    def test_moves_at_most_rate_times_years(self):
        r = make_race(tolerance=TOL)
        w, st = world_and_state(r, [0, 10, 0, 0])
        adapt(st, environment(w), [r], np.ones((1, 4)), 10)
        self.assertAlmostEqual(float(st.O[0, D, 1]), 20.0)         # 2 m/yr × 10 yr toward the deep end (400)
        self.assertEqual(float(st.O[0, D, 0]), 0.0)                # empty cells unchanged

    def test_multiplier_speeds_adaptation(self):
        r = make_race(tolerance=TOL)
        w, st = world_and_state(r, [0, 10, 0, 0])
        adapt(st, environment(w), [r], np.ones((1, 4)), 10, mult=np.full((1, 4), 2.0))
        self.assertAlmostEqual(float(st.O[0, D, 1]), 40.0)         # pressed: 2 × (2 m/yr × 10 yr)

    def test_turnover_slows_poor_habitat(self):
        r = make_race(tolerance=TOL)
        w, st = world_and_state(r, [0, 10, 0, 0])
        adapt(st, environment(w), [r], np.zeros((1, 4)), 10)
        self.assertAlmostEqual(float(st.O[0, D, 1]), 6.0)          # × 0.3

    def test_soft_limit_brakes_but_does_not_stop(self):
        r = make_race(tolerance=TOL)                               # hi 600, width 50
        w, st = world_and_state(r, [0, 0, 0, 10])
        st.O[0, D, 3] = 650.0                                      # 1 width past hi
        adapt(st, environment(w), [r], np.ones((1, 4)), 10)
        self.assertAlmostEqual(float(st.O[0, D, 3]), 650.0 + 20.0 * np.exp(-0.5), places=4)
        for _ in range(200):
            adapt(st, environment(w), [r], np.ones((1, 4)), 10)
        self.assertGreater(float(st.O[0, D, 3]), 650.0)
        self.assertLess(float(st.O[0, D, 3]), 800.0)               # far slower than the free rate (2 m/yr)

    def test_inward_moves_are_not_braked(self):
        r = make_race(tolerance=TOL)
        w = line_world(3)                                          # all land: target depth 0
        st = new_state(1, 3, native_optima([r]), seed=0)
        st.P[0] = 10
        st.O[0, D] = 700.0                                         # past hi, moving back in
        adapt(st, environment(w), [r], np.ones((1, 3)), 10)
        np.testing.assert_allclose(st.O[0, D], 680.0)

    def test_rain_does_not_drift_at_sea(self):
        r = make_race(habitat={"realm": "both", "terms": [{"p": "sea", "w": 1}]},
                      tolerance={"rain": {"optimum": 3.0, "width": 0.3, "lo": 1.5, "hi": 3.5, "rate": 0.01}})
        w = line_world(3, ocean=True, elevation_m=-100)
        st = new_state(1, 3, native_optima([r]), seed=0)
        st.P[0] = 10
        adapt(st, environment(w), [r], np.ones((1, 3)), 10)
        np.testing.assert_allclose(st.O[0, CONDITIONS.index("rain")], 3.0)

    def test_inland_returns_to_native(self):
        r = make_race(tolerance=TOL)
        w = line_world(3)                                          # all land: depth span [0, 0]
        st = new_state(1, 3, native_optima([r]), seed=0)
        st.P[0] = 10
        st.O[0, D] = 100.0
        adapt(st, environment(w), [r], np.ones((1, 3)), 10)
        np.testing.assert_allclose(st.O[0, D], 80.0)

    def test_adapting_costs_the_old(self):
        r = make_race(tolerance=TOL)
        w, st = world_and_state(r, [0, 0, 0, 10])
        env = environment(w)
        O_before = st.O[0].copy()
        for _ in range(60):
            adapt(st, env, [r], np.ones((1, 4)), 10)
        at_surface_before = fitness(env, O_before, r)[0]
        at_surface_after = fitness(env, np.repeat(st.O[0][:, 3:4], 4, 1), r)[0]
        self.assertAlmostEqual(float(at_surface_before), 1.0)
        self.assertLess(float(at_surface_after), 0.1)              # deep-adapted people cannot live at the surface


class ProspectiveTest(unittest.TestCase):
    def test_empty_cells_take_neighbour_optimum(self):
        from worldhistory.adaptation import prospective
        r = make_race(tolerance=TOL)
        w, st = world_and_state(r, [0, 30, 0, 10])
        st.O[0, D] = [0.0, 100.0, 0.0, 400.0]
        prospective(st, w)
        # cell 0: only neighbour 1 occupied -> 100; cell 2: (30*100 + 10*400) / 40 = 175; occupied cells unchanged
        np.testing.assert_allclose(st.O[0, D], [100.0, 100.0, 175.0, 400.0])

    def test_isolated_empty_cells_keep_their_value(self):
        from worldhistory.adaptation import prospective
        r = make_race(tolerance=TOL)
        w, st = world_and_state(r, [0, 0, 0, 0])
        st.O[0, D] = [5.0, 6.0, 7.0, 8.0]
        prospective(st, w)
        np.testing.assert_allclose(st.O[0, D], [5.0, 6.0, 7.0, 8.0])

    def test_fractional_groups_keep_their_optimum(self):
        # all groups count (spec §1.3): a thin frontier (P < 1) is not "empty"
        from worldhistory.adaptation import prospective
        r = make_race(tolerance=TOL)
        w, st = world_and_state(r, [100, 0.5, 0, 0])
        st.O[0, D] = [0.0, 900.0, 0.0, 0.0]
        prospective(st, w)
        self.assertEqual(float(st.O[0, D, 1]), 900.0)
        self.assertEqual(float(st.O[0, D, 2]), 900.0)          # the empty cell beyond takes the frontier's optimum