"""
Polygon to Raster Window Preparer.

This module provides a window preparer that converts polygon/multipolygon annotations
to raster labels. It can create either:
- One window per annotation task (default), using the task geometry as the window boundary
- Multiple grid windows per task (optional), with configurable pixel dimensions

Supports optional UTM CRS validation and two strategies for handling tasks that don't
divide evenly into grid windows: center crop or pad with nodata values.
"""

import logging
import math
from datetime import datetime
from typing import assert_never

import numpy as np
from pydantic import BaseModel
from rslearn.utils import Projection, STGeometry
from shapely.geometry import MultiPolygon, Polygon, box
from shapely.geometry.base import BaseGeometry

from olmoearth_run.runner.models.training.labeled_data import (
    AnnotationTask,
    LabeledWindow,
    RasterLabel,
)
from olmoearth_run.runner.tools.labeled_window_preparers.geometry_utils import (
    InconsistentUtmZoneError,
    BoxBounds,
    compute_window_bounds,
    create_raster_labels_from_annotations,
    get_utm_crs_for_geometry,
    project_geometry_to_crs,
)
from olmoearth_run.runner.tools.labeled_window_preparers.labeled_window_preparer import (
    RasterLabelsWindowPreparer,
)
from olmoearth_run.runner.tools.labeled_window_preparers.rasterization_utils import (
    DEFAULT_NODATA_VALUE,
)
from olmoearth_shared.models.olmoearth_config.labeled_data_prep.window_preparers import (
    GriddingStrategy,
)

logger = logging.getLogger(__name__)


class GridConfig(BaseModel, frozen=True):
    """Configuration for grid-based window generation.

    Attributes:
        window_size_px: Size of each square window in pixels.
        gridding_strategy: Strategy for handling non-integer divisions.
    """

    window_size_px: int
    gridding_strategy: GriddingStrategy = GriddingStrategy.CENTER_CROP


class GridCell(BaseModel, frozen=True):
    """Specification for slicing a cell from a full task raster.

    Attributes:
        pixel_bounds: Cell bounds in raster coordinates (for array slicing)
        world_bounds: Cell bounds in projected coordinates (for geometry)
    """

    pixel_bounds: BoxBounds
    world_bounds: BoxBounds


def compute_grid_cells(
    full_bounds: BoxBounds,
    cell_size_px: int,
    gridding_strategy: GriddingStrategy,
) -> list[GridCell]:
    """Compute specs for dividing a raster into grid cells.

    Args:
        full_bounds: Bounds of the full task raster in projected pixel coordinates
        cell_size_px: Size of each square cell in pixels
        gridding_strategy: Strategy for handling non-integer divisions

    Returns:
        List of GridCell objects for extracting each cell.
        Returns empty list for CENTER_CROP if raster is smaller than one cell.
    """
    raster_height = full_bounds.height
    raster_width = full_bounds.width

    match gridding_strategy:
        case GriddingStrategy.CENTER_CROP:
            n_cols = raster_width // cell_size_px
            n_rows = raster_height // cell_size_px
            if n_cols == 0 or n_rows == 0:
                return []
            # Center the grid within the raster
            x_offset = (raster_width - n_cols * cell_size_px) // 2
            y_offset = (raster_height - n_rows * cell_size_px) // 2

        case GriddingStrategy.PAD_TO_FILL:
            n_cols = max(1, math.ceil(raster_width / cell_size_px))
            n_rows = max(1, math.ceil(raster_height / cell_size_px))
            x_offset = 0
            y_offset = 0

        case _:
            assert_never(gridding_strategy)

    slices: list[GridCell] = []
    for row in range(n_rows):
        for col in range(n_cols):
            cell_y = y_offset + row * cell_size_px
            cell_x = x_offset + col * cell_size_px

            slices.append(GridCell(
                pixel_bounds=BoxBounds(
                    minx=cell_x,
                    miny=cell_y,
                    maxx=cell_x + cell_size_px,
                    maxy=cell_y + cell_size_px,
                ),
                world_bounds=BoxBounds(
                    minx=full_bounds.minx + cell_x,
                    miny=full_bounds.miny + cell_y,
                    maxx=full_bounds.minx + cell_x + cell_size_px,
                    maxy=full_bounds.miny + cell_y + cell_size_px,
                ),
            ))

    return slices


def extract_cell_rasters(
    full_raster_labels: list[RasterLabel],
    grid_slice: GridCell,
    nodata_value: int,
) -> list[RasterLabel]:
    """Extract a cell's raster data from the full task raster, with padding if needed."""
    cell_labels: list[RasterLabel] = []
    pixel_bounds = grid_slice.pixel_bounds
    cell_size = pixel_bounds.width  # assumes square cells

    for label in full_raster_labels:
        # Slice the portion from full raster (numpy handles out-of-bounds gracefully)
        sliced = label.value[pixel_bounds.miny:pixel_bounds.maxy, pixel_bounds.minx:pixel_bounds.maxx]

        # Pad if slice is smaller than cell_size (edge cell)
        if sliced.shape[0] < cell_size or sliced.shape[1] < cell_size:
            padded = np.full(
                (cell_size, cell_size),
                nodata_value,
                dtype=label.value.dtype,
            )
            padded[:sliced.shape[0], :sliced.shape[1]] = sliced
            cell_labels.append(RasterLabel(key=label.key, value=padded))
        else:
            cell_labels.append(RasterLabel(key=label.key, value=sliced))

    return cell_labels


def create_labeled_window(
    name: str,
    time_range: tuple[datetime, datetime] | None,
    raster_labels: list[RasterLabel],
    cell_bounds: BoxBounds,
    projection: Projection,
) -> LabeledWindow[list[RasterLabel]]:
    """Create a LabeledWindow from pre-computed raster labels."""
    window_polygon = box(
        cell_bounds.minx, cell_bounds.miny, cell_bounds.maxx, cell_bounds.maxy
    )
    window_st_geometry = STGeometry(projection, window_polygon, time_range)

    return LabeledWindow(
        name=name,
        st_geometry=window_st_geometry,
        labels=raster_labels,
    )


class PolygonToRasterWindowPreparer(RasterLabelsWindowPreparer):
    """
    Window preparer that converts polygon/multipolygon annotations to raster labels.

    This preparer can operate in two modes:
    1. Single window mode (default): Creates one window per annotation task, using the
       task geometry as the window boundary.
    2. Grid mode: Divides the task into multiple fixed-size square windows.

    Key characteristics:
    - Labels are uint8 raster arrays
    - Uses UTM projection for consistent resolution
    - Supports optional UTM CRS validation for tasks spanning UTM zones
    - Grid mode supports two strategies for non-integer divisions:
      - CENTER_CROP: Crop to largest centered area fitting whole windows
      - PAD_TO_FILL: Pad edge windows with nodata values
    """

    def __init__(
        self,
        window_resolution: float = 10.0,
        nodata_value: int = DEFAULT_NODATA_VALUE,
        grid_config: GridConfig | None = None,
        drop_inconsistent_utm: bool = False,
    ):
        self.window_resolution = window_resolution
        self.nodata_value = nodata_value
        self.grid_config = grid_config
        self.drop_inconsistent_utm = drop_inconsistent_utm

    def prepare_labeled_windows(
        self, annotation_task: AnnotationTask
    ) -> list[LabeledWindow[list[RasterLabel]]]:
        """Prepare labeled windows from polygon annotation tasks."""

        if not annotation_task.annotations:
            return []

        # Validate that all annotations have polygon geometries
        for annotation in annotation_task.annotations:
            geom = annotation.st_geometry.shp
            if not isinstance(geom, (Polygon, MultiPolygon)):
                raise ValueError(
                    f"Expected Polygon or MultiPolygon annotation geometry, "
                    f"got {type(geom).__name__} with geom_type {getattr(geom, 'geom_type', 'unknown')}"
                )

        # Extract the task geometry
        task_geom = annotation_task.task_st_geometry.shp
        if not isinstance(
            task_geom, (Polygon, BaseGeometry)
        ) or task_geom.geom_type not in ["Polygon", "MultiPolygon"]:
            raise ValueError(
                f"Expected Polygon or MultiPolygon for task, got {type(task_geom)} with geom_type {getattr(task_geom, 'geom_type', 'unknown')}"
            )

        # Determine UTM CRS
        try:
            utm_crs = get_utm_crs_for_geometry(task_geom, self.drop_inconsistent_utm)
        except InconsistentUtmZoneError:
            logger.warning(
                f"Task {annotation_task.task_id} spans multiple UTM zones, skipping"
            )
            return []

        # Convert to appropriate projection
        projected_geometry = project_geometry_to_crs(
            task_geom, self.window_resolution, utm_crs
        )
        full_bounds = compute_window_bounds(projected_geometry)

        # Rasterize the full task once
        full_raster_labels = create_raster_labels_from_annotations(
            annotations=annotation_task.annotations,
            window_bounds=full_bounds,
            window_resolution=self.window_resolution,
            crs=utm_crs,
            nodata_value=self.nodata_value,
        )

        if not full_raster_labels:
            return []

        # Compute grid slices (single slice covering full task when no grid config)
        if self.grid_config is None:
            grid_slices = [GridCell(
                pixel_bounds=BoxBounds(
                    minx=0, miny=0, maxx=full_bounds.width, maxy=full_bounds.height
                ),
                world_bounds=full_bounds,
            )]
        else:
            grid_slices = compute_grid_cells(
                full_bounds=full_bounds,
                cell_size_px=self.grid_config.window_size_px,
                gridding_strategy=self.grid_config.gridding_strategy,
            )

            if not grid_slices:
                logger.warning(
                    f"Task {annotation_task.task_id} is too small for grid cells "
                    f"with CENTER_CROP strategy (task smaller than "
                    f"{self.grid_config.window_size_px}x{self.grid_config.window_size_px} pixels)"
                )
                return []

        # Extract each cell from the full raster
        labeled_windows: list[LabeledWindow[list[RasterLabel]]] = []
        for idx, grid_slice in enumerate(grid_slices):
            cell_raster_labels = extract_cell_rasters(
                full_raster_labels, grid_slice, self.nodata_value
            )

            # Skip cells where all labels are entirely nodata
            if all(
                (label.value == self.nodata_value).all()
                for label in cell_raster_labels
            ):
                logger.debug(f"Skipping cell {idx} for task {annotation_task.task_id} because it is entirely nodata")
                continue

            labeled_window = create_labeled_window(
                name=f"task_{annotation_task.task_id}_cell_{idx}",
                time_range=annotation_task.task_st_geometry.time_range,
                raster_labels=cell_raster_labels,
                cell_bounds=grid_slice.world_bounds,
                projection=projected_geometry.projection,
            )
            labeled_windows.append(labeled_window)

        return labeled_windows
