from __future__ import annotations

import sys
import unittest
from pathlib import Path
from tempfile import TemporaryDirectory

import geopandas as gpd
import numpy as np
import rasterio
from rasterio.transform import from_origin
from shapely.geometry import MultiPolygon, Polygon, box

SCRIPTS_DIR = Path(__file__).resolve().parents[1] / "scripts"
sys.path.insert(0, str(SCRIPTS_DIR))

from build_county_climate_data import _touched_raster_values, _zonal_majority_class  # noqa: E402
from common.county_zonal_stats import geometry_window, split_at_antimeridian  # noqa: E402
from common.koppen_legend import DEFAULT_KOPPEN_CODE_MAP  # noqa: E402


def islands_on_both_sides() -> MultiPolygon:
    """A county shaped like Aleutians West: islands at +178..+180 and -180..-179."""
    return MultiPolygon([box(178.2, 51.2, 179.8, 51.8), box(-179.8, 51.2, -179.2, 51.8)])


class SplitAtAntimeridianTests(unittest.TestCase):
    def test_geometry_on_one_side_is_unchanged(self) -> None:
        geometry = box(-100.0, 30.0, -90.0, 40.0)

        pieces = split_at_antimeridian(geometry)

        self.assertEqual(len(pieces), 1)
        self.assertTrue(pieces[0].equals(geometry))

    def test_islands_are_grouped_by_side(self) -> None:
        geometry = MultiPolygon(
            [box(172.0, 51.0, 173.0, 52.0), box(179.0, 51.0, 179.5, 52.0), box(-179.0, 51.0, -178.0, 52.0)]
        )

        pieces = split_at_antimeridian(geometry)

        self.assertEqual(len(pieces), 2)
        bounds = sorted(piece.bounds for piece in pieces)
        self.assertEqual(bounds[0], (-179.0, 51.0, -178.0, 52.0))
        self.assertEqual(bounds[1], (172.0, 51.0, 179.5, 52.0))
        self.assertAlmostEqual(sum(piece.area for piece in pieces), geometry.area)

    def test_single_outline_crossing_the_line_is_cut(self) -> None:
        # Encoded in -180..180 coordinates, this 2-degree-wide ring looks 358 degrees wide.
        crossing = Polygon([(179.0, 50.0), (-179.0, 50.0), (-179.0, 51.0), (179.0, 51.0)])

        pieces = split_at_antimeridian(crossing)

        bounds = sorted(piece.bounds for piece in pieces)
        self.assertEqual(bounds, [(-180.0, 50.0, -179.0, 51.0), (179.0, 50.0, 180.0, 51.0)])
        self.assertAlmostEqual(sum(piece.area for piece in pieces), 2.0)


class WindowedKoppenTests(unittest.TestCase):
    def setUp(self) -> None:
        self._temp_dir = TemporaryDirectory()
        self.raster_path = Path(self._temp_dir.name) / "koppen.tif"
        # Global 1-degree raster; row 38 covers 51..52 N.
        data = np.zeros((180, 360), dtype=np.uint8)
        data[38, 358] = 29  # ET at 178..179 E
        data[38, 359] = 29  # ET at 179..180 E
        data[38, 0] = 15  # Cfb at 180..179 W
        with rasterio.open(
            self.raster_path,
            "w",
            driver="GTiff",
            height=180,
            width=360,
            count=1,
            dtype="uint8",
            crs="EPSG:4326",
            transform=from_origin(-180.0, 90.0, 1.0, 1.0),
            nodata=0,
        ) as destination:
            destination.write(data, 1)

    def tearDown(self) -> None:
        self._temp_dir.cleanup()

    def test_each_side_is_read_from_a_small_window(self) -> None:
        with rasterio.open(self.raster_path) as source:
            windows = [geometry_window(source, piece) for piece in split_at_antimeridian(islands_on_both_sides())]
            values = _touched_raster_values(source, islands_on_both_sides(), split_antimeridian=True)

        self.assertEqual(len(windows), 2)
        self.assertTrue(all(window.width <= 5 for window in windows))
        self.assertEqual(sorted(values[values != 0].tolist()), [15, 29, 29])

    def test_majority_class_counts_islands_on_both_sides(self) -> None:
        counties = gpd.GeoDataFrame(geometry=[islands_on_both_sides(), box(-100.0, 30.0, -99.0, 31.0)], crs="EPSG:4326")

        classes = _zonal_majority_class(self.raster_path, counties, DEFAULT_KOPPEN_CODE_MAP)

        self.assertEqual(classes, ["ET", "Cfa"])


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