import unittest

import numpy as np

from tests.helpers import line_world, make_history, make_race
from worldhistory.state import new_state
from worldhistory.tech import DI, effective, reach, step_land, step_tech

TP = make_history().tech


def state(n, pops):
    st = new_state(1, n, np.zeros((1, 6)), seed=0)
    st.P[0] = pops
    return st


class TechTest(unittest.TestCase):
    def test_progress_grows_with_population(self):
        w = line_world(7)
        st = state(7, [0, 0, 10000, 0, 0, 20, 0])
        step_tech(st, w, [make_race()], {**TP, "diffuse": 0.0, "loss_below": 0.0})
        big, small = st.T[0, DI["farming"], 2], st.T[0, DI["farming"], 5]
        # s = N/(N+n_half): 10000/15000 vs 20/5020; gain = rate*s
        self.assertAlmostEqual(float(big), TP["rate"] * 10000 / 15000, places=6)
        self.assertAlmostEqual(float(small), TP["rate"] * 20 / 5020, places=6)
        self.assertEqual(float(st.T[0, 0, 0]), 0.0)                    # empty cells do not learn

    def test_priority_scales(self):
        w = line_world(1)
        st = state(1, [10000])
        step_tech(st, w, [make_race(tech={"travel": 0.0, "farming": 2.0})], {**TP, "diffuse": 0.0})
        self.assertEqual(float(st.T[0, DI["travel"], 0]), 0.0)
        self.assertAlmostEqual(float(st.T[0, DI["farming"], 0]), 2 * TP["rate"] * 10000 / 15000, places=6)

    def test_small_isolated_groups_lose_tech(self):
        w = line_world(1)
        st = state(1, [50])
        st.T[0, :, 0] = 0.8
        step_tech(st, w, [make_race()], {**TP, "diffuse": 0.0, "rate": 0.0})
        self.assertAlmostEqual(float(st.T[0, 0, 0]), 0.8 * (1 - TP["loss_rate"]), places=6)

    def test_diffusion_only_upward(self):
        w = line_world(2)
        st = state(2, [1000, 1000])
        st.T[0, 0] = [1.0, 0.0]
        step_tech(st, w, [make_race()], {**TP, "rate": 0.0, "loss_below": 0.0, "diffuse": 0.5})
        self.assertAlmostEqual(float(st.T[0, 0, 0]), 1.0)
        self.assertAlmostEqual(float(st.T[0, 0, 1]), 0.25)            # 0.5 * (mean 0.5 - 0)

    def test_magic_folds_in_and_reach_capped(self):
        T = np.zeros((6, 1))
        T[DI["magic"]] = 1.0
        eff = effective(T, TP)
        self.assertAlmostEqual(float(eff["farming"][0]), TP["magic_share"])
        self.assertAlmostEqual(float(reach(eff, make_race(tech={"reach_cap": 0.1}), TP)[0]), 0.1)

    def test_land_improvement_up_and_decay(self):
        st = state(2, [100, 0])
        st.T[0, DI["land"]] = 1.0
        st.improve[0] = [0.0, 0.4]
        step_land(st, [make_race()], TP)
        self.assertAlmostEqual(float(st.improve[0, 0]), TP["land_rate"] * 1.0 * TP["land_cap"], places=6)
        self.assertAlmostEqual(float(st.improve[0, 1]), 0.4 * (1 - TP["land_decay"]), places=6)