import unittest

import numpy as np

from mapgen import crust as CR
from mapgen.plates import CONV, DIV
from tests.helpers import make_ctx

TECT = {"plate": [],
        "lip": [{"name": "l", "center": [10.0, -30.0], "radius_km": 1500.0}],
        "volcano": [{"name": "v", "center": [-10.0, 20.0], "radius_km": 700.0, "height_m": 7000.0,
                     "scar_azimuths": [90.0]}]}


def ctx3():
    ctx = make_ctx(3, tect=TECT, cfg={"crust": {"edge_noise": 0.0}})
    g = ctx.grid
    land = ((np.abs(g.lat) < 35) & (np.abs(g.lon) < 60)).astype(float)
    bt = np.zeros(g.n, np.int8)
    bt[np.abs(g.lon - 120) < 1.0] = DIV   # a mid-ocean ridge along lon 120
    bt[np.abs(g.lon - 0) < 1.0] = CONV
    ctx.data.update({"sk_land": land, "m_land_hint": np.zeros(g.n), "bnd_type": bt})
    return ctx


class CrustTest(unittest.TestCase):
    def test_classes(self):
        ctx = ctx3()
        out = CR.run(ctx)
        g = ctx.grid
        age = out["age_class"]
        self.assertTrue(out["continental"][g.cell_index(20.0, -50.0)])
        self.assertFalse(out["continental"][g.cell_index(0.0, 150.0)])
        self.assertEqual(age[g.cell_index(0.0, 150.0)], CR.OCEANIC)
        self.assertEqual(age[g.cell_index(10.0, -30.0)], CR.LIP)
        self.assertEqual(age[g.cell_index(-10.0, 20.0)], CR.VOLCANO)
        self.assertEqual(age[g.cell_index(-10.0, 24.0)], CR.SCAR)      # east sector, ~900 km out
        self.assertEqual(age[g.cell_index(20.0, 3.0)], CR.POST_OROGEN)

    def test_ocean_age_grows_from_ridge(self):
        ctx = ctx3()
        out = CR.run(ctx)
        g = ctx.grid
        ocean = ~out["continental"] & (np.abs(g.lat) < 20) & (np.abs(g.lon - 120) < 20)
        r = np.corrcoef(np.abs(g.lon[ocean] - 120), out["ocean_age_myr"][ocean])[0, 1]
        self.assertGreater(r, 0.95)
        self.assertTrue(np.all(out["ocean_age_myr"][out["continental"]] == 0))


class ShelfTest(unittest.TestCase):
    def test_continental_shelf_is_resolution_independent(self):
        from mapgen.graph import distance_to
        for res in (3, 4):
            ctx = make_ctx(res, tect={"plate": []}, cfg={"crust": {"edge_noise": 0.0}})
            g = ctx.grid
            land = ((np.abs(g.lat) < 30) & (np.abs(g.lon) < 40)).astype(float)
            ctx.data.update({"sk_land": land, "m_land_hint": np.zeros(g.n), "bnd_type": np.zeros(g.n, np.int8)})
            cont = CR.run(ctx)["continental"]
            d = distance_to(g, land > 0.5)
            self.assertTrue(np.all(cont[land > 0.5]), res)
            self.assertGreater(cont[(d > 0) & (d < 300)].mean(), 0.8, res)     # shelf
            self.assertEqual(cont[d > 700].sum(), 0, res)                      # but not far out


class FractalMarginTest(unittest.TestCase):
    @staticmethod
    def _edge(noise):
        from mapgen.graph import components
        ctx = make_ctx(4, tect={"plate": []}, cfg={"crust": {"edge_noise": noise}})
        g = ctx.grid
        land = ((np.abs(g.lat) < 30) & (np.abs(g.lon) < 60)).astype(float)
        ctx.data.update({"sk_land": land, "m_land_hint": np.zeros(g.n), "bnd_type": np.zeros(g.n, np.int8)})
        cont = CR.run(ctx)["continental"]
        lab = components(g, cont)
        sizes = np.bincount(lab[cont])
        return int(np.sum(cont[g.src] != cont[g.dst])), int(np.sum((sizes > 0) & (sizes < 0.05 * sizes.max())))

    def test_margins_are_fractal_with_fragments(self):
        smooth, _ = self._edge(0.0)
        edges, fragments = self._edge(CR.DEFAULTS["edge_noise"])
        self.assertGreater(edges, 1.4 * smooth)
        self.assertGreaterEqual(fragments, 3)


class RevisionCrustTest(unittest.TestCase):
    def test_no_patch_no_plateau_keeps_crust(self):
        out = CR.run(ctx3())
        np.testing.assert_array_equal(out["continental"], out["continental_base"])
        self.assertTrue(np.all(out["plateau_id"] == -1))

    def test_land_patch_adds_continental_crust_only_around_its_disc(self):
        ctx = ctx3()
        ctx.tect = {**TECT, "land_patch": [{"name": "fill", "center": [0.0, 150.0], "radius_km": 2000.0,
                                            "edge_noise": 0.4}]}
        out = CR.run(ctx)
        g = ctx.grid
        i = g.cell_index(0.0, 150.0)
        self.assertTrue(out["continental"][i])
        self.assertFalse(out["continental_base"][i])
        far = CR.center_dist(g, [0.0, 150.0]) > 4500.0
        np.testing.assert_array_equal(out["continental"][far], out["continental_base"][far])
        self.assertEqual(out["age_class"][i] == CR.OCEANIC, False)

    def test_plateau_cells_become_continental_crust_with_ids(self):
        ctx = ctx3()
        ctx.tect = {**TECT, "plateau": [{"name": "p", "center": [0.0, 150.0], "area_km2": 3.0e6,
                                         "top_m": [1500.0, 3000.0]}]}
        out = CR.run(ctx)
        i = ctx.grid.cell_index(0.0, 150.0)
        self.assertEqual(out["plateau_id"][i], 0)
        self.assertTrue(out["continental"][i])
        self.assertFalse(out["continental_base"][i])
        self.assertNotEqual(out["age_class"][i], CR.OCEANIC)
        self.assertGreater(out["ocean_age_myr"][i], 0.0, "the plateau keeps the age of the sea floor around it")