s2sr_pipe/utils/geo_utils.py

"""
Rasterio / geospatial helpers used across the pipeline.
"""
from __future__ import annotations

from pathlib import Path
from typing import Dict, Tuple

import numpy as np
import rasterio
from rasterio.crs import CRS
from rasterio.enums import Resampling
from rasterio.transform import Affine
from rasterio.warp import transform_bounds

from s2sr_pipe.utils.logging_utils import get_logger

logger = get_logger("geo_utils")


def assert_same_crs(profile_a: dict, profile_b: dict) -> None:
    """Raise if two rasterio profiles have different CRS."""
    crs_a = profile_a.get("crs")
    crs_b = profile_b.get("crs")
    if crs_a != crs_b:
        raise ValueError(f"CRS mismatch: {crs_a} vs {crs_b}")


def build_output_profile(
    ref_profile: dict,
    n_bands: int,
    scale: int = 1,
    dtype: str = "uint16",
    compress: str = "deflate",
    predictor: int = 2,
    tiled: bool = True,
    blockxsize: int = 512,
    blockysize: int = 512,
) -> dict:
    """
    Build a rasterio write profile suitable for large Cloud-Optimised GeoTIFF.

    Parameters
    ----------
    ref_profile  : Profile from the 10 m reference band.
    n_bands      : Number of output bands.
    scale        : Super-resolution upscale factor (1 = no change, 4 = 10m→2.5m).
    dtype        : Output data type.
    compress     : Compression codec.
    predictor    : DEFLATE predictor (2 = horizontal differencing, good for int).
    tiled        : Write as tiled TIFF (required for COG / streaming reads).
    blockxsize/blockysize : Internal tile size.

    Returns
    -------
    dict  : rasterio write profile.
    """
    profile = ref_profile.copy()

    # Scale the transform for super-resolved output (e.g. ×4: 10 m → 2.5 m)
    if scale > 1:
        original_transform = profile["transform"]
        sx = original_transform.a / scale   # pixel width  (positive)
        sy = original_transform.e / scale   # pixel height (negative)
        new_transform = Affine(
            sx, 0.0, original_transform.c,
            0.0, sy, original_transform.f,
        )
        profile["transform"] = new_transform
        profile["width"]  = profile["width"]  * scale
        profile["height"] = profile["height"] * scale

    profile.update(
        {
            "driver": "GTiff",
            "count": n_bands,
            "dtype": dtype,
            "compress": compress,
            "predictor": predictor,
            "tiled": tiled,
            "blockxsize": blockxsize,
            "blockysize": blockysize,
            "interleave": "pixel",
            "BIGTIFF": "YES",   # mandatory for files > 4 GB (43920×43920×10 ≈ 18 GB)
        }
    )
    return profile


def write_geotiff(
    array: np.ndarray,
    profile: dict,
    output_path: Path | str,
) -> None:
    """
    Write a (C, H, W) numpy array to a GeoTIFF.

    The array is clipped to the dtype's valid range before writing.

    Parameters
    ----------
    array       : (C, H, W) numpy array.
    profile     : rasterio write profile (must include crs, transform, dtype).
    output_path : Destination path.
    """
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    dtype = profile["dtype"]
    if dtype == "uint16":
        array = np.clip(array, 0, 10_000).astype(np.uint16)
    elif dtype == "float32":
        array = array.astype(np.float32)

    logger.info(f"Writing GeoTIFF -> {output_path}  shape={array.shape}  dtype={dtype}")
    with rasterio.open(output_path, "w", **profile) as dst:
        dst.write(array)
    logger.info("GeoTIFF written successfully.")


def get_pixel_size(transform: Affine) -> Tuple[float, float]:
    """Return (pixel_width, pixel_height) in CRS units from an Affine transform."""
    return abs(transform.a), abs(transform.e)


def roi_wgs84_to_pixel_window(
    roi: Tuple[float, float, float, float],
    profile: dict,
    buffer_px: int = 128,
) -> Dict:
    """
    Convertit un ROI WGS84 en fenêtres pixel sur l'image LR.

    Parameters
    ----------
    roi       : (min_lon, min_lat, max_lon, max_lat) en WGS84 (EPSG:4326).
    profile   : Profil rasterio de l'image LR (doit avoir crs, transform, width, height).
    buffer_px : Buffer en pixels LR ajouté de chaque côté pour le contexte du modèle.

    Returns
    -------
    dict :
      "read"          — (rr_min, rr_max, rc_min, rc_max) : fenêtre à lire (avec buffer)
      "inner"         — (ir_start, ir_end, ic_start, ic_end) : position du ROI dans le crop
      "roi_transform" — Affine : transform du coin sup-gauche du ROI exact
      "roi_width"     — int
      "roi_height"    — int

    Raises
    ------
    ValueError si le ROI ne chevauche pas la dalle.
    """
    min_lon, min_lat, max_lon, max_lat = roi
    T = profile["transform"]
    H, W = profile["height"], profile["width"]

    # WGS84 → CRS image
    xmin, ymin, xmax, ymax = transform_bounds(
        CRS.from_epsg(4326), profile["crs"],
        min_lon, min_lat, max_lon, max_lat,
    )

    # World coords → pixel coords (inverse affine)
    col_min_f, row_min_f = ~T * (xmin, ymax)   # coin haut-gauche
    col_max_f, row_max_f = ~T * (xmax, ymin)   # coin bas-droit

    col_min = int(np.floor(col_min_f))
    row_min = int(np.floor(row_min_f))
    col_max = int(np.ceil(col_max_f))
    row_max = int(np.ceil(row_max_f))

    # Vérifier que le ROI chevauche la dalle avant clamp
    if col_max <= 0 or col_min >= W or row_max <= 0 or row_min >= H:
        raise ValueError(
            f"Le ROI ({min_lon},{min_lat},{max_lon},{max_lat}) ne chevauche pas "
            f"la dalle (CRS: {profile.get('crs')}, taille: {W}x{H} px). "
            f"Pixel bounds calculés : cols {col_min}..{col_max}, rows {row_min}..{row_max}."
        )

    # Clamp ROI au bord de l'image
    col_min_orig, row_min_orig = col_min, row_min
    col_max_orig, row_max_orig = col_max, row_max
    col_min = max(0, col_min)
    row_min = max(0, row_min)
    col_max = min(W, col_max)
    row_max = min(H, row_max)

    clipped = []
    if col_min_orig < 0:
        clipped.append(f"gauche ({col_min_orig}px -> 0)")
    if row_min_orig < 0:
        clipped.append(f"haut ({row_min_orig}px -> 0)")
    if col_max_orig > W:
        clipped.append(f"droite ({col_max_orig}px -> {W})")
    if row_max_orig > H:
        clipped.append(f"bas ({row_max_orig}px -> {H})")
    if clipped:
        logger.warning(
            "Le ROI depasse la dalle et a ete clippe : %s. "
            "Seule la partie disponible sera traitee.",
            ", ".join(clipped),
        )

    # Fenêtre bufférisée (clampée aux bords)
    rc_min = max(0, col_min - buffer_px)
    rc_max = min(W, col_max + buffer_px)
    rr_min = max(0, row_min - buffer_px)
    rr_max = min(H, row_max + buffer_px)

    logger.debug(
        "ROI pixel bounds: cols %d..%d rows %d..%d | "
        "buffered read: cols %d..%d rows %d..%d",
        col_min, col_max, row_min, row_max,
        rc_min, rc_max, rr_min, rr_max,
    )

    # Position du ROI dans le crop bufférisé
    ic_start = col_min - rc_min
    ic_end   = col_max - rc_min
    ir_start = row_min - rr_min
    ir_end   = row_max - rr_min

    # Transform du coin sup-gauche du ROI exact (pour le profil de sortie)
    roi_transform = T * Affine.translation(col_min, row_min)

    return {
        "read":          (rr_min, rr_max, rc_min, rc_max),
        "inner":         (ir_start, ir_end, ic_start, ic_end),
        "roi_transform": roi_transform,
        "roi_width":     col_max - col_min,
        "roi_height":    row_max - row_min,
    }