geointel/backend/app/services/qa_service.py

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any
from uuid import UUID

from geoalchemy2.shape import to_shape
from shapely.geometry import GeometryCollection
from shapely.geometry.base import BaseGeometry
from shapely.strtree import STRtree
from shapely.ops import unary_union
from shapely.validation import make_valid
from app.core.errors import AppError
from app.models import Area, Dataset
from app.schemas.qa import QaProviderComparisonResult
from app.services.dataset_consumption_gate_service import DatasetConsumptionGate
from app.services.vector_operations_service import VectorOperationsService


@dataclass
class QaMatchEvidence:
    matches: int = 0
    false_positives: int = 0
    false_negatives: int = 0
    match_iou_values: list[float] = field(default_factory=list)
    warnings: list[str] = field(default_factory=list)
    unsupported: bool = False
    match_evidence: list[dict[str, Any]] = field(default_factory=list)
    false_positive_evidence: list[dict[str, Any]] = field(default_factory=list)
    false_negative_evidence: list[dict[str, Any]] = field(default_factory=list)


def _extract_crs_warnings(source_dataset: Dataset, reference_dataset: Dataset) -> list[str]:
    warnings: list[str] = []
    for dataset, label in ((source_dataset, "candidate"), (reference_dataset, "reference")):
        metadata = dataset.metadata_json
        crs_assumed = None
        if isinstance(metadata, dict):
            crs_assumed = metadata.get("crs_assumed")
        if crs_assumed:
            warnings.append(f"CRS assumption is weak for {label} dataset ({dataset.id}); geometry metrics are approximate")
        if dataset.crs is None:
            warnings.append(f"Missing CRS on {label} dataset ({dataset.id})")
    return warnings


class QaService:
    SUPPORTED_GEOMETRY_TYPES = {"Polygon", "MultiPolygon"}

    @staticmethod
    def _load_dataset_payload(db, dataset_id: UUID, *, expected_project_id: UUID | None = None) -> tuple[Dataset, dict[str, Any], list[tuple[dict[str, Any], BaseGeometry]]]:
        dataset = db.get(Dataset, dataset_id)
        if not dataset:
            raise AppError(code="DATASET_NOT_FOUND", message="Dataset not found", status_code=404)
        if expected_project_id is not None and dataset.project_id != expected_project_id:
            raise AppError(code="INVALID_DATASET_SCOPE", message="Dataset does not belong to this project", status_code=400)
        if dataset.dataset_type not in {"vector", "geojson"}:
            raise AppError(code="INVALID_DATASET_TYPE", message="Dataset is not a vector dataset", status_code=400)

        payload, raw_features = VectorOperationsService._load_dataset_payload(dataset)
        geometries = VectorOperationsService._extract_geometries(raw_features)
        return dataset, payload, geometries

    @staticmethod
    def _apply_area_filter(
        geometries: list[tuple[dict[str, Any], BaseGeometry]],
        area_geometry: BaseGeometry,
        *,
        dataset_id: UUID,
    ) -> list[tuple[dict[str, Any], BaseGeometry]]:
        area_geom = area_geometry
        if isinstance(area_geom, GeometryCollection):
            area_geom = unary_union(area_geom.geoms)

        filtered: list[tuple[dict[str, Any], BaseGeometry]] = []
        for feature, feature_geometry in geometries:
            clipped = feature_geometry.intersection(area_geom)
            if clipped.is_empty:
                continue
            if not clipped.is_valid:
                clipped = make_valid(clipped)
                if not clipped.is_valid:
                    raise AppError(
                        code="INVALID_GEOMETRY",
                        message=f"Area filtering produced invalid geometry for feature in dataset {dataset_id}",
                        status_code=400,
                    )
            filtered.append((feature, clipped))
        return filtered

    @staticmethod
    def _validate_area(db, area_id: UUID | None, project_id: UUID, *, dataset_ids: tuple[UUID, UUID]) -> BaseGeometry | None:
        if not area_id:
            return None
        area = db.get(Area, area_id)
        if not area:
            raise AppError(code="AREA_NOT_FOUND", message="Area not found", status_code=404)
        if area.project_id != project_id:
            raise AppError(code="INVALID_DATASET_SCOPE", message="Area does not belong to this project", status_code=400)
        if area.id in dataset_ids:
            raise AppError(code="INVALID_PARAMETERS", message="area_id must reference an area, not a dataset", status_code=400)

        area_geometry = to_shape(area.geometry)
        if area_geometry.is_empty:
            raise AppError(code="INVALID_GEOMETRY", message="Area geometry is empty", status_code=400)
        return area_geometry

    @staticmethod
    def _feature_identifier(feature: dict[str, Any], fallback_prefix: str, index: int) -> str:
        feature_id = feature.get("id")
        if feature_id is not None:
            return str(feature_id)
        properties = feature.get("properties")
        if isinstance(properties, dict):
            for key in ("vector_feature_id", "source_feature_id", "detection_id", "segmentation_id", "id", "name"):
                value = properties.get(key)
                if value is not None:
                    return str(value)
        for key in ("vector_feature_id", "source_feature_id", "detection_id", "segmentation_id", "class_name", "feature_class"):
            value = feature.get(key)
            if value is not None:
                return str(value)
        return f"{fallback_prefix}-{index + 1}"

    @staticmethod
    def _feature_confidence(feature: dict[str, Any]) -> float | None:
        """Read a detector confidence from a QA feature, if the source has one.

        Vector-vs-vector comparisons have no confidence at all; those fall back
        to identity ordering so the result stays reproducible either way.
        """

        candidates: list[Any] = [feature.get("confidence")]
        properties = feature.get("properties")
        if isinstance(properties, dict):
            candidates.append(properties.get("confidence"))
        for value in candidates:
            if value is None or isinstance(value, bool):
                continue
            try:
                confidence = float(value)
            except (TypeError, ValueError):
                continue
            if confidence == confidence:  # reject NaN
                return confidence
        return None

    @staticmethod
    def _candidate_match_order(
        source_supported: list[tuple[int, dict[str, Any], BaseGeometry]],
    ) -> list[tuple[int, dict[str, Any], BaseGeometry]]:
        """Order candidates the way detection benchmarks do: best score first.

        Greedy IoU matching gives the reference to whichever candidate is
        offered first, so the input order decides both the score and which
        geometry an operator sees as false-positive evidence. Database row
        order is not a defensible answer to that question — every detection in
        a run shares one transaction timestamp — so candidates are ranked by
        confidence, with feature identity as a stable tiebreaker.
        """

        def order_key(entry: tuple[int, dict[str, Any], BaseGeometry]) -> tuple[float, str, int]:
            index, feature, _ = entry
            confidence = QaService._feature_confidence(feature)
            identifier = QaService._feature_identifier(feature, "candidate", index)
            return (-(confidence if confidence is not None else 0.0), identifier, index)

        return sorted(source_supported, key=order_key)

    @staticmethod
    def _match_io_u_evidence(
        source_geometries: list[tuple[dict[str, Any], BaseGeometry]],
        reference_geometries: list[tuple[dict[str, Any], BaseGeometry]],
        iou_threshold: float,
    ) -> QaMatchEvidence:
        source_supported = [
            (index, feature, geom) for index, (feature, geom) in enumerate(source_geometries) if geom.geom_type in QaService.SUPPORTED_GEOMETRY_TYPES
        ]
        reference_supported = [
            (index, feature, geom)
            for index, (feature, geom) in enumerate(reference_geometries)
            if geom.geom_type in QaService.SUPPORTED_GEOMETRY_TYPES
        ]

        unsupported = sorted(
            {
                geom.geom_type
                for _, geom in source_geometries + reference_geometries
                if geom.geom_type not in QaService.SUPPORTED_GEOMETRY_TYPES
            }
        )
        if not source_supported or not reference_supported:
            return QaMatchEvidence(
                false_positives=len(source_supported),
                false_negatives=len(reference_supported),
                warnings=[f"Unsupported geometry types: {unsupported}"] if unsupported else [],
                unsupported=True,
                false_positive_evidence=[
                    {"candidate_feature_id": QaService._feature_identifier(feature, "candidate", source_index)}
                    for source_index, feature, _ in QaService._candidate_match_order(source_supported)
                ],
                false_negative_evidence=[
                    {"reference_feature_id": QaService._feature_identifier(feature, "reference", reference_index)}
                    for reference_index, feature, _ in reference_supported
                ],
            )

        reference_tree = STRtree([geometry for _, _, geometry in reference_supported])
        unmatched_reference_indices = {
            index for index, (_, _, geometry) in enumerate(reference_supported) if geometry.area > 0
        }
        evidence = QaMatchEvidence(warnings=[f"Unsupported geometry types: {unsupported}"] if unsupported else [], unsupported=bool(unsupported))

        for source_index, source_feature, source_geom in QaService._candidate_match_order(source_supported):
            source_feature_id = QaService._feature_identifier(source_feature, "candidate", source_index)
            if source_geom.area <= 0:
                evidence.false_positives += 1
                evidence.false_positive_evidence.append({"candidate_feature_id": source_feature_id})
                continue

            best_iou = 0.0
            best_index = None
            candidate_reference_indices = sorted(int(index) for index in reference_tree.query(source_geom))
            for reference_index in candidate_reference_indices:
                if reference_index not in unmatched_reference_indices:
                    continue
                _, _, reference_geom = reference_supported[reference_index]
                try:
                    intersection = source_geom.intersection(reference_geom)
                except Exception as exc:  # pragma: no cover - robustness path
                    raise AppError(code="GEOMETRY_OPERATION_UNSUPPORTED", message="Geometry operations failed", details={"reason": str(exc)}, status_code=422)

                if intersection.is_empty:
                    continue

                intersection_area = intersection.area
                if intersection_area < 0:
                    intersection_area = 0.0
                union_area = source_geom.area + reference_geom.area - intersection_area
                if union_area <= 0:
                    continue

                candidate_iou = intersection_area / union_area
                if candidate_iou > best_iou:
                    best_iou = candidate_iou
                    best_index = reference_index

            if best_index is not None and best_iou >= iou_threshold:
                reference_original_index, reference_feature, _ = reference_supported[best_index]
                evidence.matches += 1
                evidence.match_iou_values.append(best_iou)
                evidence.match_evidence.append(
                    {
                        "candidate_feature_id": source_feature_id,
                        "reference_feature_id": QaService._feature_identifier(reference_feature, "reference", reference_original_index),
                        "iou": best_iou,
                    }
                )
                unmatched_reference_indices.discard(best_index)
            else:
                evidence.false_positives += 1
                evidence.false_positive_evidence.append({"candidate_feature_id": source_feature_id})

        evidence.false_negatives = len(unmatched_reference_indices)
        for reference_index in sorted(unmatched_reference_indices):
            reference_original_index, reference_feature, _ = reference_supported[reference_index]
            evidence.false_negative_evidence.append(
                {"reference_feature_id": QaService._feature_identifier(reference_feature, "reference", reference_original_index)}
            )

        return evidence

    @staticmethod
    def _match_io_u_metrics(
        source_geometries: list[tuple[dict[str, Any], BaseGeometry]],
        reference_geometries: list[tuple[dict[str, Any], BaseGeometry]],
        iou_threshold: float,
    ) -> tuple[int, int, int, list[float], list[str], bool]:
        evidence = QaService._match_io_u_evidence(source_geometries, reference_geometries, iou_threshold)
        return (
            evidence.matches,
            evidence.false_positives,
            evidence.false_negatives,
            evidence.match_iou_values,
            evidence.warnings,
            evidence.unsupported,
        )

    @staticmethod
    def compare_candidate_with_reference(
        db,
        project_id: UUID,
        candidate_dataset_id: UUID,
        reference_dataset_id: UUID,
        iou_threshold: float = 0.5,
        area_id: UUID | None = None,
    ) -> QaProviderComparisonResult:
        if candidate_dataset_id == reference_dataset_id:
            raise AppError(code="INVALID_PARAMETERS", message="Candidate and reference dataset must differ", status_code=400)

        candidate_dataset, candidate_payload, candidate_geometries = QaService._load_dataset_payload(
            db,
            candidate_dataset_id,
            expected_project_id=project_id,
        )
        reference_dataset, reference_payload, reference_geometries = QaService._load_dataset_payload(
            db,
            reference_dataset_id,
            expected_project_id=project_id,
        )
        DatasetConsumptionGate.assert_eligible(candidate_dataset, purpose="quality_assessment")
        DatasetConsumptionGate.assert_eligible(
            reference_dataset,
            purpose="reference_validation",
            reference_task="building_validation",
        )

        area_geometry = QaService._validate_area(
            db,
            area_id=area_id,
            project_id=project_id,
            dataset_ids=(candidate_dataset_id, reference_dataset_id),
        )

        candidate_feature_count_raw = len(candidate_payload.get("features", [])) if isinstance(candidate_payload, dict) else 0
        reference_feature_count_raw = len(reference_payload.get("features", [])) if isinstance(reference_payload, dict) else 0

        if area_geometry is not None:
            candidate_geometries = QaService._apply_area_filter(candidate_geometries, area_geometry, dataset_id=candidate_dataset.id)
            reference_geometries = QaService._apply_area_filter(reference_geometries, area_geometry, dataset_id=reference_dataset.id)

        evidence = QaService._match_io_u_evidence(
            candidate_geometries,
            reference_geometries,
            iou_threshold,
        )

        # Report the population the metrics were computed over, not the raw
        # dataset totals: with an area filter the two differ, and a count that
        # disagrees with matches + false positives is unreadable as evidence.
        candidate_feature_count = len(candidate_geometries)
        reference_feature_count = len(reference_geometries)
        mean_iou = None if not evidence.match_iou_values else sum(evidence.match_iou_values) / len(evidence.match_iou_values)

        precision = None
        if evidence.matches + evidence.false_positives > 0:
            precision = evidence.matches / (evidence.matches + evidence.false_positives)

        recall = None
        if evidence.matches + evidence.false_negatives > 0:
            recall = evidence.matches / (evidence.matches + evidence.false_negatives)

        f1_score = None
        if precision is not None and recall is not None and precision + recall > 0:
            f1_score = (2 * precision * recall) / (precision + recall)

        status = "unsupported" if evidence.unsupported else "ok"
        return QaProviderComparisonResult(
            status=status,
            warnings=_extract_crs_warnings(candidate_dataset, reference_dataset) + evidence.warnings,
            candidate_feature_count=candidate_feature_count,
            reference_feature_count=reference_feature_count,
            candidate_feature_count_raw=candidate_feature_count_raw,
            reference_feature_count_raw=reference_feature_count_raw,
            matches=evidence.matches,
            false_positives=evidence.false_positives,
            false_negatives=evidence.false_negatives,
            precision=precision,
            recall=recall,
            f1_score=f1_score,
            mean_iou=mean_iou,
            iou_threshold=iou_threshold,
            unsupported_geometry=evidence.unsupported,
            unsupported_geometries=evidence.warnings,
            match_evidence=evidence.match_evidence,
            false_positive_evidence=evidence.false_positive_evidence,
            false_negative_evidence=evidence.false_negative_evidence,
            generated_at=datetime.now(timezone.utc),
        )