import unittest

import numpy as np

from mapgen import graph as G
from mapgen import hydrology as HY
from tests.helpers import make_ctx


def basin_ctx(p_basin, pet_basin):
    ctx = make_ctx(3)
    g = ctx.grid
    ocean = g.lat < -20
    z = np.where(ocean, -3000.0, 200.0 + 30.0 * (g.lat + 20))          # land rises northward
    c = g.xyz[g.cell_index(40.0, 0.0)]
    d = g.radius_km * np.arccos(np.clip(g.xyz @ c, -1, 1))
    z = np.where(d < 800, z - 1500 * (1 - d / 800), z)                  # closed basin
    ctx.data.update({"elevation_eroded_m": z.astype(np.float32), "P_ann": np.full(g.n, p_basin),
                     "PET": np.full(g.n, pet_basin)})
    return ctx, d


class HydrologyTest(unittest.TestCase):
    def test_strahler(self):
        recv = np.array([2, 2, 6, 5, 5, 6, 6])
        levels = G.receiver_levels(recv)
        o = HY.strahler(recv, levels, np.ones(7, bool))
        np.testing.assert_array_equal(o, [1, 1, 2, 1, 1, 2, 3])

    def test_arid_basin_is_endorheic_with_salt(self):
        ctx, d = basin_ctx(100.0, 1500.0)
        out = HY.run(ctx)
        inside = d < 300
        self.assertTrue(out["endorheic"][inside].all())
        self.assertTrue(out["salt_flat"][inside].any() or out["lake"][inside].any())

    def test_humid_basin_overflows(self):
        ctx, d = basin_ctx(2500.0, 400.0)
        out = HY.run(ctx)
        inside = d < 300
        self.assertTrue(out["lake"][inside].all())
        self.assertFalse(out["endorheic"][inside].any())

    def test_mass_balance_and_rivers(self):
        ctx, _ = basin_ctx(2500.0, 400.0)
        out = HY.run(ctx)
        g = ctx.grid
        roots = out["recv"] == np.arange(g.n)
        total = np.sum(out["runoff_mm"] * g.area_km2) * 1e-6
        self.assertAlmostEqual(out["discharge_km3_yr"][roots].sum(), total, delta=total * 1e-6)
        self.assertTrue(out["river"].any())
        self.assertTrue(np.all(out["strahler"][out["river"]] >= 1))
        land = ctx.data["elevation_eroded_m"] > 0
        nonendo = land & ~out["endorheic"] & ~roots
        self.assertTrue(np.all(out["z_filled_m"][out["recv"][nonendo]] < out["z_filled_m"][nonendo]))


class InlandDepressionTest(unittest.TestCase):
    def test_humid_depression_below_sea_level_becomes_lake(self):
        ctx, d = basin_ctx(2500.0, 400.0)
        g = ctx.grid
        z = ctx.data["elevation_eroded_m"].astype(np.float64)
        z[d < 300] = -40.0
        ctx.data["elevation_eroded_m"] = z.astype(np.float32)
        ctx.data["ocean"] = G.ocean_mask(g, z, 1.0e6)
        out = HY.run(ctx)
        self.assertTrue(out["lake"][d < 300].all())


def two_pit_ctx(p, pet):
    ctx = make_ctx(3)
    g = ctx.grid
    ocean = g.lat < -20
    z = np.where(ocean, -3000.0, 200.0 + 30.0 * (g.lat + 20))
    for lon in (-3.5, 3.5):
        c = g.xyz[g.cell_index(40.0, lon)]
        d = g.radius_km * np.arccos(np.clip(g.xyz @ c, -1, 1))
        z = np.where(d < 900, z - (1500 + 200 * (lon > 0)) * (1 - d / 900), z)   # two pits, one depression
    ctx.data.update({"elevation_eroded_m": z.astype(np.float32), "P_ann": np.full(g.n, p), "PET": np.full(g.n, pet)})
    return ctx, z


class TerminalTest(unittest.TestCase):
    def test_every_land_root_is_lake_or_salt(self):
        ctx, z = two_pit_ctx(700.0, 1400.0)
        out = HY.run(ctx)
        roots = (out["recv"] == np.arange(len(z))) & (z > 0) | (out["recv"] == np.arange(len(z))) & out["endorheic"]
        self.assertTrue(out["endorheic"].any())
        self.assertTrue(np.all(out["lake"][roots] | out["salt_flat"][roots]))
        self.assertEqual(int((roots & out["endorheic"]).sum()), 1)          # one terminal for the depression

    def test_mass_balance_with_lake_evaporation(self):
        ctx, z = two_pit_ctx(1000.0, 1200.0)
        out = HY.run(ctx)
        g = ctx.grid
        roots = out["recv"] == np.arange(g.n)
        total = np.sum(out["runoff_mm"] * g.area_km2) * 1e-6
        self.assertGreater(out["lake_loss_km3_yr"].sum(), 0)
        self.assertAlmostEqual(out["discharge_km3_yr"][roots].sum() + out["lake_loss_km3_yr"].sum(), total,
                               delta=total * 1e-6)


class LakeBedTest(unittest.TestCase):
    def test_lakes_carved_below_level_and_rerun_is_stable(self):
        ctx, d = basin_ctx(2500.0, 400.0)
        z0 = np.asarray(ctx.data["elevation_eroded_m"], dtype=np.float64)
        out = HY.run(ctx)
        lake, lev, z = out["lake"], out["lake_level_m"], out["elevation_eroded_m"].astype(np.float64)
        self.assertTrue(lake.any())
        self.assertTrue(np.all(np.isnan(lev[~lake])))
        self.assertTrue(np.all(z[lake] <= lev[lake] - HY.DEFAULTS["lake_min_m"] + 1e-3))
        np.testing.assert_array_equal(z[~lake], z0[~lake].astype(np.float32))   # only lake beds change
        ctx.data.update({"elevation_eroded_m": out["elevation_eroded_m"]})     # again (eras rerun hydrology)
        again = HY.run(ctx)
        np.testing.assert_array_equal(again["lake"], lake)
        np.testing.assert_allclose(again["lake_level_m"][lake], lev[lake])
        np.testing.assert_array_equal(again["elevation_eroded_m"], out["elevation_eroded_m"])

    def test_bigger_lakes_deeper(self):
        from tests.helpers import make_ctx as mk
        g = mk(3).grid
        P = {**HY.DEFAULTS, "lake_depth_exp": 0.3}
        z = np.zeros(g.n)
        lake = np.zeros(g.n, bool)
        small, big = g.cell_index(10.0, 0.0), np.flatnonzero(g.xyz @ g.xyz[g.cell_index(-30.0, 90.0)] > 0.97)
        lake[small] = lake[big] = True
        lid = np.where(lake, 0, -1)
        lid[big] = 1
        bed = HY.lake_beds(g, z, np.where(lake, 0.0, np.nan), lake, lid, np.zeros(g.n, bool), np.zeros(g.n, bool),
                           np.full(g.n, 1000.0), 1, P)
        self.assertLess(bed[big].min(), -P["lake_min_m"])
        self.assertLessEqual(bed[small], -P["lake_min_m"])
        self.assertTrue(np.all(bed[~lake] == 0))