import unittest

import numpy as np

from tests.helpers import line_world, make_history, make_race
from worldhistory.habitat import environment, suitability
from worldhistory.migration import drift_and_bud, long_jumps
from worldhistory.predicates import evaluate
from worldhistory.state import new_state

MP = make_history().migration


def east(w, speed=1.0):
    """Uniform current along the line (the line runs along the equator, eastward = +lon)."""
    lon = np.radians(w.lon)
    return speed * np.stack([-np.sin(lon), np.cos(lon), np.zeros(w.n)], 1)


def sea_world(n, spacing_km=100.0, current=True, **over):
    w = line_world(n, spacing_km=spacing_km, ocean=True, elevation_m=-3000.0, **over)
    if current:
        w.base["current"] = east(w)
    return w


def swimmer(bias):
    return make_race(demography={"founder": 2.0, "allee": 1.0, "jump_path": "sea", "current_bias": bias,
                                 "mobility": 0.2},
                     habitat={"realm": "sea", "terms": [{"p": "sea", "w": 1.0}]})


class SailCostTest(unittest.TestCase):
    def test_downstream_cheaper_than_still_cheaper_than_upstream(self):
        w = sea_world(3)
        c = w.sail_cost(5.0)
        k_east = list(w.nb[1]).index(2)
        k_west = list(w.nb[1]).index(0)
        edge_h = 100.0 * 1000.0 / 3600.0                                # 100 km in hours per (m/s)
        self.assertAlmostEqual(c[1, k_east], edge_h / 6.0, delta=0.01)  # 5 + 1 m/s
        self.assertAlmostEqual(c[1, k_west], edge_h / 4.0, delta=0.01)  # 5 − 1 m/s
        still = sea_world(3, current=False).sail_cost(5.0)
        self.assertAlmostEqual(still[1, k_east], edge_h / 5.0, delta=0.01)
        self.assertTrue(np.isinf(c[0, list(w.nb[0]).index(-1)]))        # no neighbour: no edge

    def test_land_edges_are_impassable(self):
        w = line_world(3, ocean=[True, True, False])
        c = w.sail_cost(5.0)
        self.assertTrue(np.isinf(c[1, list(w.nb[1]).index(2)]))


class FieldsTest(unittest.TestCase):
    def test_missing_current_is_still_water(self):
        w = sea_world(3, current=False)
        self.assertEqual(w.current.shape, (3, 3))
        self.assertTrue(np.all(w.current == 0))

    def test_productivity_predicate(self):
        w = sea_world(3, productivity=np.array([0.0, 0.5, 1.0]))
        np.testing.assert_allclose(evaluate("productivity", w), [0.0, 0.5, 1.0])
        np.testing.assert_allclose(evaluate("productivity", sea_world(3)), [0.0, 0.0, 0.0])   # absent → 0


class JumpBiasTest(unittest.TestCase):
    def landings(self, bias, current=True):
        w = sea_world(400, spacing_km=10.0, current=current)
        r = swimmer(bias)
        env, suit = environment(w), suitability(w, r)[None]
        east_n = west_n = 0
        for seed in range(150):
            st = new_state(1, w.n, np.zeros((1, 6)), seed=0)
            st.P[0, 200] = 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.0)
            hit = np.flatnonzero(st.P[0] > 0)
            east_n += int(np.sum(hit > 200))
            west_n += int(np.sum(hit < 200))
        return east_n, west_n

    def test_downstream_jumps_dominate_with_bias(self):
        e, w = self.landings(3.0)
        self.assertGreater(e, 3 * max(w, 1))

    def test_no_bias_is_symmetric(self):
        e, w = self.landings(0.0)
        self.assertLess(abs(e - w), 0.35 * (e + w) + 5)

    def test_no_current_field_ignores_bias(self):
        e, w = self.landings(3.0, current=False)
        self.assertLess(abs(e - w), 0.35 * (e + w) + 5)


class DriftBiasTest(unittest.TestCase):
    def drift(self, bias):
        w = sea_world(3)
        r = swimmer(bias)
        st = new_state(1, 3, np.zeros((1, 6)), seed=0)
        st.P[0] = [100.0, 1000.0, 100.0]                          # neighbours settled: drift, not founding
        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)
        return st.P[0]

    def test_drift_follows_the_current(self):
        P = self.drift(2.0)
        self.assertGreater(P[2] - 100.0, 3.0 * max(P[0] - 100.0, 1e-9))   # the downstream neighbour gains most

    def test_no_bias_drifts_evenly(self):
        P = self.drift(0.0)
        self.assertAlmostEqual(P[0], P[2], delta=1e-6 * P.sum())