geointel/scripts/audit_evaluation_lineage_spatial_independence.py

#!/usr/bin/env python3
"""Audit fresh evaluation AOIs against every raster in a model's corpus lineage."""

from __future__ import annotations

import argparse
import hashlib
import json
from pathlib import Path
from typing import Any

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


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def metric_box(bounds: list[float] | tuple[float, ...]) -> Any:
    transformer = Transformer.from_crs("EPSG:4326", "EPSG:31370", always_xy=True)
    return shapely_transform(transformer.transform, box(*map(float, bounds)))


def raster_bounds_epsg4326(path: Path) -> list[float]:
    import rasterio
    from rasterio.warp import transform_bounds

    with rasterio.open(path) as dataset:
        if dataset.crs is None:
            raise ValueError(f"lineage raster has no CRS: {path}")
        bounds = transform_bounds(dataset.crs, "EPSG:4326", *dataset.bounds)
    return list(map(float, bounds))


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--evaluation-manifest", type=Path, required=True)
    parser.add_argument("--lineage-summary", type=Path, action="append", required=True)
    parser.add_argument("--manifest-search-root", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--minimum-distance-m", type=float, default=2000.0)
    args = parser.parse_args()
    if args.output.exists():
        parser.error(f"output already exists: {args.output}")
    evaluation_manifest = args.evaluation_manifest.resolve(strict=True)
    evaluation = json.loads(evaluation_manifest.read_text(encoding="utf-8-sig"))
    evaluation_rows = evaluation.get("samples")
    if not isinstance(evaluation_rows, list) or not evaluation_rows:
        raise SystemExit("evaluation manifest contains no samples")

    lineage_slugs: set[str] = set()
    summary_evidence: list[dict[str, Any]] = []
    for raw in args.lineage_summary:
        summary = raw.resolve(strict=True)
        payload = json.loads(summary.read_text(encoding="utf-8"))
        tiles = payload.get("tiles")
        if not isinstance(tiles, list):
            raise SystemExit(f"lineage summary has no tiles: {summary}")
        slugs = {str(row.get("sample_slug") or "").strip() for row in tiles}
        slugs.discard("")
        lineage_slugs.update(slugs)
        summary_evidence.append(
            {"path": str(summary), "sha256": sha256_file(summary), "sample_slugs": sorted(slugs)}
        )

    candidates: dict[str, dict[str, dict[str, Any]]] = {
        slug: {} for slug in lineage_slugs
    }
    for manifest_path in sorted(
        args.manifest_search_root.rglob("operator_samples_manifest.json")
    ):
        try:
            payload = json.loads(manifest_path.read_text(encoding="utf-8-sig"))
        except (OSError, json.JSONDecodeError):
            continue
        for sample in payload.get("samples") or []:
            if not isinstance(sample, dict):
                continue
            slug = str(sample.get("sample_slug") or "").strip()
            raster_raw = sample.get("raster_path")
            if slug not in candidates or not isinstance(raster_raw, str):
                continue
            raster = Path(raster_raw)
            if not raster.is_file():
                continue
            key = str(raster.resolve())
            candidates[slug].setdefault(
                key,
                {
                    "sample_slug": slug,
                    "raster_path": key,
                    "declared_raster_sha256": sample.get("raster_sha256"),
                    "source_manifest_path": str(manifest_path),
                    "source_manifest_sha256": sha256_file(manifest_path),
                },
            )

    missing = sorted(slug for slug, rows in candidates.items() if not rows)
    lineage_geometries: list[tuple[dict[str, Any], Any]] = []
    for rows in candidates.values():
        for record in rows.values():
            bounds = raster_bounds_epsg4326(Path(record["raster_path"]))
            record["bbox_epsg4326"] = bounds
            lineage_geometries.append((record, metric_box(bounds)))

    results: list[dict[str, Any]] = []
    for sample in evaluation_rows:
        slug = str(sample.get("sample_slug") or "").strip()
        bounds = sample.get("bbox_epsg4326")
        if not isinstance(bounds, list) or len(bounds) != 4:
            raise SystemExit(f"evaluation sample has no governed bbox: {slug}")
        geometry = metric_box(bounds)
        distances = [
            (float(geometry.distance(other)), record)
            for record, other in lineage_geometries
        ]
        distances.sort(key=lambda row: row[0])
        minimum, nearest = distances[0]
        results.append(
            {
                "sample_slug": slug,
                "bbox_epsg4326": bounds,
                "minimum_lineage_distance_m": minimum,
                "nearest_lineage_sample_slug": nearest["sample_slug"],
                "nearest_lineage_raster_path": nearest["raster_path"],
                "passes_minimum_distance": minimum >= args.minimum_distance_m,
            }
        )

    status = "independent" if not missing and all(
        row["passes_minimum_distance"] for row in results
    ) else "failed"
    output = {
        "schema_version": 1,
        "status": status,
        "minimum_distance_m": args.minimum_distance_m,
        "evaluation_manifest_path": str(evaluation_manifest),
        "evaluation_manifest_sha256": sha256_file(evaluation_manifest),
        "lineage_summaries": summary_evidence,
        "lineage_sample_count": len(lineage_slugs),
        "resolved_lineage_sample_count": len(lineage_slugs) - len(missing),
        "missing_lineage_sample_slugs": missing,
        "evaluation_samples": results,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(output, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    print(json.dumps(output, indent=2, sort_keys=True))
    return 0 if status == "independent" else 2


if __name__ == "__main__":
    raise SystemExit(main())