import unittest

import numpy as np

from tests.helpers import line_world, make_race
from worldhistory.config import CONDITIONS, link
from worldhistory.habitat import native_optima
from worldhistory.lineage import birth_step, chance, init_rules, native, progress, return_step
from worldhistory.regions import Regions
from worldhistory.state import new_state
from tests.helpers import make_history

D, G = CONDITIONS.index("depth"), CONDITIONS.index("gravity")
PARENT_TOL = {"depth": {"optimum": 0, "width_lo": 60, "width_hi": 100, "lo": 0, "hi": 6000, "rate": 2.0},
              "gravity": {"optimum": 1.0, "width": 0.5, "lo": 0.3, "hi": 1.3, "rate": 0.01}}
SUB_TOL = {"depth": {"optimum": 2000, "width_lo": 400, "width_hi": 1500, "lo": 500, "hi": 6000, "rate": 2.0},
           "gravity": {"optimum": 0.5, "width": 0.5, "lo": 0.3, "hi": 1.3, "rate": 0.01}}


def pair(n=6, emerge=None, **fields):
    w = line_world(n, **fields)
    races = link(make_history(), [make_race("p", tolerance=PARENT_TOL),
                                  make_race("s", tolerance=SUB_TOL, emerge={"parent": "p", **(emerge or {})})])
    st = new_state(2, n, native_optima(races), seed=0)
    return w, races, st


class ProgressTest(unittest.TestCase):
    def test_weighted_mean_over_distinct_axes(self):
        w, (p, s), st = pair(1)
        O = native(p)[:, None].copy()
        O[D, 0] = 1000.0                      # halfway on depth (weight 2000/100 = 20)
        a = progress(O, p, s)                 # gravity untouched (weight 0.5/0.5 = 1)
        self.assertAlmostEqual(float(a[0]), 20 * 0.5 / 21, places=6)

    def test_overshoot_clips_and_equal_axes_ignored(self):
        w, (p, s), st = pair(1)
        O = native(p)[:, None].copy()
        O[D, 0], O[G, 0] = 3000.0, 0.2       # both past the sub-species' optimum
        self.assertAlmostEqual(float(progress(O, p, s)[0]), 1.0)
        self.assertEqual(float(progress(O, p, p)[0]), 0.0)                # no distinct axis → 0

    def test_back_toward_parent(self):
        w, (p, s), st = pair(1)
        O = native(s)[:, None].copy()
        O[D, 0] = 1000.0                      # Havs risen halfway; toward the parent the width is s's lower side (400)
        a = progress(O, s, p)                 # weights: depth 2000/400 = 5, gravity 1
        self.assertAlmostEqual(float(a[0]), 5 * 0.5 / 6, places=6)


class ChanceTest(unittest.TestCase):
    def test_floor_ceiling_and_monotone(self):
        a = np.array([0.0, 0.39, 0.4, 0.5, 0.6, 0.7, 0.8, 1.0])
        c = chance(a, 0.4, 0.8, 0.2)
        np.testing.assert_allclose(c[[0, 1, 2]], 0.0)
        np.testing.assert_allclose(c[[6, 7]], 0.2)
        self.assertTrue(np.all(np.diff(c) >= 0))
        self.assertAlmostEqual(float(c[4]), 0.1)                         # smoothstep(0.5) = 0.5

    def test_equal_bounds_is_a_step(self):
        np.testing.assert_allclose(chance(np.array([0.0, 0.5]), 0.5, 0.5, 1.0), [0.0, 1.0])


class BirthTest(unittest.TestCase):
    def deep(self, st, cells, depth=2000.0):
        st.O[0, D, cells] = depth

    def test_no_birth_below_min(self):
        w, races, st = pair()
        st.P[0] = 100
        self.deep(st, slice(None), 500.0)                                # a = 20·0.25/21 < 0.4
        rules = init_rules(races)
        ok = np.ones((2, 6), bool)
        for s in range(200):
            self.assertEqual(birth_step(st, w, races, rules, Regions(w, {}), np.random.default_rng(s), ok), [])

    def test_birth_resets_to_native_and_is_once(self):
        w, races, st = pair(emerge={"birth_rate": 1.0})
        st.P[0] = [100, 100, 0, 0, 100, 100]
        self.deep(st, slice(None))
        st.T[0, 0, :] = 0.5
        rules = init_rules(races)
        ok = np.ones((2, 6), bool)
        new = birth_step(st, w, races, rules, Regions(w, {}), np.random.default_rng(0), ok)
        self.assertEqual(len(new), 1)
        c = new[0]["cell"]
        self.assertEqual(float(st.P[0, c]), 0.0)                         # converted wholly
        self.assertAlmostEqual(float(st.P[1, c]), 100.0)
        np.testing.assert_allclose(st.O[1, :, c], native(races[1]))     # baseline shift = 1
        self.assertAlmostEqual(float(st.T[1, 0, c]), 0.5)                # tech kept
        self.assertGreater(new[0]["a"], 0.8)
        for s in range(20):
            self.assertEqual(birth_step(st, w, races, rules, Regions(w, {}), np.random.default_rng(s + 1), ok), [])

    def test_birth_needs_native_fit(self):
        # Review Focus 1: where the sub-species' own curves do not fit, no birth
        w, races, st = pair(emerge={"birth_rate": 1.0})
        st.P[0] = 100
        self.deep(st, slice(None))
        ok = np.ones((2, 6), bool)
        ok[1] = False
        rules = init_rules(races)
        for s in range(50):
            self.assertEqual(birth_step(st, w, races, rules, Regions(w, {}), np.random.default_rng(s), ok), [])

    def test_origin_rules_still_apply(self):
        w, races, st = pair(emerge={"birth_rate": 1.0, "settlement_density": 0.001})
        st.P[0] = [10, 20, 90, 30, 10, 10]
        self.deep(st, slice(None))
        rules = init_rules(races)
        birth_step(st, w, races, rules, Regions(w, {}), np.random.default_rng(0), np.ones((2, 6), bool))
        self.assertEqual(rules[0]["origin"], 2)                          # the densest settlement

    def test_small_groups_rarely_found(self):
        counts = []
        for P in (1.0, 1000.0):
            n = 0
            for s in range(300):
                w = line_world(1)
                races = link(make_history(), [make_race("p", tolerance=PARENT_TOL, demography={"founder": 50.0}),
                                              make_race("s", tolerance=SUB_TOL, emerge={"parent": "p"})])
                st = new_state(2, 1, native_optima(races), seed=0)
                st.P[0] = P
                st.O[0, D] = 2000.0
                if birth_step(st, w, races, init_rules(races), Regions(w, {}), np.random.default_rng(s),
                              np.ones((2, 1), bool)):
                    n += 1
            counts.append(n)
        self.assertLess(counts[0], counts[1] / 5)                        # 1/51 vs 0.95 of 0.2 per step


class ReturnTest(unittest.TestCase):
    def test_return_births_parents_many_times_capped(self):
        # Review Focus 3: returns may fire in several cells, never convert more than the group, never log emergence
        w, races, st = pair(emerge={"birth_rate": 1.0})
        rules = init_rules(races)
        rules[0]["origin"], rules[0]["year"] = 0, 0                       # already born
        st.P[1] = [50, 50, 50, 0, 0, 0]
        st.O[1, :, :] = native(races[1])[:, None]
        st.O[1, D, :3] = 0.0                                             # risen all the way back
        ok = np.ones((2, 6), bool)
        rets = return_step(st, w, races, rules, np.random.default_rng(0), ok)
        self.assertEqual(sorted(r["cell"] for r in rets), [0, 1, 2])
        np.testing.assert_allclose(st.P[1, :3], 0.0)
        np.testing.assert_allclose(st.P[0, :3], 50.0)
        np.testing.assert_allclose(st.O[0, :, 0], native(races[0]))
        self.assertTrue((st.P >= 0).all())

    def test_no_return_before_birth_or_for_ritual(self):
        w, races, st = pair(emerge={"birth_rate": 1.0})
        rules = init_rules(races)
        st.P[1] = 50
        st.O[1, D] = 0.0
        self.assertEqual(return_step(st, w, races, rules, np.random.default_rng(0), np.ones((2, 6), bool)), [])


class RitualTest(unittest.TestCase):
    def test_ritual_birth_ignores_progress(self):
        w, races, st = pair(emerge={"mode": "ritual", "trigger": 1.0})
        st.P[0] = 100                                                    # optima untouched: a = 0
        rules = init_rules(races)
        new = birth_step(st, w, races, rules, Regions(w, {}), np.random.default_rng(0), np.ones((2, 6), bool))
        self.assertEqual(len(new), 1)


from worldhistory.lineage import displacement, inherit, rituals


class DisplacementTest(unittest.TestCase):
    def test_less_fit_race_squeezed(self):
        w, races, st = pair(3)
        rules = init_rules(races)
        rules[0]["origin"] = 0
        K = np.full((2, 3), 100.0)
        P = np.array([[30.0, 50.0, 10.0], [10.0, 0.0, 30.0]])
        fit = np.array([[0.2, 0.9, 0.9], [0.9, 0.1, 0.1]])
        displacement(K, P, fit, rules, races)
        np.testing.assert_allclose(K[0], [100 * (1 - 0.9 * 10 / 40), 100.0, 100.0])      # parent less fit in cell 0
        np.testing.assert_allclose(K[1], [100.0, 100.0, 100 * (1 - 0.9 * 10 / 40)])      # sub less fit in cell 2
        self.assertTrue((K >= 0).all())

    def test_no_effect_before_birth_or_on_strangers(self):
        # Review Focus 5
        w, races, st = pair(3)
        rules = init_rules(races)
        K = np.full((2, 3), 100.0)
        displacement(K, np.ones((2, 3)), np.array([[0.1] * 3, [0.9] * 3]), rules, races)
        np.testing.assert_allclose(K, 100.0)


class InheritTest(unittest.TestCase):
    def test_adapted_neighbours_bear_the_dominant_race(self):
        w, races, st = pair(4)
        rules = init_rules(races)
        rules[0]["origin"] = 0
        st.P[1] = [100, 0, 0, 0]
        st.P[0] = [0, 100, 100, 100]
        st.O[0, D, 1] = 2000.0                    # adapted (a = 20/21 ≥ sure)
        st.O[0, D, 2] = 2000.0                    # adapted but not in contact
        births = np.zeros((2, 4))
        births[0] = 10.0
        inherit(st, w, races, rules, births, Regions(w, {}), np.ones((2, w.n), bool))
        share = 100 / 300                         # around cell 1 (cells 0..2): 100 sub-species, 200 parents
        self.assertAlmostEqual(float(st.P[1, 1]), 10 * share, places=6)
        np.testing.assert_allclose(st.O[1, :, 1], native(races[1]))
        self.assertEqual(float(st.P[1, 3]), 0.0)

    def test_no_dominant_births_where_native_curves_do_not_fit(self):
        w, races, st = pair(2)
        rules = init_rules(races)
        rules[0]["origin"] = 0
        st.P[1] = [100, 0]
        st.P[0] = [0, 100]
        st.O[0, D, 1] = 2000.0
        births = np.zeros((2, 2))
        births[0] = 10.0
        ok = np.ones((2, 2), bool)
        ok[1, 1] = False
        inherit(st, w, races, rules, births, Regions(w, {}), ok)
        self.assertEqual(float(st.P[1, 1]), 0.0)

    def test_holdouts_refuse_to_mix(self):
        w, races, st = pair(2)
        rules = init_rules(races)
        rules[0]["origin"] = 0
        st.P[1] = [100, 0]
        st.P[0] = [0, 100]                        # optima native: a = 0 < mix_min
        births = np.zeros((2, 2))
        births[0] = 10.0
        inherit(st, w, races, rules, births, Regions(w, {}), np.ones((2, w.n), bool))
        self.assertEqual(float(st.P[1, 1]), 0.0)

    def test_ritual_race_is_plainly_dominant(self):
        w, races, st = pair(2, emerge={"mode": "ritual"})
        rules = init_rules(races)
        rules[0]["origin"] = 0
        st.P[1] = [100, 0]
        st.P[0] = [0, 100]                        # a = 0 but no gate in ritual mode
        births = np.zeros((2, 2))
        births[0] = 10.0
        inherit(st, w, races, rules, births, Regions(w, {}), np.ones((2, w.n), bool))
        self.assertAlmostEqual(float(st.P[1, 1]), 10 * 100 / 200, places=6)


class RitualConvertTest(unittest.TestCase):
    def test_rituals_cost_hundreds_and_scale_sublinearly(self):
        w, races, st = pair(2, emerge={"mode": "ritual", "ritual_rate": 0.5, "ritual_cost": 300,
                                       "ritual_converts": 30})
        rules = init_rules(races)
        rules[0]["origin"] = 0
        st.P[1] = [400, 0]
        st.P[0] = [0, 100000]
        out = rituals(st, w, races, rules, np.random.default_rng(0))
        n = sum(r["converts"] for r in out) / 30
        self.assertGreater(n, 0)
        self.assertAlmostEqual(float(st.P[0, 1]), 100000 - 330 * n)
        self.assertAlmostEqual(float(st.P[1].sum()), 400 + 30 * n)

    def test_rituals_never_overdraw(self):
        # Review Focus 4
        w, races, st = pair(2, emerge={"mode": "ritual", "ritual_rate": 50.0})
        rules = init_rules(races)
        rules[0]["origin"] = 0
        st.P[1] = [10000, 0]
        st.P[0] = [0, 500]
        rituals(st, w, races, rules, np.random.default_rng(0))
        self.assertTrue((st.P >= 0).all())


class SummaryTest(unittest.TestCase):
    def test_events_aggregate_per_step(self):
        from worldhistory.lineage import summarise
        ev = [{"race": "p", "from": "s", "cell": 1, "year": 10, "people": 5.0},
              {"race": "p", "from": "s", "cell": 2, "year": 10, "people": 7.0}]
        self.assertEqual(summarise(ev, ("people",)), [{"race": "p", "from": "s", "year": 10, "groups": 2, "people": 12.0}])
        self.assertEqual(summarise([], ("people",)), [])

    def test_split_by_source_and_dicts_summed(self):
        from worldhistory.lineage import summarise
        ev = [{"race": "p", "from": "s", "cell": 1, "year": 10, "people": 5.0},
              {"race": "p", "from": "t", "cell": 2, "year": 10, "people": 7.0},
              {"race": "p", "from": "s", "cell": 3, "year": 10, "people": 1.0}]
        self.assertEqual(summarise(ev, ("people",)),
                         [{"race": "p", "from": "s", "year": 10, "groups": 2, "people": 6.0},
                          {"race": "p", "from": "t", "year": 10, "groups": 1, "people": 7.0}])
        ev = [{"race": "b", "cell": 1, "year": 10, "victims": 4.0, "dead": {"p": 1.0, "o": 3.0}},
              {"race": "b", "cell": 2, "year": 10, "victims": 2.0, "dead": {"p": 2.0}}]
        self.assertEqual(summarise(ev, ("victims", "dead")),
                         [{"race": "b", "year": 10, "groups": 2, "victims": 6.0, "dead": {"p": 3.0, "o": 3.0}}])
        ev.append({"race": "b", "cell": -1, "year": 10, "victims": 0.0, "backlash": 9.0})    # no "dead" field
        self.assertEqual(summarise(ev, ("victims", "dead", "backlash"))[0]["dead"], {"p": 3.0, "o": 3.0})

    def test_backlash_record_is_summed_not_counted(self):
        from worldhistory.lineage import summarise
        ev = [{"race": "b", "cell": 3, "year": 10, "victims": 300.0, "converts": 30.0},
              {"race": "b", "cell": -1, "year": 10, "victims": 0.0, "converts": 0.0, "backlash": 70.0}]
        self.assertEqual(summarise(ev, ("victims", "converts", "backlash")),
                         [{"race": "b", "year": 10, "groups": 1, "victims": 300.0, "converts": 30.0, "backlash": 70.0}])

def trio(n=3, emerge=None):
    """Parent p, sub-species s, and an unrelated race o."""
    w = line_world(n)
    races = link(make_history(), [make_race("p", tolerance=PARENT_TOL), make_race("o"),
                                  make_race("s", tolerance=SUB_TOL, emerge={"parent": "p", **(emerge or {})})])
    st = new_state(3, n, native_optima(races), seed=0)
    return w, races, st


class AfterTest(unittest.TestCase):
    def test_no_birth_before_the_earliest_year(self):
        for mode in ("ritual", "adapt"):
            em = {"mode": mode, "trigger": 1.0, "after": 100} if mode == "ritual" else \
                 {"after": 100, "birth_min": 0.0, "birth_sure": 0.0, "birth_rate": 1.0}
            w, races, st = pair(emerge=em)
            st.P[0] = 1e6
            rules = init_rules(races)
            ok = np.ones((2, 6), bool)
            st.t = 90
            self.assertEqual(birth_step(st, w, races, rules, Regions(w, {}), np.random.default_rng(0), ok), [], mode)
            st.t = 100
            self.assertEqual(len(birth_step(st, w, races, rules, Regions(w, {}), np.random.default_rng(0), ok)), 1,
                             mode)


class CrisisTest(unittest.TestCase):
    def test_ritual_birth_only_where_parents_were_almost_eradicated(self):
        # cell 1 falls from 10,000 to 500 (−95 %), cell 3 from 10,000 to 5,000 (−50 %), cell 4 tiny all along
        w, races, st = pair(emerge={"mode": "ritual", "trigger": 1.0, "crisis_drop": 0.9, "crisis_min": 1000,
                                    "crisis_years": 200})
        rules = init_rules(races)
        ok = np.ones((2, 6), bool)
        rng = np.random.default_rng(0)
        st.P[0] = [0, 10000, 0, 10000, 50, 0]
        st.t = 0
        self.assertEqual(birth_step(st, w, races, rules, Regions(w, {}), rng, ok), [])   # peaks recorded, no crisis
        st.P[0] = [0, 500, 0, 5000, 5, 0]
        st.t = 10
        new = birth_step(st, w, races, rules, Regions(w, {}), rng, ok)
        self.assertEqual([e["cell"] for e in new], [1])

    def _off_habitat(self, km):
        # crisis in cell 1 (10,000 → 500), where s cannot live; s can live only in cell 3, 200 km away
        w, races, st = pair(emerge={"mode": "ritual", "trigger": 1.0, "crisis_drop": 0.9, "crisis_min": 1000,
                                    "convert": 0.5, **({"crisis_km": km} if km is not None else {})})
        rules = init_rules(races)
        ok = np.ones((2, 6), bool)
        ok[1] = [False, False, False, True, False, False]
        rng = np.random.default_rng(0)
        st.P[0] = [0, 10000, 0, 0, 0, 0]
        st.t = 0
        birth_step(st, w, races, rules, Regions(w, {}), rng, ok)
        st.P[0] = [0, 500, 0, 0, 0, 0]
        st.t = 10
        return st, birth_step(st, w, races, rules, Regions(w, {}), rng, ok)

    def test_crisis_off_habitat_births_in_nearest_habitat_within_range(self):
        st, new = self._off_habitat(250)
        self.assertEqual([e["cell"] for e in new], [3])
        self.assertEqual(new[0]["crisis"], 1)
        np.testing.assert_allclose(st.P[1], [0, 0, 0, 250, 0, 0])      # half the 500 survivors, moved to cell 3
        np.testing.assert_allclose(st.P[0], [0, 250, 0, 0, 0, 0])

    def test_crisis_off_habitat_no_birth_out_of_range_or_by_default(self):
        for km in (150, None):
            st, new = self._off_habitat(km)
            self.assertEqual(new, [], km)
            np.testing.assert_allclose(st.P[1], 0)

    def test_no_crisis_no_birth(self):
        w, races, st = pair(emerge={"mode": "ritual", "trigger": 1.0, "crisis_drop": 0.9})
        st.P[0] = 10000
        rules = init_rules(races)
        self.assertEqual(birth_step(st, w, races, rules, Regions(w, {}), np.random.default_rng(0),
                                    np.ones((2, 6), bool)), [])


class VictimsTest(unittest.TestCase):
    def test_victims_from_every_nearby_race_converts_from_parents(self):
        # Blod (s) in cell 0; parents 30,000 in cell 1; other race 90,000 in cell 1 → victims split 1:3
        w, races, st = trio(emerge={"mode": "ritual", "ritual_rate": 0.5, "ritual_cost": 300,
                                    "ritual_converts": 30, "victims": "all"})
        rules = init_rules(races)
        rules[0]["origin"] = 0
        p, o, s = 0, 1, 2
        st.P[s] = [400, 0, 0]
        st.P[p] = [0, 30000, 0]
        st.P[o] = [0, 90000, 0]
        out = rituals(st, w, races, rules, np.random.default_rng(0))
        k = sum(r["converts"] for r in out) / 30
        self.assertGreater(k, 0)
        self.assertAlmostEqual(float(st.P[p, 1]), 30000 - 30 * k - 300 * k * 0.25)
        self.assertAlmostEqual(float(st.P[o, 1]), 90000 - 300 * k * 0.75)
        self.assertAlmostEqual(float(st.P[s].sum()), 400 + 30 * k)
        self.assertTrue((st.P >= 0).all())
        dead = {}
        for r in out:
            for key, x in r["dead"].items():
                dead[key] = dead.get(key, 0.0) + x
        self.assertAlmostEqual(dead["p"], 300 * k * 0.25)
        self.assertAlmostEqual(dead["o"], 300 * k * 0.75)


class BacklashTest(unittest.TestCase):
    def test_smackdown_once_then_calmer(self):
        w, races, st = pair(2, emerge={"mode": "ritual", "ritual_rate": 0.5, "ritual_cost": 300,
                                       "ritual_converts": 30, "backlash_victims": 600, "backlash_loss": 0.7,
                                       "backlash_calm": 0.1})
        rules = init_rules(races)
        rules[0]["origin"] = 0
        st.P[1] = [400, 0]
        st.P[0] = [0, 100000]
        rng = np.random.default_rng(0)
        rituals(st, w, races, rules, rng)
        v = rules[0]["victims"]
        self.assertGreaterEqual(v, 600)                                  # ≥ 2 rituals: the smackdown fires
        self.assertTrue(rules[0]["backlash"])
        blod = 400 + v / 300 * 30
        self.assertAlmostEqual(float(st.P[1].sum()), blod * 0.3)
        self.assertAlmostEqual(rules[0]["rate"], 0.05)
        before = st.P[1].sum()
        rules[0]["victims"] += 1e6
        rituals(st, w, races, rules, rng)
        self.assertGreaterEqual(float(st.P[1].sum()), before)            # no second smackdown


class ReturnRateTest(unittest.TestCase):
    def test_return_rate_sets_the_chance(self):
        w, races, st = pair(emerge={"birth_rate": 1.0, "return_rate": 0.0})
        rules = init_rules(races)
        rules[0]["origin"] = 0
        st.P[1] = 1e6
        st.O[1, :, :] = native(races[1])[:, None]
        st.O[1, D, :] = 0.0
        self.assertEqual(return_step(st, w, races, rules, np.random.default_rng(0), np.ones((2, 6), bool)), [])


from worldhistory.lineage import raids

RAID = {"mode": "ritual", "raid_rate": 5.0, "raid_size": 100, "raid_min": 1000, "raid_km": 250}


class RaidTest(unittest.TestCase):
    def setUp(self):
        self.w, self.races, self.st = pair(6, emerge=RAID)
        self.rules = init_rules(self.races)
        self.rules[0]["origin"] = 0
        self.st.P[1] = [5000, 0, 0, 0, 0, 0]
        self.st.P[0] = [0, 0, 800, 0, 900, 0]         # cell 1 empty, cell 2 in reach (200 km), cell 4 out (400 km)
        self.ok = np.ones((2, 6), bool)

    def test_cells_raid_occupied_non_blod_land_in_reach(self):
        out = raids(self.st, self.w, self.races, self.rules, np.random.default_rng(0), self.ok)
        n = sum(r["raids"] for r in out)
        self.assertGreater(n, 0)
        np.testing.assert_allclose(self.st.P[1], [5000 - 100 * n, 0, 100 * n, 0, 0, 0])
        np.testing.assert_allclose(self.st.P[0], [0, 0, 800, 0, 900, 0])     # the killing is the rituals' job

    def test_no_raids_after_the_smackdown(self):
        self.rules[0]["backlash"] = True
        self.assertEqual(raids(self.st, self.w, self.races, self.rules, np.random.default_rng(0), self.ok), [])
        self.assertEqual(float(self.st.P[1, 0]), 5000.0)

    def test_small_groups_do_not_raid(self):
        self.st.P[1, 0] = 900
        self.assertEqual(raids(self.st, self.w, self.races, self.rules, np.random.default_rng(0), self.ok), [])

    def test_no_raids_where_blod_cannot_live(self):
        self.ok[1, 2] = False
        self.assertEqual(raids(self.st, self.w, self.races, self.rules, np.random.default_rng(0), self.ok), [])


class FrenzyTest(unittest.TestCase):
    def test_linear_rituals_scale_with_cultists(self):
        # ritual_power 1: rituals ≈ rate·P per step; 10,000 Blod at 1/15 → ≈ 667 rituals, 200k dead
        w, races, st = trio(emerge={"mode": "ritual", "ritual_rate": 1 / 15, "ritual_power": 1.0,
                                    "ritual_cost": 300, "ritual_converts": 30, "victims": "all"})
        rules = init_rules(races)
        rules[0]["origin"] = 0
        st.P[2] = [10000, 0, 0]
        st.P[0] = [0, 1e7, 0]
        st.P[1] = [0, 1e7, 0]
        out = rituals(st, w, races, rules, np.random.default_rng(0))
        dead = sum(r["victims"] for r in out)
        self.assertAlmostEqual(dead / 200000, 1.0, delta=0.08)                 # Poisson(667): sd ≈ 4 %


class SmackdownTest(unittest.TestCase):
    def rig(self, **em):
        w, races, st = trio(5, emerge={"mode": "ritual", "ritual_rate": 0.0, "backlash_loss": 0.7, **em})
        rules = init_rules(races)
        rules[0]["origin"], rules[0]["year"] = 0, 100
        return w, races, st, rules

    def test_time_limit_fires_without_the_body_count(self):
        w, races, st, rules = self.rig(backlash_victims=3e6, backlash_years=50)
        st.P[2] = [1000, 0, 0, 0, 0]
        st.t = 140
        rituals(st, w, races, rules, np.random.default_rng(0))
        self.assertFalse(rules[0]["backlash"])
        st.t = 150
        rituals(st, w, races, rules, np.random.default_rng(0))
        self.assertTrue(rules[0]["backlash"])
        self.assertAlmostEqual(float(st.P[2].sum()), 300.0)

    def test_small_remote_cells_survive_big_exposed_ones_pay(self):
        # cell 0: 10,000 Blod beside 50,000 enemies; cell 4: 100 Blod alone → survivors 30 % of 10,100 = 3,030
        w, races, st, rules = self.rig(backlash_years=1)
        st.P[2] = [10000, 0, 0, 0, 100]
        st.P[0] = [0, 50000, 0, 0, 0]
        st.t = 101
        rituals(st, w, races, rules, np.random.default_rng(0))
        self.assertAlmostEqual(float(st.P[2].sum()), 3030.0, places=6)
        # kill share ∝ exposure (own + nearby enemies): 60,000 vs 100 → cell 4 loses 100·100·κ ≈ 0.12
        kappa = 7070 / (10000 * 60000 + 100 * 100)
        self.assertAlmostEqual(float(st.P[2, 4]), 100 - 100 * 100 * kappa, places=6)
        self.assertGreater(float(st.P[2, 4]), 99.8)