Defining custom pyproj transformation pipelines

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from pyproj import Transformer
from pyproj.transformer import AreaOfInterest, TransformerGroup


@dataclass(frozen=True)
class OperationChoice:
    """A single candidate transformation between two CRS."""
    index: int
    name: str
    accuracy_m: float | None
    available: bool


def list_candidate_operations(
    src_crs: str,
    dst_crs: str,
    aoi: AreaOfInterest | None = None,
) -> list[OperationChoice]:
    """
    Enumerate every datum-shift operation PROJ can build between two CRS.

    Parameters
    ----------
    src_crs, dst_crs : str
        CRS identifiers, e.g. "EPSG:4326", "EPSG:25832".
    aoi : AreaOfInterest, optional
        Bounding box used to rank operations by regional validity.

    Returns
    -------
    list[OperationChoice]
        One entry per candidate, in PROJ's ranked order (best first).

    Raises
    ------
    ValueError
        If PROJ can build no operation between the two CRS.
    """
    group = TransformerGroup(src_crs, dst_crs, area_of_interest=aoi)
    if not group.transformers:
        raise ValueError(f"No transformation found: {src_crs} -> {dst_crs}")

    choices: list[OperationChoice] = []
    for i, t in enumerate(group.transformers):
        op = t.operations[0] if t.operations else None
        acc = getattr(op, "accuracy", None) if op else None
        choices.append(
            OperationChoice(
                index=i,
                name=t.description,
                accuracy_m=float(acc) if acc not in (None, "") else None,
                available=t.is_available,
            )
        )
    return choices


def build_pinned_transformer(
    src_crs: str,
    dst_crs: str,
    bbox: tuple[float, float, float, float] | None = None,
) -> Transformer:
    """
    Build a reproducible Transformer, pinning the best *available* operation.

    bbox is (west, south, east, north) in degrees. Supplying it lets PROJ
    prefer a regional grid over a global low-accuracy fallback.

    always_xy=True forces longitude-latitude (easting-northing) order so
    coordinate arrays never silently swap between CRS definitions.

    Raises
    ------
    ValueError
        If no *available* operation exists (e.g. required grid not installed).
    """
    aoi = AreaOfInterest(*bbox) if bbox is not None else None
    candidates = list_candidate_operations(src_crs, dst_crs, aoi)

    available = [c for c in candidates if c.available]
    if not available:
        raise ValueError(
            f"No installed operation for {src_crs} -> {dst_crs}. "
            f"Install the required PROJ transformation grids."
        )

    # Pin the top-ranked *available* candidate. from_crs with the same AOI
    # reconstructs exactly this operation deterministically.
    return Transformer.from_crs(
        src_crs, dst_crs, always_xy=True, area_of_interest=aoi
    )


def build_pipeline_transformer(pipeline: str) -> Transformer:
    """
    Build a Transformer from an explicit +proj=pipeline string.

    Use this when you must pin an exact multi-step chain, e.g. a named
    NTv2 grid, regardless of which grids PROJ would otherwise rank first.
    """
    if "+proj=pipeline" not in pipeline:
        raise ValueError("pipeline must be a '+proj=pipeline' string")
    return Transformer.from_pipeline(pipeline)


def transform_arrays(
    transformer: Transformer,
    lon: np.ndarray,
    lat: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Vectorized reprojection of longitude/latitude arrays (lon, lat order).

    Returns projected (x, y) in the target CRS units — metres for a UTM
    or national metric grid. Never loop per row; pass whole arrays.
    """
    lon = np.asarray(lon, dtype=float)
    lat = np.asarray(lat, dtype=float)
    if lon.shape != lat.shape:
        raise ValueError("lon and lat must share the same shape")
    if lon.size == 0:
        return lon.copy(), lat.copy()  # empty in, empty out
    x, y = transformer.transform(lon, lat)
    return np.asarray(x), np.asarray(y)

+proj=pipeline
  +step +proj=unitconvert +xy_in=deg +xy_out=rad
  +step +proj=utm +zone=32 +ellps=WGS84

def validate_round_trip(
    fwd: Transformer,
    inv: Transformer,
    lon: np.ndarray,
    lat: np.ndarray,
    tol_m: float = 1e-3,
) -> dict[str, float]:
    """
    Reproject forward then inverse and assert the residual is negligible.

    The residual is measured in the metric projected CRS (metres), NOT in
    degrees: a fixed degree error is a different ground distance at every
    latitude, so degree-space tolerances are meaningless.

    Raises AssertionError if the maximum residual exceeds tol_m.
    """
    if lon.size == 0:
        return {"max_residual_m": 0.0, "n_points": 0}

    # Forward: WGS84 -> metric grid (x, y in metres)
    x, y = transform_arrays(fwd, lon, lat)
    # Inverse: metric grid -> WGS84 (degrees)
    lon_rt, lat_rt = inv.transform(x, y)
    # Re-project the round-tripped degrees back to the metric grid so the
    # residual is a true ground distance in metres.
    x_rt, y_rt = transform_arrays(fwd, lon_rt, lat_rt)

    residual_m = np.sqrt((x - x_rt) ** 2 + (y - y_rt) ** 2)
    max_res = float(np.max(residual_m))
    assert max_res <= tol_m, (
        f"Round-trip residual {max_res:.6f} m exceeds tolerance {tol_m} m — "
        f"the forward and inverse operations are not consistent."
    )
    return {"max_residual_m": max_res, "n_points": int(lon.size)}