import unittest

import numpy as np

from mapgen import fields as FI
from mapgen import ice as IC
from tests.helpers import make_ctx


class IceTest(unittest.TestCase):
    def test_sheet_glacier_sea_ice(self):
        ctx = make_ctx(3)
        g = ctx.grid
        polar_land = g.lat > 70
        z = np.where(polar_land, 500.0, -3000.0)
        peak = g.cell_index(0.0, 0.0)
        z[peak] = 6000.0
        t_mean = np.where(g.lat > 60, -25.0, 15.0)
        t_mean[peak] = -9.0          # too warm for a sheet, summer below zero → glacier
        t_sum = t_mean + 8
        t_win = t_mean - 8
        ctx.data.update({"elevation_eroded_m": z, "T_mean": t_mean, "T_jun": t_sum, "T_dec": t_win,
                         "P_ann": np.full(g.n, 300.0)})
        out = IC.run(ctx)
        self.assertEqual(out["ice"][g.cell_index(85.0, 0.0)], IC.ICE_SHEET)
        self.assertEqual(out["ice"][peak], IC.ICE_GLACIER)
        self.assertEqual(out["ice"][g.cell_index(65.0, 0.0)], IC.ICE_SEA_PERENNIAL)
        self.assertEqual(out["ice"][g.cell_index(0.0, 150.0)], IC.ICE_NONE)
        i = g.cell_index(85.0, 0.0)
        self.assertGreater(out["z_surface_m"][i], z[i] + 200)


class FieldsTest(unittest.TestCase):
    def test_o2_and_gravity(self):
        ctx = make_ctx(1)
        n = ctx.grid.n
        zs = np.zeros(n)
        zs[1] = 8000.0
        m_o2 = np.zeros(n)
        m_o2[2] = 1.0
        m_g = np.zeros(n)
        m_g[3] = -1.0
        ctx.data.update({"z_surface_m": zs, "m_o2_zones": m_o2, "m_gravity_zones": m_g})
        out = FI.run(ctx)
        self.assertAlmostEqual(out["po2_bar"][0], 0.21)
        self.assertAlmostEqual(out["po2_bar"][1], 0.21 * np.exp(-1), places=6)
        self.assertAlmostEqual(out["po2_bar"][2], 0.21 * 1.5)
        self.assertAlmostEqual(out["gravity_g"][0], 1.05)
        self.assertAlmostEqual(out["gravity_g"][3], 1.05 * 0.3)

    def test_pressure_o2_fraction_and_fire(self):
        ctx = make_ctx(1)
        n = ctx.grid.n
        zs = np.zeros(n)
        zs[1] = 8000.0
        m_o2 = np.zeros(n)
        m_o2[2] = 1.0
        ctx.data.update({"z_surface_m": zs, "m_o2_zones": m_o2, "m_gravity_zones": np.zeros(n)})
        out = FI.run(ctx)
        self.assertAlmostEqual(out["pressure_bar"][0], 1.0)
        self.assertAlmostEqual(out["pressure_bar"][1], np.exp(-1))
        self.assertAlmostEqual(out["o2_fraction"][2], 0.21 * 1.5)
        np.testing.assert_allclose(out["po2_bar"], out["o2_fraction"] * out["pressure_bar"])
        self.assertTrue(np.all(out["fire_reactivity"] == 1.0))
        np.testing.assert_allclose(out["plant_height_x"], ctx.cfg["planet"]["gravity_g"] / out["gravity_g"])
        np.testing.assert_allclose(FI.pressure(2.0, np.array([-100.0, 8000.0]), 8000.0), [2.0, 2.0 * np.exp(-1)])

class IceOceanMaskTest(unittest.TestCase):
    def test_cold_inland_basin_is_not_sea_ice(self):
        ctx = make_ctx(3)
        g = ctx.grid
        land = g.lat > 50
        z = np.where(land, 300.0, -3000.0)
        basin = (g.lat > 60) & (g.lat < 68) & (np.abs(g.lon) < 20)
        z[basin] = -20.0
        t = np.where(g.lat > 50, -5.0, 10.0)
        ctx.data.update({"elevation_eroded_m": z, "ocean": ~land, "T_mean": t, "T_jun": t + 8, "T_dec": t - 8,
                         "P_ann": np.full(g.n, 300.0)})
        out = IC.run(ctx)
        self.assertFalse(np.isin(out["ice"][basin], [IC.ICE_SEA_SEASONAL, IC.ICE_SEA_PERENNIAL]).any())


class SummerMeltRuleTest(unittest.TestCase):
    def test_ice_sheets_need_summers_below_freezing(self):
        ctx = make_ctx(3)
        g = ctx.grid
        land = g.lat > 55
        z = np.where(land, 300.0, -3000.0)
        t_mean = np.where(g.lat > 75, -25.0, -12.0)           # 55–75°: cold on average, but summers melt
        t_sum = np.where(g.lat > 75, -8.0, 4.0)
        peak = g.cell_index(0.0, 0.0)
        z[peak], t_mean[peak], t_sum[peak] = 6000.0, -9.0, -1.0
        ctx.data.update({"elevation_eroded_m": z, "ocean": ~(land | (np.arange(g.n) == peak)), "T_mean": t_mean,
                         "T_jun": t_sum, "T_dec": t_mean - 10, "P_ann": np.full(g.n, 300.0)})
        out = IC.run(ctx)
        self.assertEqual(out["ice"][g.cell_index(65.0, 0.0)], IC.ICE_NONE)        # tundra, not ice
        self.assertEqual(out["ice"][g.cell_index(85.0, 0.0)], IC.ICE_SHEET)
        self.assertEqual(out["ice"][peak], IC.ICE_GLACIER)                          # small cold spot = glacier