import unittest

from autark import JsonReader, OnceFrame, P


class TestStringWork(unittest.TestCase):
    def test_group_by_count_on_string_column(self) -> None:
        rows = [
            {"city": "A", "v": 1.0},
            {"city": "B", "v": 2.0},
            {"city": "A", "v": 3.0},
            {"city": "C", "v": 4.0},
            {"city": "B", "v": 5.0},
        ]

        frame = OnceFrame(JsonReader(rows))
        p = P()

        p.columns("city").group_reduce("count").alias("city_count")
        realized = frame.realize(p)
        print(realized)
        table = realized["city_count"]

        self.assertIn("city", table)
        self.assertIn("city_count", table)

        cities = table["city"]
        counts = table["city_count"]
        self.assertEqual(len(cities), len(counts))

        observed = dict(zip(cities, map(int, counts)))
        self.assertEqual(observed, {"A": 2, "B": 2, "C": 1})

    def test_string_equality_filter(self) -> None:
        rows = [
            {"city": "A", "v": 1.0},
            {"city": "B", "v": 2.0},
            {"city": "A", "v": 3.0},
            {"city": "C", "v": 4.0},
            {"city": "B", "v": 5.0},
        ]

        frame = OnceFrame(JsonReader(rows))
        p = P()

        city = p.columns("city")
        value = p.columns("v")
        (city == "A").alias("is_a")
        (city != "A").alias("is_not_a")
        (city == "Z").alias("is_z")
        (value == "A").alias("num_eq_str")

        realized = frame.realize(p)
        self.assertEqual(realized["is_a"]["is_a"], [True, False, True, False, False])
        self.assertEqual(realized["is_not_a"]["is_not_a"], [False, True, False, True, True])
        self.assertEqual(realized["is_z"]["is_z"], [False, False, False, False, False])
        self.assertEqual(realized["num_eq_str"]["num_eq_str"], [False, False, False, False, False])


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