import unittest

import numpy as np

from tests.helpers import fields, line_world
from worldhistory.predicates import check, depth, evaluate, parse
from worldhistory.regions import Regions, region_mask


class PredicateTest(unittest.TestCase):
    def setUp(self):
        # cells: 0 forest land, 1 desert land by a river, 2 mountain land, 3 shelf sea, 4 deep sea with a vent
        self.w = line_world(5, ocean=[False, False, False, True, True], holdridge=[20, 23, 12, 0, 0],
                            landform=[1, 1, 3, 0, 0], river=[False, True, False, False, False],
                            strahler=[0, 4, 0, 0, 0], elevation_m=[100, 50, 2000, -150, -3000],
                            vent_potential=[0, 0, 0, 0, 0.8], T_mean=[10, 30, 0, 12, 5], lithology=[1, 3, 4, 0, 0])

    def test_parse(self):
        self.assertEqual(parse("land"), ("land", []))
        self.assertEqual(parse("warm:5:20"), ("warm", [5.0, 20.0]))
        self.assertEqual(parse("lithology:3,4"), ("lithology", [[3, 4]]))
        self.assertEqual(parse("above:o2_fraction:0.3"), ("above", ["o2_fraction", 0.3]))

    def test_basic(self):
        e = lambda p: evaluate(p, self.w).tolist()
        self.assertEqual(e("land"), [1, 1, 1, 0, 0])
        self.assertEqual(e("forest"), [1, 0, 0, 0, 0])
        self.assertEqual(e("desert"), [0, 1, 0, 0, 0])
        self.assertEqual(e("mountain"), [0, 0, 1, 0, 0])
        self.assertEqual(e("coast"), [0, 0, 1, 0, 0])
        self.assertEqual(e("shelf"), [0, 0, 0, 1, 0])
        self.assertEqual(e("deep_sea"), [0, 0, 0, 0, 1])
        self.assertAlmostEqual(e("water")[1], 1.0)              # 0.4 + 0.15*4 = 1.0
        self.assertAlmostEqual(e("vent")[4], 0.8)
        self.assertEqual(e("lithology:3,4"), [0, 1, 1, 0, 0])

    def test_ranges_and_products(self):
        self.assertEqual(evaluate("warm:5:25", self.w).tolist(), [0.25, 1.0, 0.0, 0.35, 0.0])
        self.assertEqual(evaluate(["land", "above:T_mean:20"], self.w).tolist(), [0, 1, 0, 0, 0])
        np.testing.assert_array_equal(depth(self.w.fields), [0, 0, 0, 150, 3000])

    def test_check(self):
        check(["land", "warm:1:2", "lithology:1"])
        for bad in ("nosuch", "warm:1", "above:nofield:1"):
            with self.assertRaises(ValueError):
                check(bad, field_names=set(fields(1)))


class RegionTest(unittest.TestCase):
    def test_kinds(self):
        w = line_world(4, ocean=[False, False, False, True], plate=[1, 1, 2, 2], gravity_g=1.0)
        w.eras["late"] = fields(4, ocean=[False, True, False, True], gravity_g=[1.0, 0.35, 0.35, 0.35])
        self.assertEqual(region_mask(w, {"kind": "all"}).sum(), 4)
        self.assertEqual(region_mask(w, {"kind": "plate", "ids": [2]}).tolist(), [False, False, True, True])
        self.assertEqual(region_mask(w, {"kind": "changed", "era": "late", "fields": ["gravity_g"]}).tolist(),
                         [False, True, True, True])
        self.assertEqual(region_mask(w, {"kind": "changed", "era": "late", "fields": ["gravity_g"], "land": "base"})
                         .tolist(), [False, True, True, False])
        box = region_mask(w, {"kind": "box", "lat": [-1, 1], "lon": [0.5, 10]})
        self.assertEqual(box.tolist(), [False, True, True, True])
        self.assertEqual(region_mask(w, {"kind": "predicate", "p": "land"}).tolist(), [True, True, True, False])

    def test_box_with_predicate_filter(self):
        w = line_world(4, ocean=[False, False, True, True])
        m = region_mask(w, {"kind": "box", "lat": [-1, 1], "lon": [0.5, 10], "p": "land"})
        self.assertEqual(m.tolist(), [False, True, False, False])

    def test_exclude_other_region(self):
        w = line_world(4)
        r = Regions(w, {"east": {"kind": "box", "lat": [-1, 1], "lon": [0.5, 10]},
                        "land_not_east": {"kind": "predicate", "p": "land", "exclude": "east"}})
        self.assertEqual(r("land_not_east").tolist(), [True, False, False, False])

    def test_registry(self):
        w = line_world(3)
        r = Regions(w, {"west": {"kind": "box", "lat": [-1, 1], "lon": [-1, 0.5]}})
        self.assertEqual(r("west").tolist(), [True, False, False])
        self.assertTrue(r("all").all())
        with self.assertRaises(KeyError):
            r("nowhere")