import json
import unittest

import numpy as np

from mapgen import seabed as SB
from mapgen.graph import distance_to
from tests.helpers import make_ctx


def sea_ctx(tect=None):
    ctx = make_ctx(3, tect=tect or {"plate": []})
    g = ctx.grid
    land = (np.abs(g.lat) < 30) & (np.abs(g.lon) < 40)
    ridge = np.abs(g.lon - 120) < 1.0
    d_div = distance_to(g, ridge)
    age = np.where(land, 0.0, np.minimum(d_div / 30.0, 200.0))
    ctx.data.update({"elevation_eroded_m": np.where(land, 500.0, -4500.0).astype(np.float32), "ocean": ~land,
                     "continental": land, "ocean_age_myr": age.astype(np.float32), "d_div_km": d_div,
                     "d_over_km": np.full(g.n, np.inf), "d_sub_km": np.full(g.n, np.inf),
                     "T_mean": 25.0 - 0.5 * np.abs(g.lat), "vel": np.zeros((g.n, 3))})
    return ctx


class SeabedTest(unittest.TestCase):
    def test_land_has_no_sea_floor(self):
        ctx = sea_ctx()
        out = SB.run(ctx)
        land = ~ctx.data["ocean"]
        for k in ("vent_potential", "seabed_type", "seabed_mineral", "bottom_temp_c", "sediment_m"):
            self.assertTrue(np.all(np.asarray(out[k])[land] == 0), k)

    def test_vents_on_the_ridge_not_far_from_it(self):
        ctx = sea_ctx()
        out = SB.run(ctx)
        d = ctx.data["d_div_km"]
        self.assertGreaterEqual(out["vent_potential"][d == 0].min(), 0.5)
        self.assertLess(out["vent_potential"][ctx.data["ocean"] & (d > 3000)].max(), 0.05)

    def test_sediment_thickens_with_age_away_from_land(self):
        ctx = sea_ctx()
        out = SB.run(ctx)
        g = ctx.grid
        young = ctx.data["ocean"] & (ctx.data["d_div_km"] < 300) & (np.abs(g.lon) > 90)
        old = ctx.data["ocean"] & (ctx.data["ocean_age_myr"] > 100) & (np.abs(g.lon) > 90)
        self.assertLess(out["sediment_m"][young].mean(), out["sediment_m"][old].mean())

    def test_bottom_temperature(self):
        ctx = sea_ctx()
        g = ctx.grid
        eq, polar = g.cell_index(0.0, 170.0), g.cell_index(-80.0, 170.0)
        z = np.asarray(ctx.data["elevation_eroded_m"]).copy()
        shallow = g.cell_index(-20.0, 170.0)
        z[shallow] = -100.0
        ctx.data["elevation_eroded_m"] = z
        out = SB.run(ctx)
        self.assertAlmostEqual(out["bottom_temp_c"][eq], 1.0 + 3.0 * np.cos(np.radians(g.lat[eq])) ** 2, places=4)
        self.assertLess(out["bottom_temp_c"][polar], 1.3)
        deep = 1.0 + 3.0 * np.cos(np.radians(g.lat[shallow])) ** 2
        t = 25.0 - 0.5 * abs(g.lat[shallow])
        self.assertAlmostEqual(out["bottom_temp_c"][shallow], deep + 0.875 * (t - deep), places=3)

    def test_plateau_volcanic_field_makes_volcanic_sea_floor(self):
        plat = {"name": "p", "center": [0.0, -150.0], "area_km2": 3.0e6, "top_m": [1500.0, 3000.0]}
        ctx = sea_ctx({"plate": [], "plateau": [plat]})
        out = SB.run(ctx)
        from mapgen.crust import center_dist
        near = center_dist(ctx.grid, [0.0, -150.0]) < 800.0
        self.assertTrue(np.isin(out["seabed_type"][near], [SB.SB_VOLCANIC, SB.SB_VENTS]).any())
        self.assertGreater(out["vent_potential"][near].max(), 0.3)

    def test_deterministic(self):
        a, b = SB.run(sea_ctx()), SB.run(sea_ctx())
        for k in a:
            np.testing.assert_array_equal(a[k], b[k])


class SeabedRenderTest(unittest.TestCase):
    def test_rasters_legends_and_plateaus_in_the_build(self):
        from mapgen.testing import built_world
        out = built_world() / "out" / "r2"
        fields = json.loads((out / "fields.json").read_text())
        for name in ("pressure", "o2_fraction", "fire", "vent_potential", "bottom_temp", "sediment",
                     "plant_height"):
            self.assertIn(name, fields["continuous"])
            self.assertTrue((out / "raster" / f"{name}.png").exists(), name)
        for name in ("seabed_type", "seabed_mineral"):
            self.assertIn(name, fields["categorical"])
        meta = json.loads((out / "cells_meta.json").read_text())
        self.assertEqual(meta["legends"]["seabed_type"], SB.SEABED_NAMES)
        self.assertEqual(meta["plateaus"], [])
        with np.load(out / "cells.npz") as z:
            for k in ("vent_potential", "seabed_type", "pressure_bar", "fire_reactivity", "plateau_id"):
                self.assertIn(k, z.files)