import unittest

import numpy as np

from mapgen import plates as PL
from mapgen.pipeline import StageError
from tests.helpers import make_ctx


def two_plates(az_a, az_b, speed=5.0, kinds=("oceanic", "oceanic")):
    return {"plate": [
        {"id": "a", "seed": [0.0, -40.0], "kind": kinds[0], "motion": [az_a, speed]},
        {"id": "b", "seed": [0.0, 40.0], "kind": kinds[1], "motion": [az_b, speed]},
    ]}


def ctx_for(tect, land=None, res=2):
    ctx = make_ctx(res, tect=tect)
    n = ctx.grid.n
    ctx.data.update({"sk_land": np.zeros(n) if land is None else land, "m_land_hint": np.zeros(n)})
    return ctx


class PlatesTest(unittest.TestCase):
    def test_convergent_then_divergent(self):
        for az_a, az_b, want in ((90.0, 270.0, PL.CONV), (270.0, 90.0, PL.DIV)):
            ctx = ctx_for(two_plates(az_a, az_b))
            out = PL.run(ctx)
            g = ctx.grid
            near = (np.abs(g.lon) < 12) & (np.abs(g.lat) < 30) & (out["bnd_type"] > 0)
            self.assertGreater(near.sum(), 3)
            self.assertTrue(np.mean(out["bnd_type"][near] == want) > 0.8, (az_a, want))

    def test_every_cell_assigned_and_seeds_own_cells(self):
        ctx = ctx_for(two_plates(90.0, 270.0))
        out = PL.run(ctx)
        g = ctx.grid
        self.assertEqual(set(np.unique(out["plate"])), {0, 1})
        self.assertEqual(out["plate"][g.cell_index(0.0, -40.0)], 0)
        self.assertEqual(out["plate"][g.cell_index(0.0, 40.0)], 1)

    def test_continent_stays_on_its_plate(self):
        tect = {"plate": [
            {"id": "c", "seed": [0.0, 0.0], "kind": "continental", "motion": [0.0, 1.0]},
            {"id": "o", "seed": [0.0, 60.0], "kind": "oceanic", "motion": [0.0, 1.0]},
        ]}
        ctx = make_ctx(2, tect=tect)
        g = ctx.grid
        land = ((np.abs(g.lat) < 20) & (np.abs(g.lon) < 40)).astype(float)
        ctx.data.update({"sk_land": land, "m_land_hint": np.zeros(g.n)})
        out = PL.run(ctx)
        self.assertTrue(np.all(out["plate"][land > 0.5] == 0))

    def test_duplicate_seed_cell_errors(self):
        tect = two_plates(90.0, 270.0)
        tect["plate"][1]["seed"] = [0.01, -40.01]
        with self.assertRaisesRegex(StageError, "same cell"):
            PL.run(ctx_for(tect))

    def test_zero_speed_plate(self):
        out = PL.run(ctx_for(two_plates(90.0, 270.0, speed=0.0)))
        self.assertTrue(np.all(np.isfinite(out["vel"])))
        self.assertTrue(np.allclose(out["vel"], 0.0))


class MeanderTest(unittest.TestCase):
    def test_boundaries_meander(self):
        ctx = ctx_for(two_plates(90.0, 270.0, kinds=("continental", "continental")), res=4)
        g = ctx.grid
        ctx.data["sk_land"] = ((np.abs(g.lat) < 40) & (np.abs(g.lon) < 70)).astype(float)
        out = PL.run(ctx)
        b = (out["bnd_type"] > 0) & (np.abs(g.lat) < 35) & (np.abs(g.lon) < 40)
        self.assertGreater(g.lon[b].std(), 3.0)