import json
import math
import tempfile
import unittest
from pathlib import Path

import numpy as np
from PIL import Image

import export as EX
import serve
import tiles as T
from mapgen import config as C
from tests.test_serve import built_world


def _spot():
    """A river mouth on the test world: land, sea and a river in one small square."""
    a = np.load(built_world() / "out" / "r2" / "cells.npz")
    mouths = np.where(a["river"].astype(bool) & a["ocean"][a["recv"]].astype(bool))[0]
    i = int(mouths[np.argmax(a["discharge_km3_yr"][mouths])])
    t = np.linspace(0, 1, 101)[:, None]                          # from the mouth toward its sea cell …
    p = a["g_xyz"][i] * (1 - t) + a["g_xyz"][a["recv"][i]] * t
    p /= np.linalg.norm(p, axis=1, keepdims=True)
    lat, lon = np.degrees(np.arcsin(p[:, 2])), np.degrees(np.arctan2(p[:, 1], p[:, 0]))
    src = T.TileSource(serve.World(built_world(), 2), seed=int(C.load(built_world())[0]["build"]["seed"]),
                       cache_dir=Path(tempfile.mkdtemp()), regions_dir=Path(tempfile.mkdtemp()))
    wet = src.surface_at(lat, lon, 1.0)[1]
    k = int(np.argmax(wet)) if wet.any() else 50               # … to where the map shows the sea: the coast
    return float(lat[k]), float(lon[k])


class ExportTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.root = built_world()
        cls.lat, cls.lon = _spot()
        cls.pins = Path(tempfile.mkdtemp()) / "pins.json"
        cls.pins.write_text(json.dumps({"pins": [{"id": "pin-1", "name": "Here", "lat": cls.lat, "lon": cls.lon,
                                                  "lore": "somewhere", "note": "", "epoch": None},
                                                 {"id": "pin-2", "name": "Far", "lat": -cls.lat, "lon": cls.lon + 90}]}))
        cls.out = Path(tempfile.mkdtemp())
        cls.d = EX.export(cls.root, 2, cls.lat, cls.lon, size_km=200, res_m=2000, out_dir=cls.out, name="t",
                          pins_path=cls.pins, regions_dir=Path(tempfile.mkdtemp()))
        cls.meta = json.loads((cls.d / "meta.json").read_text())
        cls.n = cls.meta["samples"]
        cls.h = np.fromfile(cls.d / "height.f32", dtype="<f4").reshape(cls.n, cls.n)

    def test_files_and_grid(self):
        self.assertEqual(self.d, self.out / "t")
        for f in ("height.f32", "height.png", "height.r16", "water.png", "biome.png", "ground.png", "landform.png",
                  "preview.png", "rivers.json", "pins.json", "legend.json", "meta.json", "README.txt"):
            self.assertTrue((self.d / f).exists(), f)
        self.assertEqual(self.n, 129, "2^n + 1 samples, the smallest covering 200 km at 2 km")
        self.assertEqual(self.meta["res_m"], 2000)
        self.assertEqual(self.meta["size_m"], 128 * 2000)
        for f in ("water.png", "biome.png", "preview.png"):
            self.assertEqual(Image.open(self.d / f).size, (self.n, self.n), f)
        self.assertEqual(self.meta["sea_level_m"], 0)
        self.assertIn("world_build", self.meta["sources"])

    def test_samples_are_res_apart_on_the_ground(self):
        lat, lon = EX.grid_latlon(self.lat, self.lon, self.n, 2000.0, 12742.0)
        def gc(i1, j1, i2, j2):
            p1, p2 = np.radians([lat[i1, j1], lon[i1, j1]]), np.radians([lat[i2, j2], lon[i2, j2]])
            return 12742e3 * 2 * math.asin(math.sqrt(math.sin((p2[0] - p1[0]) / 2) ** 2 + math.cos(p1[0]) * math.cos(p2[0])
                                                      * math.sin((p2[1] - p1[1]) / 2) ** 2))
        c = self.n // 2
        for a, b in (((c, c), (c, c + 1)), ((0, 0), (0, 1)), ((0, 0), (1, 0)), ((self.n - 1, 5), (self.n - 2, 5))):
            self.assertAlmostEqual(gc(*a, *b), 2000.0, delta=4.0)
        self.assertAlmostEqual(lat[c, c], self.lat, places=9)
        self.assertGreater(lat[0, c], lat[-1, c], "row 0 is the north edge")

    def test_heights_are_the_viewers_terrain(self):
        src = T.TileSource(serve.World(self.root, 2), seed=int(C.load(self.root)[0]["build"]["seed"]),
                           cache_dir=Path(tempfile.mkdtemp()), regions_dir=Path(tempfile.mkdtemp()))
        lat, lon = EX.grid_latlon(self.lat, self.lon, self.n, 2000.0, 12742.0)
        k = np.s_[::8, ::8]
        ref = src.z_at(lat[k].ravel(), lon[k].ravel(), self.meta["tile_px_m"] / 1000.0).reshape(lat[k].shape)
        self.assertGreater(np.corrcoef(ref.ravel(), self.h[k].ravel())[0, 1], 0.99)
        self.assertLess(np.median(np.abs(ref - self.h[k])), 0.1 * ref.std() + 1.0)

    def test_16_bit_heights_decode_to_the_float_heights(self):
        lo, hi = self.meta["height"]["min"], self.meta["height"]["max"]
        v = np.asarray(Image.open(self.d / "height.png"), dtype=np.float64)
        raw = np.fromfile(self.d / "height.r16", dtype="<u2").reshape(self.n, self.n)
        np.testing.assert_array_equal(v, raw)
        np.testing.assert_allclose(lo + v / 65535 * (hi - lo), self.h, atol=(hi - lo) / 65535 + 1e-3)

    def test_water_classes(self):
        w = np.asarray(Image.open(self.d / "water.png"))
        self.assertTrue(set(np.unique(w)) <= {0, 1, 2, 3})
        self.assertTrue((w == 1).any() and (w == 0).any(), "the square holds sea and land")
        self.assertLess(np.median(self.h[w == 1]), 0.0, "the sea lies below sea level")
        self.assertEqual(json.loads((self.d / "legend.json").read_text())["water"], ["land", "sea", "lake", "river"])

    def test_rivers_and_pins_in_local_metres(self):
        rivers = json.loads((self.d / "rivers.json").read_text())
        self.assertTrue(rivers, "the river mouth's river")
        half = self.meta["size_m"] / 2
        for r in rivers:
            self.assertGreater(r["width_m"], 0)
            p = np.asarray(r["points"])
            self.assertTrue(np.all(np.abs(p) <= half + 1e-6))
        pins = json.loads((self.d / "pins.json").read_text())
        self.assertEqual([p["name"] for p in pins], ["Here"])
        self.assertLess(abs(pins[0]["x"]) + abs(pins[0]["y"]), 1.0)
        self.assertEqual(pins[0]["lore"], "somewhere")

    def test_preview_shades_the_relief_so_gentle_land_still_shows(self):
        pv = self.meta["preview"]
        self.assertGreaterEqual(pv["exaggeration"], 1.0)
        rgb = np.asarray(Image.open(self.d / "preview.png"), dtype=np.float64)
        w = np.asarray(Image.open(self.d / "water.png"))
        lum = rgb.mean(axis=2)[w == 0]
        self.assertGreater(lum.std(), 2.0, "not a flat colour field")

    def test_deterministic_in_parallel_too(self):
        d2 = EX.export(self.root, 2, self.lat, self.lon, size_km=200, res_m=2000, out_dir=self.out, name="t2",
                       pins_path=self.pins, regions_dir=Path(tempfile.mkdtemp()), workers=2)
        for f in ("height.f32", "water.png", "biome.png", "preview.png", "rivers.json", "pins.json"):
            self.assertEqual((self.d / f).read_bytes(), (d2 / f).read_bytes(), f)

    def test_exports_into_a_new_folder_leave_no_work_files(self):
        out = Path(tempfile.mkdtemp()) / "new" / "exports"
        d = EX.export(self.root, 2, self.lat, self.lon, size_km=50, res_m=2000, out_dir=out, name="n",
                      regions_dir=Path(tempfile.mkdtemp()))
        self.assertEqual([p.name for p in out.iterdir()], ["n"])
        self.assertTrue((d / "meta.json").exists())

    def test_a_region_change_during_an_export_stops_it_cleanly(self):
        import refine as RF
        import shutil
        from tests.test_refine import square
        regions = Path(tempfile.mkdtemp())
        reg = regions / "r.json"
        reg.write_text(json.dumps({"regions": [square(self.lat, self.lon, 3)]}))
        built = Path(tempfile.mkdtemp())
        RF.build_areas(self.root, 2, reg, log=lambda s: None, regions_root=built)
        src = T.TileSource(serve.World(self.root, 2), seed=int(C.load(self.root)[0]["build"]["seed"]),
                           cache_dir=Path(tempfile.mkdtemp()), regions_dir=regions / "live")
        real, calls = src.export_tile, []

        def tile(*a):
            calls.append(a)
            if len(calls) == 2:                                  # a region build finishes mid-export
                shutil.copytree(built, regions / "live")
                src.reload_regions()
            return real(*a)
        src.export_tile = tile
        out = Path(tempfile.mkdtemp())
        with self.assertRaises(RuntimeError):
            EX.export(self.root, 2, self.lat, self.lon, size_km=200, res_m=2000, out_dir=out, name="mixed", src=src)
        self.assertEqual(list(out.iterdir()), [], "nothing half-mixed written")

    def test_an_existing_export_is_not_overwritten(self):
        with self.assertRaises(FileExistsError):
            EX.export(self.root, 2, self.lat, self.lon, size_km=200, res_m=2000, out_dir=self.out, name="t",
                      pins_path=self.pins, regions_dir=Path(tempfile.mkdtemp()))


class CliTest(unittest.TestCase):
    def test_export_command(self):
        import importlib.util
        spec = importlib.util.spec_from_file_location("mapview_cli", Path(EX.__file__).parent / "mapview.py")
        cli = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(cli)
        lat, lon = _spot()
        out = Path(tempfile.mkdtemp())
        code = cli.main(["export", "--res", "2", "--lat", str(lat), "--lon", str(lon), "--size-km", "100",
                            "--res-m", "2000", "--name", "cli", "--out", str(out), "--workers", "0"], root=built_world())
        self.assertEqual(code, 0)
        self.assertEqual(json.loads((out / "cli" / "meta.json").read_text())["samples"], 65)

    def test_exact_grid_512_m_at_1_m(self):
        import importlib.util
        spec = importlib.util.spec_from_file_location("mapview_cli", Path(EX.__file__).parent / "mapview.py")
        cli = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(cli)
        lat, lon = _spot()
        out = Path(tempfile.mkdtemp())
        self.assertEqual(cli.main(["export", "--res", "2", "--lat", str(lat), "--lon", str(lon), "--size-km", "0.512", "--res-m", "1", "--pixels", "513",
                                      "--name", "c", "--out", str(out), "--workers", "0"], root=built_world()), 0)
        meta = json.loads((out / "c" / "meta.json").read_text())
        self.assertEqual((meta["samples"], meta["res_m"], meta["size_m"]), (513, 1, 512))
        h = np.fromfile(out / "c" / "height.f32", dtype="<f4").reshape(513, 513)
        self.assertAlmostEqual(meta["center_height_m"], float(h[256, 256]), places=3)

class ExportForkTest(unittest.TestCase):
    def test_region_trees_are_built_before_export_workers_fork(self):
        import refine
        from unittest import mock
        order = []

        class FakePool:                                      # stands in for the forked pool: records when it forks
            def __init__(self, n):
                order.append("fork")

            def __enter__(self):
                return self

            def __exit__(self, *a):
                return False

            def imap(self, f, jobs, chunksize=1):
                return map(f, jobs)
        warm = refine.RegionSet.warm
        a = np.load(built_world() / "out" / "r2" / "cells.npz")
        i = int(np.argmax(a["z_surface_m"]))
        with mock.patch.object(refine.RegionSet, "warm", autospec=True,
                               side_effect=lambda rs: (order.append("warm"), warm(rs))[1]), \
             mock.patch.object(EX.multiprocessing, "get_context", return_value=mock.Mock(Pool=FakePool)):
            EX.export(built_world(), 2, float(a["g_lat"][i]), float(a["g_lon"][i]), size_km=50, res_m=2000,
                      out_dir=Path(tempfile.mkdtemp()), name="w", workers=2)
        self.assertIn("warm", order)
        self.assertLess(order.index("warm"), order.index("fork"))


if __name__ == "__main__":
    unittest.main()


class EraExportTest(unittest.TestCase):
    def test_export_of_an_era_with_sea_and_zone_metadata(self):
        import json
        import export
        from tests.test_refine import era_world
        root, _, cut, _ = era_world()
        out = Path(tempfile.mkdtemp())
        a = export.export(root, 2, cut[0], cut[1], size_km=20.0, res_m=500.0, out_dir=out, name="after-x")
        b = export.export(root, 2, cut[0], cut[1], size_km=20.0, res_m=500.0, out_dir=out, name="before-x",
                          era="before")
        ma, mb = (json.loads((d / "meta.json").read_text()) for d in (a, b))
        self.assertEqual((ma["era"]["name"], mb["era"]["name"]), ("after", "before"))
        self.assertLess(ma["center_height_m"], mb["center_height_m"], "the cut lowered the ground")
        for m in (ma, mb):
            self.assertEqual(set(m["water"]), {"sea_fraction", "max_depth_m", "center_depth_m"})
            self.assertEqual(set(m["zones"]), {"gravity_g", "o2_fraction", "po2_bar", "pressure_bar", "fire_reactivity"})
            self.assertIsInstance(m["vents"], list)
            self.assertEqual(m["format"], "worldmap-terrain-export")
        self.assertAlmostEqual(ma["zones"]["fire_reactivity"], 0.4, places=3)
        with self.assertRaises(SystemExit):
            export.export(root, 2, cut[0], cut[1], size_km=20.0, res_m=500.0, out_dir=out, name="n", era="nope")
        from unittest import mock
        import serve
        real, made = serve.World, []
        with mock.patch.object(serve, "World", side_effect=lambda *a, **k: made.append(k.get("era")) or real(*a, **k)):
            export.export(root, 2, cut[0], cut[1], size_km=20.0, res_m=500.0, out_dir=out, name="one", era="before")
        self.assertEqual([e["name"] for e in made], ["before"], "only the chosen era's world is loaded")