import unittest

import numpy as np

from tests.helpers import line_world, make_race
from worldhistory.config import CONDITIONS
from worldhistory.habitat import comfort, environment, fitness, intervals, native_optima, quality, suitability

DEPTH = CONDITIONS.index("depth")
GRAV = CONDITIONS.index("gravity")


class HabitatTest(unittest.TestCase):
    def test_intervals_span_to_neighbour_midpoints(self):
        w = line_world(4)
        lo, hi = intervals(w, np.array([0.0, 100.0, 300.0, 300.0]))
        np.testing.assert_array_equal(lo, [0, 50, 200, 300])
        np.testing.assert_array_equal(hi, [50, 200, 300, 300])

    def test_environment_depth_and_temperature(self):
        w = line_world(3, ocean=[False, True, True], elevation_m=[10, -200, -1000], T_mean=[20, 15, 15],
                       bottom_temp_c=[0, 8, 3])
        env = environment(w)
        np.testing.assert_array_equal(env.x["depth"], [0, 200, 1000])
        np.testing.assert_array_equal(env.x["temperature"], [20, 8, 3])
        np.testing.assert_array_equal(env.hi["depth"], [100, 600, 1000])

    def test_suitability_terms_gates_realm(self):
        w = line_world(4, ocean=[False, False, False, True], holdridge=[20, 23, 20, 0], T_mean=[20, 20, 0, 20])
        r = make_race(habitat={"realm": "land",
                               "terms": [{"p": "land", "w": 0.25}, {"p": "forest", "w": 0.5}],
                               "gates": [{"require": "warm:0:20", "realm": "land"}]})
        np.testing.assert_allclose(suitability(w, r), [0.75, 0.25, 0.0, 0.0])
        pen = make_race(habitat={"realm": "both", "terms": [{"p": "land", "w": 1}, {"p": "sea", "w": 1}],
                                 "gates": [{"penalty": "sea", "factor": 0.5}]})
        np.testing.assert_allclose(suitability(w, pen), [1, 1, 1, 0.5])

    def test_quality_normalised_to_p99(self):
        s = np.r_[np.linspace(0, 1, 101), 50.0]         # one outlier must not squash everyone else
        q = quality(s, make_race(habitat={"qmax": 0.7}), np.ones_like(s))
        self.assertAlmostEqual(float(q[100]), 0.7, places=6)
        self.assertLessEqual(float(q.max()), 0.7)

    def test_quality_all_zero(self):
        q = quality(np.zeros(5), make_race(), np.ones(5))
        np.testing.assert_array_equal(q, 0)

    def test_fitness_uses_distance_to_interval_and_reach(self):
        w = line_world(3, ocean=[False, True, True], elevation_m=[0, -100, -400])
        env = environment(w)
        r = make_race(tolerance={"depth": {"optimum": 0, "width": 50, "lo": 0, "hi": 5000}})
        O = native_optima([r])[0]
        O = np.repeat(O[:, None], 3, 1)
        f = fitness(env, O, r)
        self.assertAlmostEqual(float(f[0]), 1.0)                                  # interval [0, 50] holds 0
        self.assertAlmostEqual(float(f[1]), np.exp(-0.5), places=6)               # interval [50, 250]: d = 50 = 1 width
        self.assertAlmostEqual(float(fitness(env, O, r, reach=1.0)[1]), np.exp(-0.125), places=6)

    def test_unlisted_condition_has_no_effect(self):
        w = line_world(2, gravity_g=[1.0, 0.35])
        r = make_race()
        O = np.repeat(native_optima([r])[0][:, None], 2, 1)
        np.testing.assert_array_equal(fitness(environment(w), O, r), [1, 1])

    def test_comfort_piecewise(self):
        r = make_race(tolerance={"gravity": {"optimum": 1.0, "width": 0.2, "lo": 0.3, "hi": 1.3,
                                             "comfort": [[1.0, 1.0], [0.35, 1.2]]}})
        O = np.zeros((len(CONDITIONS), 3))
        O[GRAV] = [1.0, 0.675, 0.35]
        np.testing.assert_allclose(comfort(O, r), [1.0, 1.1, 1.2])

    def test_lopsided_widths(self):
        w = line_world(5, gravity_g=[1.05, 1.05, 0.75, 0.45, 0.45])
        r = make_race(tolerance={"gravity": {"optimum": 0.75, "width_lo": 0.3, "width_hi": 0.15, "lo": 0.3, "hi": 1.0}})
        O = np.repeat(native_optima([r])[0][:, None], 5, 1)
        f = fitness(environment(w), O, r)
        # cell 0 spans [1.05, 1.05]: 0.3 above = 2 upper widths; cell 4 spans [0.45, 0.45]: 0.3 below = 1 lower width
        self.assertAlmostEqual(float(f[0]), np.exp(-2.0), places=6)
        self.assertAlmostEqual(float(f[4]), np.exp(-0.5), places=6)

    def test_width_sets_both_sides(self):
        r = make_race(tolerance={"depth": {"optimum": 0, "width": 50, "lo": 0, "hi": 5000}})
        t = r.tolerance["depth"]
        self.assertEqual((t.width_lo, t.width_hi), (50, 50))

    def test_rain_is_log_and_ignored_at_sea(self):
        w = line_world(3, ocean=[False, False, True], P_ann=[100.0, 1000.0, 1000.0], elevation_m=[10, 10, -100])
        env = environment(w)
        np.testing.assert_allclose(env.x["rain"][:2], [2.0, 3.0])
        self.assertEqual(env.lo["rain"][2], -np.inf)
        self.assertEqual(env.hi["rain"][2], np.inf)
        r = make_race(habitat={"realm": "both", "terms": [{"p": "land", "w": 1}]},
                      tolerance={"rain": {"optimum": 2.0, "width": 0.3, "lo": 1.5, "hi": 3.5}})
        O = np.repeat(native_optima([r])[0][:, None], 3, 1)
        f = fitness(env, O, r)
        self.assertEqual(float(f[2]), 1.0)                     # no rain effect in the sea
        self.assertLess(float(f[1]), float(f[0]))

    def test_default_strain_off_native(self):
        # adapted away from the native optimum: livable, not good — the peak falls to 0.35 at the range edge
        r = make_race(tolerance={"gravity": {"optimum": 1.0, "width": 0.2, "lo": 0.3, "hi": 1.3}})
        O = np.zeros((len(CONDITIONS), 4))
        O[GRAV] = [1.0, 0.3, 0.65, 1.3]
        np.testing.assert_allclose(comfort(O, r), [1.0, 0.35, 0.35 ** 0.25, 0.35], rtol=1e-9)