import unittest

import numpy as np

from tests.helpers import line_world, make_race
from worldhistory.demography import (allee, capacity, crowding, grow, harvest_shocks, overlap_matrix,
                                     stochastic_round, veins)


class DemographyTest(unittest.TestCase):
    def test_capacity_formula(self):
        w = line_world(2)
        r = make_race(demography={"density": 2.0})
        K = capacity(w, r, np.array([1.0, 0.5]), np.ones(2), np.ones(2), np.array([11.0, 11.0]), np.ones(2),
                     np.ones(2))
        # density * area * q * (1 + gain q²): 2*10000*1*12, 2*10000*0.5*(1+11*0.25)
        np.testing.assert_allclose(K, [240000, 37500])

    def test_logistic_reaches_K(self):
        P = np.array([[10.0]])
        K = np.array([[1000.0]])
        r = [make_race(demography={"growth": 0.5})]
        for _ in range(200):
            grow(P, K, crowding(P, K, np.zeros((1, 1)), r, 2.0), np.ones((1, 1)), r, np.ones((1, 1)))
        self.assertAlmostEqual(float(P[0, 0]), 1000.0, places=3)

    def test_growth_rate_low_gear(self):
        P = np.array([[10.0, 10.0]])
        K = np.array([[1e9, 1e9]])
        r = [make_race(demography={"growth": 0.1})]
        grow(P, K, crowding(P, K, np.zeros((1, 1)), r, 2.0), np.array([[1.0, 0.0]]), r, np.ones((1, 2)))
        np.testing.assert_allclose(P[0], [11.0, 10.3], rtol=1e-6)      # r*(0.3+0.7q)

    def test_grow_returns_births(self):
        r = [make_race()]
        P = np.array([[100.0, 100.0]])
        b = grow(P, np.array([[1e6, 50.0]]), np.array([[0.0, 2.0]]), np.ones((1, 2)), r, np.ones((1, 2)))
        # gross births (generational turnover), not net growth: a crowded, shrinking group still bears children
        np.testing.assert_allclose(b[0], [0.1 * 100, 0.1 * 100])

    def test_allee(self):
        self.assertLess(allee(np.array(10.0), 75.0), 0)
        self.assertGreater(allee(np.array(100.0), 75.0), 0)
        self.assertAlmostEqual(float(allee(np.array(1.0), 1000.0)), -0.3)
        P = np.array([[10.0]])
        r = [make_race(demography={"growth": 0.1, "allee": 75.0})]
        grow(P, np.array([[1e6]]), np.zeros((1, 1)), np.ones((1, 1)), r, np.ones((1, 1)))
        self.assertLess(float(P[0, 0]), 10.0)

    def test_no_capacity_halves(self):
        P = np.array([[10.0]])
        grow(P, np.array([[0.0]]), np.zeros((1, 1)), np.ones((1, 1)), [make_race()], np.ones((1, 1)))
        self.assertAlmostEqual(float(P[0, 0]), 5.0)

    def test_competition(self):
        P = np.array([[100.0], [100.0]])
        K = np.array([[1000.0], [1000.0]])
        rs = [make_race("a"), make_race("b")]
        c0 = crowding(P, K, np.zeros((2, 2)), rs, 1.0)
        c1 = crowding(P, K, np.array([[0, 1.0], [1.0, 0]]), rs, 1.0)
        np.testing.assert_allclose(c0[:, 0], [0.1, 0.1])
        np.testing.assert_allclose(c1[:, 0], [0.2, 0.2])

    def test_tolerated_share_sharpens(self):
        P = np.array([[100.0], [100.0]])
        K = np.array([[1000.0], [1000.0]])
        rs = [make_race("a", temperament={"tolerated_share": 0.1}), make_race("b", temperament={"tolerated_share": 0.9})]
        c = crowding(P, K, np.array([[0, 1.0], [1.0, 0]]), rs, 2.0)
        np.testing.assert_allclose(c[:, 0], [0.3, 0.2])               # a: others' share 0.5 > 0.1 → ×2

    def test_overlap_matrix(self):
        S = np.array([[1.0, 0, 0], [1.0, 0, 0], [0, 0, 1.0]])
        calm = {"xenophobia": 0.0, "pace": 0.5, "structure": 0.5, "warlike": 0.0}
        rs = [make_race("a", temperament=calm), make_race("b", temperament=calm), make_race("c", temperament=calm)]
        a = overlap_matrix(S, rs)
        self.assertAlmostEqual(a[0, 1], 0.5)          # sim 1 × friction 0.5
        self.assertAlmostEqual(a[0, 2], 0.0)          # disjoint habitats
        rs[0].temperament = {**calm, "xenophobia": 1.0}
        a = overlap_matrix(S, rs)
        self.assertAlmostEqual(a[0, 1], 1.0)          # a minds b more than b minds a
        self.assertAlmostEqual(a[1, 0], 0.5)
        rs[1].overlap = {"a": 0.05}
        self.assertAlmostEqual(overlap_matrix(S, rs)[1, 0], 0.05)

    def test_shocks_mean_one_and_scale(self):
        w = line_world(20000)
        e = np.random.default_rng(1).normal(size=w.n)
        m = harvest_shocks(w, np.ones(w.n), np.ones(w.n), e)
        self.assertAlmostEqual(float(m.mean()), 1.0, places=2)
        np.testing.assert_array_equal(harvest_shocks(w, np.zeros(w.n), np.ones(w.n), e), 1.0)

    def test_veins_need_viable_community(self):
        np.testing.assert_allclose(veins(np.array([50.0, 100.0]), np.array([1000.0, 1000.0]), 0.01, 75.0),
                                   [50.0, 109.0])

    def test_stochastic_round_expectation(self):
        P = np.full((1, 100000), 2.3)
        stochastic_round(P, np.random.default_rng(0))
        self.assertTrue(set(np.unique(P)) <= {2.0, 3.0})
        self.assertAlmostEqual(float(P.mean()), 2.3, places=2)
        Q = np.array([[7.5, 0.0]])
        stochastic_round(Q, np.random.default_rng(0))
        np.testing.assert_array_equal(Q, [[7.5, 0.0]])


class RangeTest(unittest.TestCase):
    def test_wide_range_crowding(self):
        # a clan drawing on a wide range: crowding = people over capacity summed across the neighbourhood
        from tests.helpers import line_world
        w = line_world(3)
        P, K = np.array([[300.0, 0.0, 0.0]]), np.array([[100.0, 100.0, 100.0]])
        narrow = crowding(P, K, np.zeros((1, 1)), [make_race()], 2.0)
        wide = crowding(P, K, np.zeros((1, 1)), [make_race(demography={"range": 1})], 2.0, world=w)
        self.assertAlmostEqual(float(narrow[0, 0]), 3.0)
        self.assertAlmostEqual(float(wide[0, 0]), 1.5)                 # 300 over cells 0+1 (200)

    def test_world_range_crowding(self):
        # range -1: crowding is the whole people's numbers against its whole capacity (few, widely scattered clans)
        from tests.helpers import line_world
        w = line_world(3)
        P, K = np.array([[300.0, 0.0, 0.0]]), np.array([[100.0, 100.0, 0.0]])
        c = crowding(P, K, np.zeros((1, 1)), [make_race(demography={"range": -1})], 2.0, world=w)
        np.testing.assert_allclose(c[0], [1.5, 1.5, 0.0])                # 300 / 200 where it can live