import unittest

import numpy as np

from tests.helpers import fields, line_world, make_history, make_race
from worldhistory.events import apply_event, cull, era_switch, seed
from worldhistory.regions import Regions
from worldhistory.state import new_state


class EventTest(unittest.TestCase):
    def test_seed_totals_and_fringe(self):
        w = line_world(200)
        r = make_race(seed={"clusters": 20, "heads": 100})
        q = np.linspace(0, 1, 200)[None]
        st = new_state(1, 200, np.zeros((1, 6)), seed=0)
        placed = seed(st, w, [r], q, np.ones((1, 200), bool), 0.25, np.random.default_rng(0))
        self.assertAlmostEqual(float(st.P.sum()), 2000.0)
        cells = placed["a"]
        fringe_cells = [c for c in cells if 0.05 < q[0, c] < 0.3]
        self.assertGreaterEqual(len(fringe_cells), 5)                     # 25 % of 20 on the fringe
        self.assertEqual(float(st.P[0, 0]), 0.0)                          # q = 0: never

    def test_seed_realm_land_skips_sea(self):
        ocean = np.arange(200) >= 100                                     # right half is sea, and the best q
        w = line_world(200, ocean=ocean)
        r = make_race(seed={"clusters": 20, "heads": 100, "realm": "land"}, habitat={"realm": "both"})
        q = np.linspace(0.1, 1, 200)[None]
        st = new_state(1, 200, np.zeros((1, 6)), seed=0)
        placed = seed(st, w, [r], q, np.ones((1, 200), bool), 0.25, np.random.default_rng(0))
        self.assertTrue(placed["a"])
        self.assertTrue(all(c < 100 for c in placed["a"]))

    def test_cull_exact_and_by_race(self):
        rs = [make_race("a"), make_race("b")]
        st = new_state(2, 3, np.zeros((2, 6)), seed=0)
        st.P[:] = 100.0
        cull(st, rs, np.array([True, True, False]), 0.6, {"b": 0.8})
        np.testing.assert_allclose(st.P, [[40, 40, 100], [20, 20, 100]])

    def test_cull_noise(self):
        w = line_world(500)
        st = new_state(1, 500, np.zeros((1, 6)), seed=0)
        st.P[:] = 100.0
        cull(st, [make_race()], np.ones(500, bool), 0.3, {}, noise=w.noise(1), amp=0.15)
        self.assertAlmostEqual(float(st.P.mean()), 70.0, delta=2.0)
        self.assertGreater(float(st.P.std()), 5.0)

    def test_era_switch_drowns_land_races(self):
        w = line_world(3)
        w.eras["late"] = fields(3, ocean=[False, True, True])
        land, sea = make_race("l"), make_race("s", habitat={"realm": "sea", "terms": [{"p": "sea", "w": 1}]})
        st = new_state(2, 3, np.zeros((2, 6)), seed=0)
        st.P[0] = 10.0
        st.P[1, 2] = 5.0
        era_switch(st, w, [land, sea], "late")
        np.testing.assert_array_equal(st.P[0], [10, 0, 0])
        np.testing.assert_array_equal(st.P[1], [0, 0, 5])
        self.assertEqual(w.era, "late")

    def test_targeted_seed_places_one_race_in_region(self):
        w = line_world(200)                                               # lon 0 .. ~178°, 0.9° apart
        h = make_history(regions={"west": {"kind": "box", "lat": [-1, 1], "lon": [-1, 45]}},
                         event=[{"year": 0, "kind": "seed", "race": "b", "region": "west", "clusters": 3,
                                 "heads": 50}])
        rs = [make_race("a", seed={"clusters": 5, "heads": 100}), make_race("b", seed={"clusters": 5, "heads": 100})]
        q = np.ones((2, 200)) * 0.5
        st = new_state(2, 200, np.zeros((2, 6)), seed=0)
        log = apply_event(st, w, rs, h.events[0], Regions(w, h.regions), q, np.ones((2, 200), bool), h, 0)
        self.assertEqual(float(st.P[0].sum()), 0.0)                       # race a: untouched
        self.assertAlmostEqual(float(st.P[1].sum()), 150.0)               # 3 clusters × 50
        west = np.flatnonzero(w.lon <= 45)
        self.assertTrue(set(np.flatnonzero(st.P[1])) <= set(west))
        self.assertEqual(len(log["placed"]["b"]), 3)

    def test_apply_event_log(self):
        w = line_world(2)
        h = make_history(regions={"west": {"kind": "box", "lat": [-1, 1], "lon": [-1, 0.5]}},
                         event=[{"year": 20, "kind": "die_off", "region": "west", "share": 0.9}])
        st = new_state(1, 2, np.zeros((1, 6)), seed=0)
        st.P[:] = 100.0
        log = apply_event(st, w, [make_race()], h.events[0], Regions(w, h.regions), None, None, h, 0)
        np.testing.assert_allclose(st.P[0], [10, 100])
        self.assertEqual(log["before"]["a"], 200.0)
        self.assertEqual(log["after"]["a"], 110.0)