import unittest

import numpy as np

from mapgen import climate as CL
from tests.helpers import small_grid


class InsolationTest(unittest.TestCase):
    def test_insolation_finite_at_poles(self):
        q = CL.insolation(np.array([90.0, -90.0, 89.999]), 20.0)
        self.assertTrue(np.all(np.isfinite(q)))
        self.assertAlmostEqual(q[1], 0.0, places=6)          # polar night
        self.assertGreater(q[0], 400)                          # polar day

    def test_symmetry(self):
        lat = np.linspace(-80, 80, 17)
        np.testing.assert_allclose(CL.insolation(lat, 20.0), CL.insolation(-lat, -20.0), rtol=1e-12)


class TemperatureTest(unittest.TestCase):
    def setUp(self):
        self.g = small_grid(3)
        self.P = dict(CL.DEFAULTS)

    def test_aquaplanet_zonal_means_fall_poleward(self):
        g = self.g
        z = np.full(g.n, -4000.0)
        T, _ = CL.temperatures(g, z, z > 0, self.P, 20.0)
        tm = (T["jun"] + T["dec"] + 2 * T["eq"]) / 4
        bands = [tm[(np.abs(g.lat) >= a) & (np.abs(g.lat) < a + 10)].mean() for a in range(0, 90, 10)]
        self.assertTrue(all(b1 > b2 for b1, b2 in zip(bands, bands[1:])), bands)
        self.assertTrue(20 < bands[0] < 32 and bands[-1] < -5, bands)

    def test_lapse_rate(self):
        g = self.g
        land = (np.abs(g.lat) < 30) & (np.abs(g.lon) < 60)
        flat = np.where(land, 10.0, -4000.0)
        high = np.where(land & (np.abs(g.lon) < 20), 3000.0, flat)
        T0, _ = CL.temperatures(g, flat, land, self.P, 20.0)
        T1, _ = CL.temperatures(g, high, land, self.P, 20.0)
        i = g.cell_index(0.0, 0.0)
        self.assertAlmostEqual(T0["eq"][i] - T1["eq"][i], 6.5 * 2.99, delta=0.5)

    def test_cold_west_coast_warm_east_coast(self):
        g = self.g
        land = (g.lat > 15) & (g.lat < 45) & (np.abs(g.lon) < 30)
        a = CL.current_anomaly(g, land, self.P)
        self.assertLess(a[g.cell_index(30.0, -33.0)], -0.5)
        self.assertGreater(a[g.cell_index(30.0, 33.0)], 0.5)

    def test_continental_interior_has_bigger_seasons(self):
        g = self.g
        land = (g.lat > 20) & (g.lat < 70) & (np.abs(g.lon) < 90)
        z = np.where(land, 200.0, -4000.0)
        T, _ = CL.temperatures(g, z, land, self.P, 20.0)
        rng = np.abs(T["jun"] - T["dec"])
        self.assertGreater(rng[g.cell_index(50.0, 0.0)], rng[g.cell_index(50.0, 150.0)] + 10)

    def test_biotemperature(self):
        bio = CL.biotemperature(np.array([10.0, -5.0, 35.0, 0.0]), np.array([0.0, 0.0, 0.0, 20.0]))
        np.testing.assert_allclose(bio[:3], [10.0, 0.0, 30.0])
        self.assertTrue(2.0 < bio[3] < 4.0)


class CalibrationTest(unittest.TestCase):
    def test_earthlike_midlatitudes_and_sea_ice_edge(self):
        g = small_grid(3)
        z = np.full(g.n, -4000.0)
        T, _ = CL.temperatures(g, z, z > 0, CL.DEFAULTS, 20.0)
        tm = (T["jun"] + T["dec"] + 2 * T["eq"]) / 4
        band = lambda a: np.abs(np.abs(g.lat) - a) < 2
        self.assertTrue(10 < tm[band(45)].mean() < 18, tm[band(45)].mean())
        self.assertTrue(-1 < tm[band(60)].mean() < 8, tm[band(60)].mean())
        winter = np.where(g.lat >= 0, T["dec"], T["jun"])
        self.assertGreater(winter[band(48)].mean(), -1.8)     # no sea ice at 48°
        self.assertLess(winter[band(66)].mean(), -1.8)        # seasonal sea ice by 66°

    def test_subtropical_land_heats_in_summer(self):
        g = small_grid(3)
        land = (g.lat > 10) & (g.lat < 40) & (np.abs(g.lon) < 50)
        z = np.where(land, 200.0, -4000.0)
        T, _ = CL.temperatures(g, z, land, CL.DEFAULTS, 20.0)
        anom = T["jun"] - CL.zonal_mean(g, T["jun"])
        self.assertGreater(anom[land].mean(), 2.5)
        self.assertLess(T["jun"][land].max(), 42.0)

    def test_ocean_seasonal_range_earthlike(self):
        g = small_grid(3)
        z = np.full(g.n, -4000.0)
        T, _ = CL.temperatures(g, z, z > 0, CL.DEFAULTS, 20.0)
        rng = np.abs(T["jun"] - T["dec"])[np.abs(np.abs(g.lat) - 50) < 2].mean()
        self.assertTrue(4 < rng < 10, rng)