# map/tests/test_zones.py
import tempfile
import unittest
from pathlib import Path

import numpy as np
from PIL import Image

from mapgen import sketch as SK, zones as ZN
from mapgen.sphere import east_north, great_circle_point, latlon_to_xyz
from tests.helpers import make_ctx

R = 12742.0
R1 = {"name": "R1", "field": "o2", "center": [-5.6, -118.8], "radius_km": 3200.0, "v": 1.0}
L1 = {"name": "L1", "field": "gravity", "center": [-26.3, -133.7], "radius_km": 2200.0, "v": -0.6}


def ray(zone, az_deg, step_km=5.0):
    p0 = latlon_to_xyz(*zone["center"])
    e, n = east_north(p0[None])
    a = np.radians(az_deg)
    t = np.sin(a) * e[0] + np.cos(a) * n[0]
    s = np.arange(0.0, 1.4 * zone["radius_km"], step_km)
    return s, np.array([great_circle_point(p0, t, si, R) for si in s])


class ZoneProfileTest(unittest.TestCase):
    def test_peak_at_the_centre_and_nothing_beyond_the_warped_radius(self):
        for z in (R1, L1):
            s, pts = ray(z, 40.0)
            f = ZN.contribution(pts, z, 1296, R)
            self.assertAlmostEqual(f[0], z["v"])
            self.assertTrue(np.all(f[s > z["radius_km"] / (1 - ZN.WARP) + 1] == 0.0))

    def test_steepest_change_is_hard_to_notice_over_a_days_walk(self):
        """≤ 0.23 O₂ points and ≤ 0.014 g per 30 km, for the strongest O₂ and low-gravity zones."""
        for z, per_unit, limit in ((R1, 0.21 * 0.5 * 100.0, 0.23), (L1, 1.05 * 0.7, 0.014)):
            worst = 0.0
            for az in range(0, 360, 5):
                _, pts = ray(z, az)
                f = ZN.contribution(pts, z, 1296, R) * per_unit
                worst = max(worst, float(np.max(np.abs(np.diff(f)))) / 5.0 * 30.0)
            self.assertLessEqual(worst, limit, z["name"])

    def test_outline_is_not_a_circle(self):
        at = [ZN.contribution(ray(R1, az)[1][[400]], R1, 1296, R)[0] for az in range(0, 360, 30)]   # 2,000 km out
        self.assertGreater(max(at) - min(at), 0.02)


class SketchZonesTest(unittest.TestCase):
    def test_config_zones_add_to_painted_masks_and_clip(self):
        root = Path(tempfile.mkdtemp())
        (root / "sketch").mkdir()
        (root / "masks").mkdir()
        for n in SK.SKETCH:
            Image.fromarray(np.zeros((10, 20), np.uint8), "L").save(root / "sketch" / f"{n}.png")
        Image.fromarray(np.full((10, 20), 191, np.uint8), "L").save(root / "masks" / "o2_zones.png")   # ≈ +0.5
        zones = [{**R1, "center": [0.0, 0.0]}, {**L1, "center": [0.0, 0.0]}]
        ctx = make_ctx(2, tect={"plate": [], "zone": zones}, root=root)
        out = SK.run(ctx)
        g = ctx.grid
        c, far = g.cell_index(0.0, 0.0), g.cell_index(0.0, 180.0)
        self.assertEqual(out["m_o2_zones"][c], 1.0, "mask + zone, clipped")
        self.assertAlmostEqual(out["m_o2_zones"][far], 0.498, places=2)
        self.assertAlmostEqual(out["m_gravity_zones"][c], -0.6, delta=0.05)
        self.assertEqual(out["m_gravity_zones"][far], 0.0)