geointel/backend/app/services/raster_partition_analysis_service.py

from __future__ import annotations

import math
from contextlib import ExitStack
from dataclasses import dataclass
from pathlib import Path
from typing import Any
from uuid import UUID

from pyproj import Transformer
from shapely.ops import transform as shapely_transform

from app.core.errors import AppError
from app.services.raster_cell_selection import select_cells
from app.models import Dataset


@dataclass(frozen=True)
class RasterPartitionSelection:
    datasets: list[Dataset]
    values: Any
    selected_cells: Any
    resolution_x: float
    resolution_y: float
    # Set when the selection is finer than one source cell and the analysis was
    # widened to the touched cells, so it covers more ground than was drawn.
    cell_selection_warning: str | None = None


class RasterPartitionAnalysisService:
    MAX_PARTITIONS = 64

    @staticmethod
    def _bbox_intersects(dataset: Dataset, bbox: tuple[float, float, float, float]) -> bool:
        source_bbox = (dataset.source_metadata or {}).get("bbox_epsg4326")
        if not isinstance(source_bbox, list) or len(source_bbox) != 4:
            return True
        try:
            min_x, min_y, max_x, max_y = (float(value) for value in source_bbox)
        except (TypeError, ValueError):
            return True
        return not (
            max_x <= bbox[0]
            or min_x >= bbox[2]
            or max_y <= bbox[1]
            or min_y >= bbox[3]
        )

    @staticmethod
    def _candidate_datasets(
        db,
        project_id: UUID,
        *,
        source_name: str,
        product_key: str,
        bbox: tuple[float, float, float, float],
        dataset_ids: list[UUID] | None = None,
    ) -> list[Dataset]:
        query = db.query(Dataset).filter(
            Dataset.project_id == project_id,
            Dataset.source_name == source_name,
            Dataset.dataset_type == "raster",
            Dataset.status == "ready",
        )
        if dataset_ids is not None:
            query = query.filter(Dataset.id.in_(dataset_ids))
        rows = query.all()
        candidates = [
            dataset
            for dataset in rows
            if str((dataset.source_metadata or {}).get("product_key") or "") == product_key
            and dataset.storage_path
            and Path(dataset.storage_path).is_file()
            and RasterPartitionAnalysisService._bbox_intersects(dataset, bbox)
        ]
        candidates.sort(key=lambda dataset: (str(dataset.area_id or ""), str(dataset.id)))
        if not candidates:
            raise AppError(
                code="RASTER_PARTITIONS_NOT_FOUND",
                message="No persisted raster partitions cover this selection",
                details={"source_name": source_name, "product_key": product_key},
                status_code=404,
            )
        if dataset_ids is not None and {dataset.id for dataset in candidates} != set(dataset_ids):
            raise AppError(
                code="RASTER_PARTITION_SOURCE_MISMATCH",
                message="Every requested raster partition must match the governed source product and selection",
                details={"requested_count": len(dataset_ids), "eligible_count": len(candidates)},
                status_code=409,
            )
        if len(candidates) > RasterPartitionAnalysisService.MAX_PARTITIONS:
            raise AppError(
                code="RASTER_PARTITION_LIMIT_EXCEEDED",
                message="The selection intersects too many raster partitions",
                details={
                    "partition_count": len(candidates),
                    "max_partitions": RasterPartitionAnalysisService.MAX_PARTITIONS,
                },
                status_code=422,
            )
        return candidates

    @staticmethod
    def select(
        db,
        project_id: UUID,
        *,
        source_name: str,
        product_key: str,
        selection_geometry_4326,
        nodata: float,
        max_pixels: int,
        dataset_ids: list[UUID] | None = None,
    ) -> RasterPartitionSelection:
        try:
            import numpy as np
            import rasterio
            from rasterio.merge import merge
        except ImportError as exc:
            raise AppError(
                code="RASTER_PROCESSING_UNAVAILABLE",
                message="Rasterio and numpy are required for partitioned raster analysis",
                status_code=503,
            ) from exc

        bbox = tuple(float(value) for value in selection_geometry_4326.bounds)
        datasets = RasterPartitionAnalysisService._candidate_datasets(
            db,
            project_id,
            source_name=source_name,
            product_key=product_key,
            bbox=bbox,
            dataset_ids=dataset_ids,
        )
        transformer = Transformer.from_crs("EPSG:4326", "EPSG:31370", always_xy=True)
        selection_metric = shapely_transform(transformer.transform, selection_geometry_4326)
        min_x, min_y, max_x, max_y = selection_metric.bounds

        try:
            with ExitStack() as stack:
                sources = [stack.enter_context(rasterio.open(dataset.storage_path)) for dataset in datasets]
                invalid_sources = [
                    index
                    for index, source in enumerate(sources)
                    if source.crs is None or source.crs.to_epsg() != 31370 or source.count != 1
                ]
                if invalid_sources:
                    raise AppError(
                        code="RASTER_PARTITION_MISMATCH",
                        message="Raster partitions do not share the governed CRS and band layout",
                        details={"invalid_partition_indexes": invalid_sources},
                        status_code=409,
                    )
                target_resolution = max(abs(float(sources[0].res[0])), abs(float(sources[0].res[1])))
                invalid_resolutions = [
                    {
                        "partition_index": index,
                        "resolution": [abs(float(source.res[0])), abs(float(source.res[1]))],
                    }
                    for index, source in enumerate(sources)
                    if not all(
                        math.isclose(abs(float(value)), target_resolution, rel_tol=0.001, abs_tol=0.01)
                        for value in source.res
                    )
                ]
                if invalid_resolutions:
                    raise AppError(
                        code="RASTER_PARTITION_MISMATCH",
                        message="Raster partitions do not share one analysis resolution",
                        details={"invalid_resolutions": invalid_resolutions},
                        status_code=409,
                    )
                width = max(1, math.ceil((max_x - min_x) / target_resolution))
                height = max(1, math.ceil((max_y - min_y) / target_resolution))
                if width * height > max_pixels:
                    raise AppError(
                        code="RASTER_PARTITION_SELECTION_TOO_LARGE",
                        message="Select a smaller rectangle for regional raster analysis",
                        details={"pixel_count": width * height, "max_pixels": max_pixels},
                        status_code=422,
                    )
                mosaic, transform = merge(
                    sources,
                    bounds=(min_x, min_y, max_x, max_y),
                    res=(target_resolution, target_resolution),
                    nodata=nodata,
                    dtype="float32",
                )
                values = np.asarray(mosaic[0], dtype="float64")
                cell_selection = select_cells(
                    selection_metric,
                    out_shape=values.shape,
                    transform=transform,
                    cell_area_m2=target_resolution * target_resolution,
                )
                selected_cells = cell_selection.mask
                return RasterPartitionSelection(
                    datasets=datasets,
                    values=values,
                    selected_cells=selected_cells,
                    resolution_x=target_resolution,
                    resolution_y=target_resolution,
                    cell_selection_warning=cell_selection.warning,
                )
        except AppError:
            raise
        except Exception as exc:
            raise AppError(
                code="RASTER_PARTITION_ANALYSIS_FAILED",
                message="Persisted raster partitions could not be assembled for this selection",
                details={"reason": str(exc)},
                status_code=500,
            ) from exc