import unittest

import numpy as np

from mapgen import climate as CL
from mapgen.sphere import east_north
from tests.helpers import make_ctx, small_grid


class WindTest(unittest.TestCase):
    def test_bands(self):
        g = small_grid(3)
        w = CL.band_winds(g, 0.0, CL.DEFAULTS)
        e, n = east_north(g.xyz)
        u = np.sum(w * e, axis=1)
        v = np.sum(w * n, axis=1)
        self.assertLess(u[np.abs(g.lat - 10) < 2].mean(), -3)     # trades from the east
        self.assertGreater(u[np.abs(g.lat - 40) < 2].mean(), 3)   # westerlies
        self.assertLess(v[np.abs(g.lat - 10) < 2].mean(), 0)      # NH trades flow equatorward
        self.assertGreater(v[np.abs(g.lat + 10) < 2].mean(), 0)   # SH trades flow equatorward

    def test_monsoon_reverses_onshore_flow(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)
        _, n = east_north(g.xyz)
        south = (g.lat > 0) & (g.lat < 8) & (np.abs(g.lon) < 40)
        vj = np.sum(CL.monsoon_winds(g, T["jun"], land, CL.DEFAULTS) * n, axis=1)[south].mean()
        vd = np.sum(CL.monsoon_winds(g, T["dec"], land, CL.DEFAULTS) * n, axis=1)[south].mean()
        self.assertGreater(vj, 0.5)
        self.assertLess(vd, -0.5)


class RainTest(unittest.TestCase):
    def test_itcz_wetter_than_subtropics_on_aquaplanet(self):
        g = small_grid(3)
        z = np.full(g.n, -4000.0)
        T, _ = CL.temperatures(g, z, z > 0, CL.DEFAULTS, 20.0)
        w = CL.band_winds(g, 0.0, CL.DEFAULTS)
        p = CL.precipitation(g, w, T["eq"], z, z > 0, 0.0, CL.DEFAULTS)
        self.assertGreater(p[np.abs(g.lat) < 5].mean(), 1.5 * p[np.abs(np.abs(g.lat) - 21) < 3].mean())

    def test_windward_wetter_than_lee(self):
        g = small_grid(3)
        land = (g.lat > 30) & (g.lat < 50) & (np.abs(g.lon) < 60)
        z = np.where(land, 200.0, -4000.0)
        z = np.where(land & (np.abs(g.lon) < 5), 3000.0, z)
        T, _ = CL.temperatures(g, z, land, CL.DEFAULTS, 20.0)
        w = CL.band_winds(g, 0.0, CL.DEFAULTS)
        p = CL.precipitation(g, w, T["eq"], z, land, 0.0, CL.DEFAULTS)
        band = (g.lat > 35) & (g.lat < 45)
        west = p[band & (g.lon > -12) & (g.lon < -3)].mean()
        east = p[band & (g.lon > 3) & (g.lon < 12)].mean()
        self.assertGreater(west, 1.3 * east)

    def test_stage_outputs(self):
        ctx = make_ctx(2)
        g = ctx.grid
        land = (np.abs(g.lat) < 40) & (np.abs(g.lon) < 50)
        ctx.data["elevation_eroded_m"] = np.where(land, 300.0, -4000.0).astype(np.float32)
        out = CL.run(ctx)
        for k in ("T_jun", "T_dec", "T_eq", "T_mean", "T_range", "T_min", "P_jun", "P_dec", "P_eq", "P_ann",
                  "wind_jun", "biotemp", "PET", "dist_ocean_km"):
            self.assertIn(k, out)
            self.assertTrue(np.all(np.isfinite(out[k])), k)
        mean = np.sum(out["P_ann"] * g.area_km2) / g.area_km2.sum()
        self.assertAlmostEqual(mean, 1000.0, delta=1.0)
        self.assertTrue(np.all(out["P_ann"] >= 0))


class StormTrackTest(unittest.TestCase):
    def test_midlatitudes_wetter_than_subtropics(self):
        g = small_grid(3)
        z = np.full(g.n, -4000.0)
        T, _ = CL.temperatures(g, z, z > 0, CL.DEFAULTS, 20.0)
        p = CL.precipitation(g, CL.band_winds(g, 0.0, CL.DEFAULTS), T["eq"], z, z > 0, 0.0, CL.DEFAULTS)
        a = np.abs(g.lat)
        h = CL.DEFAULTS["hadley_edge_deg"]        # storm track ≈ h+12..h+25, dry belt ≈ h−2..h+8
        self.assertGreater(p[(a > h + 12) & (a < h + 25)].mean(), 1.2 * p[(a > h - 2) & (a < h + 8)].mean())


class LandMoistureTest(unittest.TestCase):
    def test_equatorial_land_stays_wet(self):
        g = small_grid(3)
        land = (np.abs(g.lat) < 12) & (np.abs(g.lon) < 40)
        z = np.where(land, 300.0, -4000.0)
        T, _ = CL.temperatures(g, z, land, CL.DEFAULTS, 20.0)
        p = CL.precipitation(g, CL.band_winds(g, 0.0, CL.DEFAULTS), T["eq"], z, land, 0.0, CL.DEFAULTS)
        eq = np.abs(g.lat) < 8
        self.assertGreater(p[land & eq].mean(), 0.5 * p[~land & eq].mean())


class OceanMaskUseTest(unittest.TestCase):
    def test_inland_basin_is_land_for_climate(self):
        ctx = make_ctx(3)
        g = ctx.grid
        land = (np.abs(g.lat) < 40) & (np.abs(g.lon) < 60)
        z = np.where(land, 300.0, -4000.0)
        basin = (np.abs(g.lat) < 8) & (np.abs(g.lon) < 8)
        z[basin] = -30.0
        ctx.data.update({"elevation_eroded_m": z.astype(np.float32), "ocean": ~land})
        out = CL.run(ctx)
        self.assertTrue(np.all(out["dist_ocean_km"][basin] > 1000))


class InteriorRainTest(unittest.TestCase):
    def test_continental_interior_keeps_a_third_of_coastal_rain(self):
        from mapgen.graph import distance_to
        g = small_grid(3)
        land = (g.lat > 10) & (g.lat < 50) & (np.abs(g.lon) < 70)
        z = np.where(land, 300.0, -4000.0)
        T, _ = CL.temperatures(g, z, land, CL.DEFAULTS, 20.0)
        p = CL.precipitation(g, CL.band_winds(g, 0.0, CL.DEFAULTS), T["eq"], z, land, 0.0, CL.DEFAULTS)
        d = distance_to(g, ~land)
        coast, interior = p[land & (d < 500)].mean(), p[land & (d > 1500)].mean()
        self.assertGreater(interior, 0.33 * coast)          # was ≈0.25 before the desert retune