import unittest
from pathlib import Path

import numpy as np

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

GOLDEN = Path(__file__).parent / "data" / "climate_golden_r2.npz"
NEW = ("current", "current_speed", "sst", "upwelling", "productivity")


def golden_ctx(cfg=None):
    ctx = make_ctx(2, cfg=cfg)
    g = ctx.grid
    land = (np.abs(g.lat) < 40) & (np.abs(g.lon) < 50)
    ctx.data["elevation_eroded_m"] = np.where(land, np.where(np.abs(g.lon) < 10, 2500.0, 300.0),
                                              -4000.0).astype(np.float32)
    return ctx, land


class OceanDisabledTest(unittest.TestCase):
    def test_disabled_reproduces_pre_ocean_climate(self):
        ctx, land = golden_ctx({"ocean": {"enabled": False}})
        out = CL.run(ctx)
        ref = np.load(GOLDEN)
        for k in ref.files:
            np.testing.assert_array_equal(out[k], ref[k], err_msg=k)
        self.assertTrue(np.all(out["current"] == 0) and np.all(out["productivity"] == 0))
        np.testing.assert_array_equal(out["sst"], out["T_mean"].astype(np.float32))


class OceanCoupledTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.ctx, cls.land = golden_ctx()
        cls.out = CL.run(cls.ctx)

    def test_new_fields(self):
        g, out = self.ctx.grid, self.out
        for k in NEW:
            self.assertIn(k, out)
            self.assertEqual(out[k].dtype, np.float32, k)
            self.assertTrue(np.all(np.isfinite(out[k])), k)
        self.assertEqual(out["current"].shape, (g.n, 3))
        self.assertTrue(np.all(out["current"][self.land] == 0))
        np.testing.assert_allclose(out["current_speed"], np.linalg.norm(out["current"], axis=1), rtol=1e-5)
        sea = ~self.land
        self.assertTrue(np.all((out["sst"][sea] >= -1.8) & (out["sst"][sea] < 40)))            # sea water freezes at −1.8 °C

    def test_currents_change_the_climate(self):
        ref = np.load(GOLDEN)
        self.assertGreater(np.abs(self.out["T_mean"] - ref["T_mean"]).max(), 0.5)   # the old gyre rule is gone

    def test_winds_come_from_current_free_temperatures(self):
        ctx, _ = golden_ctx({"ocean": {"relax_days": 30.0}})                      # other SST, same winds
        np.testing.assert_array_equal(CL.run(ctx)["wind_jun"], self.out["wind_jun"])

    def test_unknown_ocean_key_is_an_error(self):
        from mapgen.config import ConfigError
        ctx, _ = golden_ctx({"ocean": {"frictoin_days": 3.0}})
        with self.assertRaises(ConfigError):
            CL.run(ctx)

    def test_all_land_world(self):
        ctx = make_ctx(2)
        ctx.data["elevation_eroded_m"] = np.full(ctx.grid.n, 500.0, np.float32)
        out = CL.run(ctx)
        self.assertTrue(np.all(out["current"] == 0))

    def test_locked_world_has_still_ocean(self):
        ctx, land = golden_ctx({"climate": {"lock": True}})
        out = CL.run(ctx)
        self.assertTrue(np.all(out["current"] == 0) and np.all(out["upwelling"] == 0))
        np.testing.assert_array_equal(out["sst"], out["T_mean"].astype(np.float32))

    def test_upwelling_and_productivity_filled(self):
        sea = ~self.land
        self.assertGreater(np.abs(self.out["upwelling"][sea]).max(), 1.0)
        self.assertGreater(self.out["productivity"][sea].max(), 0.05)
        self.assertTrue(np.all(self.out["productivity"][self.land] == 0))