import unittest

import numpy as np

from mapgen import ocean as OC
from mapgen.sphere import east_north
from tests.helpers import small_grid

DAY = 31.149
KM_PER_DEG = 12742.0 * np.pi / 180.0


def basin(g):
    """Ocean box 15–45° N between two meridional coasts at ±60° lon (everything else land)."""
    return (g.lat > 15) & (g.lat < 45) & (np.abs(g.lon) < 60)


def gyre_wind(g):
    """Trades (from the east) at 15° |lat|, westerlies at 45° |lat|: −8 … +8 m/s east."""
    e, _ = east_north(g.xyz)
    U = -8.0 * np.cos(np.radians((np.abs(g.lat) - 15.0) / 30.0 * 180.0))
    return U[:, None] * e


def params(**over):
    P = dict(OC.DEFAULTS)
    P["friction_days"] = 2.0          # r/β ≈ 750 km: resolved on the res-3 test grid (≈240 km spacing)
    P.update(over)
    return P


def northward(g, u):
    _, n = east_north(g.xyz)
    return np.sum(u * n, axis=1)


class GyreTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.g = small_grid(3)
        cls.ocean = basin(cls.g)
        cls.u = OC.currents(cls.g, cls.ocean, gyre_wind(cls.g), params(), DAY)

    def test_clockwise_with_western_intensification(self):
        g, v = self.g, northward(self.g, self.u)
        band = self.ocean & (g.lat > 25) & (g.lat < 35)
        west_km = (g.lon + 60.0) * KM_PER_DEG * np.cos(np.radians(30.0))
        east_km = (60.0 - g.lon) * KM_PER_DEG * np.cos(np.radians(30.0))
        west, east = band & (west_km < 1500), band & (east_km < 1500)
        self.assertGreater(v[west].max(), 0.0)                          # northward along the west coast
        self.assertLess(v[east].min(), 0.0)                             # southward in the east
        self.assertGreater(v[west].max(), 3.0 * -v[east].min())         # the western boundary current is the fast one

    def test_land_is_still(self):
        g = self.g
        psi = OC.streamfunction(g, self.ocean, gyre_wind(g), params(), DAY)
        sp = np.linalg.norm(OC.velocity(g, psi), axis=1)
        inland = ~self.ocean
        for _ in range(2):                                              # two cells away from any sea
            inland = inland & ~np.bincount(g.src, weights=(~inland[g.dst]).astype(float), minlength=g.n).astype(bool)
        self.assertLess(sp[inland].max(), 0.01 * sp[self.ocean].max())

    def test_zero_on_land_in_output(self):
        self.assertTrue(np.all(self.u[~self.ocean] == 0.0))

    def test_southern_hemisphere_is_anticlockwise(self):
        g = self.g
        ocean = (g.lat < -15) & (g.lat > -45) & (np.abs(g.lon) < 60)
        v = northward(g, OC.currents(g, ocean, gyre_wind(g), params(), DAY))
        band = ocean & (g.lat < -25) & (g.lat > -35)
        west = band & (g.lon < -45)
        self.assertLess(v[west].min(), 0.0)                             # southward along the west coast
        self.assertGreater(-v[west].min(), 3.0 * max(v[band & (g.lon > 45)].max(), 1e-9))


class ChannelTest(unittest.TestCase):
    def test_westerlies_drive_eastward_flow_downwind_without_rotation(self):
        g = small_grid(3)
        ocean = (g.lat > 30) & (g.lat < 60)
        e, _ = east_north(g.xyz)
        wind = 8.0 * e
        mid = (g.lat > 40) & (g.lat < 50)
        u = OC.currents(g, ocean, wind, params(), 1e7)                  # ~no rotation: flow is downwind
        ue, un = np.sum(u * e, axis=1), northward(g, u)
        self.assertGreater(ue[mid].mean(), 0.0)
        self.assertLess(np.abs(un[mid]).mean(), 0.05 * ue[mid].mean())
        u = OC.currents(g, ocean, wind, params(), DAY)
        self.assertGreater(np.sum(u * e, axis=1)[mid].mean(), 0.0)


class DayLengthTest(unittest.TestCase):
    def test_slower_rotation_widens_boundary_current(self):
        g = small_grid(3)
        ocean = basin(g)
        band = ocean & (g.lat > 27) & (g.lat < 33)
        widths = []
        for day in (DAY, 4 * DAY):
            v = northward(g, OC.currents(g, ocean, gyre_wind(g), params(friction_days=4.0), day))
            half = band & (v > 0.5 * v[band].max())
            widths.append(g.lon[half].max() - g.lon[band].min())
        self.assertGreater(widths[1], widths[0])


class EdgeCaseTest(unittest.TestCase):
    def test_no_ocean_gives_zeros(self):
        g = small_grid(2)
        u = OC.currents(g, np.zeros(g.n, bool), gyre_wind(g), params(), DAY)
        self.assertTrue(np.all(u == 0.0))

    def test_all_ocean_globe_is_finite_at_the_poles(self):
        g = small_grid(3)
        e, _ = east_north(g.xyz)
        wind = (-6.0 * np.cos(np.radians(3 * g.lat)))[:, None] * e
        u = OC.currents(g, np.ones(g.n, bool), wind, params(), DAY)
        self.assertTrue(np.all(np.isfinite(u)))
        self.assertLess(np.linalg.norm(u, axis=1).max(), 20.0)


class CoarseTest(unittest.TestCase):
    def test_coarse_fallback_keeps_the_gyre(self):
        g = small_grid(3)
        ocean = basin(g)
        v = northward(g, OC.currents(g, ocean, gyre_wind(g), params(friction_days=1.0, direct_max_cells=1000), DAY))
        band = ocean & (g.lat > 25) & (g.lat < 35)
        self.assertTrue(np.all(np.isfinite(v)))
        self.assertGreater(v[band & (g.lon < -40)].mean(), 0.0)
        self.assertLess(v[band & (g.lon > -20)].mean(), 0.0)


class SSTTest(unittest.TestCase):
    def test_still_water_keeps_equilibrium(self):
        g = small_grid(3)
        ocean = basin(g)
        T_eq = 30.0 - 0.5 * np.abs(g.lat)
        T = OC.sst(g, ocean, np.zeros((g.n, 3)), T_eq, params(kappa_m2s=0.0))
        np.testing.assert_allclose(T, T_eq, atol=1e-6)

    def test_boundary_current_warms_the_west_side(self):
        g = small_grid(3)
        ocean = basin(g)
        u = OC.currents(g, ocean, gyre_wind(g), params(), DAY)
        T_eq = 30.0 - 0.5 * np.abs(g.lat)
        a = OC.sst(g, ocean, u, T_eq, params()) - T_eq
        band = ocean & (g.lat > 32) & (g.lat < 40)
        self.assertGreater(a[band & (g.lon < -45)].mean(), a[band & (g.lon > 45)].mean() + 0.2)
        self.assertGreater(a[band & (g.lon < -45)].mean(), 0.0)


class UpwellingTest(unittest.TestCase):
    def setUp(self):
        self.g = small_grid(3)
        _, self.n = east_north(self.g.xyz)
        self.P = params(upwell_coast_km=0.0)

    def coast(self, ocean, lo, hi):
        g = self.g
        touches_land = np.bincount(g.src, weights=(~ocean[g.dst]).astype(float), minlength=g.n) > 0
        return ocean & touches_land & (g.lat > lo) & (g.lat < hi)

    def test_west_coast_upwells(self):
        g = self.g
        ocean = ~((g.lon > 0) & (g.lon < 90) & (g.lat > 10) & (g.lat < 50))   # continent east of the sea
        w = OC.upwelling(g, ocean, -6.0 * self.n, self.P, DAY)                 # equatorward wind (north)
        self.assertGreater(w[self.coast(ocean, 20, 40) & (g.lon < 0)].mean(), 0.0)

    def test_east_coast_downwells(self):
        g = self.g
        ocean = ~((g.lon > -90) & (g.lon < 0) & (g.lat > 10) & (g.lat < 50))  # continent west of the sea
        w = OC.upwelling(g, ocean, -6.0 * self.n, self.P, DAY)
        self.assertLess(w[self.coast(ocean, 20, 40) & (g.lon > 0)].mean(), 0.0)

    def test_southern_west_coast_upwells_under_equatorward_wind(self):
        g = self.g
        ocean = ~((g.lon > 0) & (g.lon < 90) & (g.lat < -10) & (g.lat > -50))
        w = OC.upwelling(g, ocean, 6.0 * self.n, self.P, DAY)                  # equatorward in the south = north
        self.assertGreater(w[self.coast(ocean, -40, -20) & (g.lon < 0)].mean(), 0.0)

    def test_trades_upwell_at_the_equator(self):
        g = self.g
        e, _ = east_north(g.xyz)
        w = OC.upwelling(g, np.ones(g.n, bool), -6.0 * e, self.P, DAY)
        self.assertGreater(w[np.abs(g.lat) < 3].mean(), 0.0)
        self.assertLess(w[(np.abs(g.lat) > 8) & (np.abs(g.lat) < 20)].mean(), w[np.abs(g.lat) < 3].mean())

    def test_zero_on_land(self):
        g = self.g
        ocean = g.lat < 0
        self.assertTrue(np.all(OC.upwelling(g, ocean, -6.0 * self.n, params(), DAY)[~ocean] == 0.0))


class ProductivityTest(unittest.TestCase):
    def test_upwelling_coast_beats_gyre_centre(self):
        g = small_grid(3)
        ocean = ~((g.lon > 0) & (g.lon < 90) & (g.lat > 10) & (g.lat < 50))
        _, n = east_north(g.xyz)
        w = OC.upwelling(g, ocean, -6.0 * n, params(), DAY)
        p = OC.productivity(g, ocean, w, np.full(g.n, -4000.0), np.zeros(g.n), np.full(g.n, 20.0), params())
        touches_land = np.bincount(g.src, weights=(~ocean[g.dst]).astype(float), minlength=g.n) > 0
        coast = ocean & touches_land & (g.lat > 25) & (g.lat < 35) & (g.lon < 0)
        centre = g.cell_index(30.0, -60.0)
        self.assertGreater(p[coast].max(), 0.1)
        self.assertGreater(p[coast].max(), 10.0 * p[centre])

    def test_shelf_and_light(self):
        g = small_grid(3)
        ocean = g.lat < 0
        zero = np.zeros(g.n)
        shelf = OC.productivity(g, ocean, zero, np.full(g.n, -100.0), zero, np.full(g.n, 25.0), params())
        deep = OC.productivity(g, ocean, zero, np.full(g.n, -3000.0), zero, np.full(g.n, 25.0), params())
        i = g.cell_index(-10.0, 0.0)
        self.assertAlmostEqual(shelf[i], 0.4 * (0.3 + 0.7 * np.cos(np.radians(g.lat[i]))), places=6)  # a_s·light
        self.assertEqual(deep[i], 0.0)
        self.assertGreater(shelf[i], shelf[g.cell_index(-80.0, 0.0)])

    def test_sharp_sst_front_is_productive(self):
        g = small_grid(3)
        ocean = g.lat < 0
        zero, deep = np.zeros(g.n), np.full(g.n, -4000.0)
        front = 15.0 + 8.0 * np.tanh((g.lat + 35.0) / 2.0)             # 16 °C across ≈ 4° (≈ 900 km): ≈ 2 °C/100 km
        flat = 25.0 + 0.1 * g.lat                                       # background pole-ward cooling, ≈ 0.05 °C/100 km
        pf = OC.productivity(g, ocean, zero, deep, zero, front, params())
        pb = OC.productivity(g, ocean, zero, deep, zero, flat, params())
        at, far = g.cell_index(-35.0, 0.0), g.cell_index(-60.0, 0.0)
        self.assertGreater(pf[at], 0.15)
        self.assertLess(pf[far], 0.02)
        self.assertLess(pb.max(), 1e-9)                                 # gentle background gradients add nothing

    def test_land_sea_contrast_is_not_a_front(self):
        g = small_grid(3)
        ocean = g.lat < 0
        zero, deep = np.zeros(g.n), np.full(g.n, -4000.0)
        sst_c = np.where(ocean, 20.0, -10.0)                            # land temperatures differ wildly
        self.assertLess(OC.productivity(g, ocean, zero, deep, zero, sst_c, params()).max(), 1e-9)

    def test_range_and_land(self):
        g = small_grid(3)
        ocean = g.lat < 0
        p = OC.productivity(g, ocean, np.full(g.n, 1e4), np.zeros(g.n), np.full(g.n, 40.0), np.zeros(g.n), params())
        self.assertTrue(np.all((p >= 0) & (p <= 1)))
        self.assertTrue(np.all(p[~ocean] == 0))


class GaugeTest(unittest.TestCase):
    def test_no_spurious_vortex_at_the_gauge(self):
        # cell 0 (79° N, 38° E on H3 grids) at sea in a world with a continent: pinning ψ there must not leave a point
        # vortex (the solve's compatibility residual) — the speed around cell 0 stays within the ocean's p99
        g = small_grid(3)
        land = (np.abs(g.lat) < 40) & (np.abs(g.lon) < 50)
        ocean = ~land
        self.assertTrue(ocean[0])
        e, _ = east_north(g.xyz)
        wind = (-8.0 * np.cos(np.radians(3.0 * g.lat)))[:, None] * e
        sp = np.linalg.norm(OC.currents(g, ocean, wind, params(), DAY), axis=1)
        near = np.zeros(g.n, bool)
        near[0] = True
        for _ in range(2):
            near = near | (np.bincount(g.src, weights=near[g.dst].astype(float), minlength=g.n) > 0)
        self.assertLessEqual(sp[near].max(), np.percentile(sp[ocean], 99))