"""
Landcover Classification Utilities

This module provides utilities for discrete landcover classification workflows,
including tile export with background filtering and radiometric normalization
for multi-temporal image comparability.

Key Features:
- Enhanced tile filtering with configurable feature ratio thresholds
- Separate statistics tracking for different skip reasons
- LIRRN (Location-Independent Relative Radiometric Normalization)
- Maintains full compatibility with base geoai workflow
- Optimized for discrete landcover classification tasks

Date: November 2025
"""

import logging
import os
import warnings
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

logger = logging.getLogger(__name__)

import geopandas as gpd
import numpy as np
import rasterio
from rasterio import features
from rasterio.windows import Window
from skimage.filters import threshold_multiotsu
from sklearn.linear_model import LinearRegression
from tqdm import tqdm

__all__ = [
    "export_landcover_tiles",
    "normalize_radiometric",
]


def export_landcover_tiles(
    in_raster: str,
    out_folder: str,
    in_class_data: Optional[Union[str, gpd.GeoDataFrame]] = None,
    tile_size: int = 256,
    stride: int = 128,
    class_value_field: str = "class",
    buffer_radius: float = 0,
    max_tiles: Optional[int] = None,
    quiet: bool = False,
    all_touched: bool = True,
    create_overview: bool = False,
    skip_empty_tiles: bool = False,
    min_feature_ratio: Union[bool, float] = False,
    metadata_format: str = "PASCAL_VOC",
) -> Dict[str, Any]:
    """
    Export GeoTIFF tiles optimized for landcover classification training.

    This function extends the base export_geotiff_tiles with enhanced filtering
    capabilities specifically designed for discrete landcover classification.
    It can filter out tiles dominated by background pixels to improve training
    data quality and reduce dataset size.

    Args:
        in_raster: Path to input raster (image to tile)
        out_folder: Output directory for tiles
        in_class_data: Path to vector mask or GeoDataFrame (optional for image-only export)
        tile_size: Size of output tiles in pixels (default: 256)
        stride: Stride for sliding window (default: 128)
        class_value_field: Field name containing class values (default: "class")
        buffer_radius: Buffer radius around features in pixels (default: 0)
        max_tiles: Maximum number of tiles to export (default: None)
        quiet: Suppress progress output (default: False)
        all_touched: Include pixels touched by geometry (default: True)
        create_overview: Create overview image showing tile locations (default: False)
        skip_empty_tiles: Skip tiles with no features (default: False)
        min_feature_ratio: Minimum ratio of non-background pixels required to keep tile
            - False: Disable ratio filtering (default)
            - 0.0-1.0: Minimum ratio threshold (e.g., 0.1 = 10% features required)
        metadata_format: Annotation format ("PASCAL_VOC" or "YOLO")

    Returns:
        Dictionary containing:
            - tiles_exported: Number of tiles successfully exported
            - tiles_skipped_empty: Number of completely empty tiles skipped
            - tiles_skipped_ratio: Number of tiles filtered by min_feature_ratio
            - output_dirs: Dictionary with paths to images and labels directories

    Examples:
        # Original behavior (no filtering)
        export_landcover_tiles(
            "input.tif",
            "output",
            "mask.shp",
            skip_empty_tiles=True
        )

        # Light filtering (keep tiles with ≥5% features)
        export_landcover_tiles(
            "input.tif",
            "output",
            "mask.shp",
            skip_empty_tiles=True,
            min_feature_ratio=0.05
        )

        # Moderate filtering (keep tiles with ≥15% features)
        export_landcover_tiles(
            "input.tif",
            "output",
            "mask.shp",
            skip_empty_tiles=True,
            min_feature_ratio=0.15
        )

    Note:
        This function is designed for discrete landcover classification where
        class 0 typically represents background/no data. The min_feature_ratio
        parameter counts non-zero pixels as "features".
    """

    # Validate min_feature_ratio parameter
    if min_feature_ratio is not False:
        if not isinstance(min_feature_ratio, (int, float)):
            warnings.warn(
                f"min_feature_ratio must be a number between 0.0 and 1.0, got {type(min_feature_ratio)}. "
                "Disabling ratio filtering."
            )
            min_feature_ratio = False
        elif not (0.0 <= min_feature_ratio <= 1.0):
            warnings.warn(
                f"min_feature_ratio must be between 0.0 and 1.0, got {min_feature_ratio}. "
                "Disabling ratio filtering."
            )
            min_feature_ratio = False

    # Create output directories
    out_folder = Path(out_folder)
    out_folder.mkdir(parents=True, exist_ok=True)

    images_dir = out_folder / "images"
    labels_dir = out_folder / "labels"
    images_dir.mkdir(exist_ok=True)
    labels_dir.mkdir(exist_ok=True)

    if metadata_format == "PASCAL_VOC":
        ann_dir = out_folder / "annotations"
        ann_dir.mkdir(exist_ok=True)

    # Initialize statistics
    stats = {
        "tiles_exported": 0,
        "tiles_skipped_empty": 0,
        "tiles_skipped_ratio": 0,
        "output_dirs": {"images": str(images_dir), "labels": str(labels_dir)},
    }

    # Open raster
    with rasterio.open(in_raster) as src:
        height, width = src.shape

        # Detect if in_class_data is raster or vector
        is_class_data_raster = False
        class_src = None
        gdf = None
        mask_array = None

        if in_class_data is not None:
            if isinstance(in_class_data, str):
                file_ext = Path(in_class_data).suffix.lower()
                if file_ext in [
                    ".tif",
                    ".tiff",
                    ".img",
                    ".jp2",
                    ".png",
                    ".bmp",
                    ".gif",
                ]:
                    try:
                        # Try to open as raster
                        class_src = rasterio.open(in_class_data)
                        is_class_data_raster = True

                        # Verify CRS match
                        if class_src.crs != src.crs:
                            if not quiet:
                                logger.warning(
                                    f"CRS mismatch between image ({src.crs}) and mask ({class_src.crs})"
                                )
                    except Exception as e:
                        is_class_data_raster = False
                        if not quiet:
                            logger.debug(
                                f"Could not open as raster, trying vector: {e}"
                            )

                # If not raster or raster open failed, try vector
                if not is_class_data_raster:
                    gdf = gpd.read_file(in_class_data)

                    # Reproject if needed
                    if gdf.crs != src.crs:
                        if not quiet:
                            logger.info(
                                f"Reprojecting mask from {gdf.crs} to {src.crs}"
                            )
                        gdf = gdf.to_crs(src.crs)

                    # Apply buffer if requested
                    if buffer_radius > 0:
                        gdf.geometry = gdf.geometry.buffer(buffer_radius)

                    # For vector data, rasterize entire mask up front for efficiency
                    shapes = [
                        (geom, value)
                        for geom, value in zip(gdf.geometry, gdf[class_value_field])
                    ]
                    mask_array = features.rasterize(
                        shapes,
                        out_shape=(height, width),
                        transform=src.transform,
                        all_touched=all_touched,
                        fill=0,
                        dtype=np.uint8,
                    )
            else:
                # Assume GeoDataFrame passed directly
                gdf = in_class_data

                # Reproject if needed
                if gdf.crs != src.crs:
                    if not quiet:
                        logger.info("Reprojecting mask from %s to %s", gdf.crs, src.crs)
                    gdf = gdf.to_crs(src.crs)

                # Apply buffer if requested
                if buffer_radius > 0:
                    gdf.geometry = gdf.geometry.buffer(buffer_radius)

                # Rasterize entire mask up front
                shapes = [
                    (geom, value)
                    for geom, value in zip(gdf.geometry, gdf[class_value_field])
                ]
                mask_array = features.rasterize(
                    shapes,
                    out_shape=(height, width),
                    transform=src.transform,
                    all_touched=all_touched,
                    fill=0,
                    dtype=np.uint8,
                )

        # Calculate tile positions
        tile_positions = []
        for y in range(0, height - tile_size + 1, stride):
            for x in range(0, width - tile_size + 1, stride):
                tile_positions.append((x, y))

        if max_tiles:
            tile_positions = tile_positions[:max_tiles]

        # Process tiles
        pbar = tqdm(tile_positions, desc="Exporting tiles", disable=quiet)

        for tile_idx, (x, y) in enumerate(pbar):
            window = Window(x, y, tile_size, tile_size)

            # Read image tile
            image_tile = src.read(window=window)

            # Read mask tile based on data type
            mask_tile = None
            has_features = False

            if is_class_data_raster and class_src is not None:
                # For raster masks, read directly from the raster source
                # Get window transform and bounds
                window_transform = src.window_transform(window)
                minx = window_transform[2]
                maxy = window_transform[5]
                maxx = minx + tile_size * window_transform[0]
                miny = maxy + tile_size * window_transform[4]

                # Get corresponding window in class raster
                window_class = rasterio.windows.from_bounds(
                    minx, miny, maxx, maxy, class_src.transform
                )

                try:
                    # Read label data from raster
                    mask_tile = class_src.read(
                        1,
                        window=window_class,
                        boundless=True,
                        out_shape=(tile_size, tile_size),
                    )

                    # Check if tile has features
                    has_features = np.any(mask_tile > 0)
                except Exception as e:
                    if not quiet:
                        pbar.write(f"Error reading mask tile at ({x}, {y}): {e}")
                    continue

            elif mask_array is not None:
                # For vector masks (pre-rasterized)
                mask_tile = mask_array[y : y + tile_size, x : x + tile_size]
                has_features = np.any(mask_tile > 0)

            # Skip empty tiles if requested
            if skip_empty_tiles and not has_features:
                stats["tiles_skipped_empty"] += 1
                continue

            # Apply min_feature_ratio filtering if enabled
            if skip_empty_tiles and has_features and min_feature_ratio is not False:
                # Calculate ratio of non-background pixels
                total_pixels = mask_tile.size
                feature_pixels = np.sum(mask_tile > 0)
                feature_ratio = feature_pixels / total_pixels

                # Skip tile if below threshold
                if feature_ratio < min_feature_ratio:
                    stats["tiles_skipped_ratio"] += 1
                    continue

            # Save image tile
            tile_name = f"tile_{tile_idx:06d}.tif"
            image_path = images_dir / tile_name

            # Get transform for this tile
            tile_transform = src.window_transform(window)

            # Write image
            with rasterio.open(
                image_path,
                "w",
                driver="GTiff",
                height=tile_size,
                width=tile_size,
                count=src.count,
                dtype=src.dtypes[0],
                crs=src.crs,
                transform=tile_transform,
                compress="lzw",
            ) as dst:
                dst.write(image_tile)

            # Save mask tile if available
            if mask_tile is not None:
                mask_path = labels_dir / tile_name
                with rasterio.open(
                    mask_path,
                    "w",
                    driver="GTiff",
                    height=tile_size,
                    width=tile_size,
                    count=1,
                    dtype=np.uint8,
                    crs=src.crs,
                    transform=tile_transform,
                    compress="lzw",
                ) as dst:
                    dst.write(mask_tile, 1)

            stats["tiles_exported"] += 1

            # Update progress bar description with selection count
            if not quiet:
                pbar.set_description(
                    f"Exporting tiles ({stats['tiles_exported']}/{tile_idx + 1})"
                )

    # Close raster class source if opened
    if class_src is not None:
        class_src.close()

    # Print summary
    if not quiet:
        logger.info(f"\n{'='*60}")
        logger.info("TILE EXPORT SUMMARY")
        logger.info(f"{'='*60}")
        logger.info(f"Tiles exported: {stats['tiles_exported']}/{len(tile_positions)}")
        if skip_empty_tiles:
            logger.info(f"Tiles skipped (empty): {stats['tiles_skipped_empty']}")
        if min_feature_ratio is not False:
            logger.info(
                f"Tiles skipped (low feature ratio < {min_feature_ratio}): {stats['tiles_skipped_ratio']}"
            )
        logger.info(f"\nOutput directories:")
        logger.info(f"  Images: {stats['output_dirs']['images']}")
        logger.info(f"  Labels: {stats['output_dirs']['labels']}")
        logger.info(f"{'='*60}\n")

    return stats


# ---------------------------------------------------------------------------
# Radiometric Normalization
# ---------------------------------------------------------------------------


def _load_raster(filepath: str) -> Tuple[np.ndarray, dict]:
    """Load a multi-band raster as a (H, W, B) float64 array.

    Args:
        filepath: Path to the raster file.

    Returns:
        Tuple of (image_array, profile) where image_array has shape (H, W, B)
        and profile is the rasterio dataset profile dict.

    Raises:
        FileNotFoundError: If the file does not exist.
    """
    if not os.path.isfile(filepath):
        raise FileNotFoundError(f"Raster file not found: {filepath}")
    with rasterio.open(filepath) as src:
        img = src.read()  # (B, H, W)
        profile = src.profile.copy()
    img = np.moveaxis(img, 0, -1).astype(np.float64)
    if np.any(~np.isfinite(img)):
        warnings.warn(
            "Raster contains NaN or infinite values; replacing with 0.",
            stacklevel=2,
        )
        img = np.nan_to_num(img, nan=0.0, posinf=0.0, neginf=0.0)
    return img, profile


def _save_raster(filepath: str, img: np.ndarray, profile: dict) -> None:
    """Save a (H, W, B) array as a multi-band GeoTIFF.

    Args:
        filepath: Output file path.
        img: Image array with shape (H, W, B).
        profile: Rasterio profile dict (from ``_load_raster``).

    Raises:
        ValueError: If *img* is not 3-dimensional.
    """
    if img.ndim != 3:
        raise ValueError(f"Expected 3-D array (H, W, B), got shape {img.shape}")
    out_profile = profile.copy()
    out_profile.update(
        dtype="float64",
        count=img.shape[2],
        height=img.shape[0],
        width=img.shape[1],
        compress="lzw",
    )
    with rasterio.open(filepath, "w", **out_profile) as dst:
        for i in range(img.shape[2]):
            dst.write(img[:, :, i], i + 1)


def _compute_distances(
    p_n: int,
    a1: np.ndarray,
    b1: np.ndarray,
    id_indices: np.ndarray,
) -> Tuple[np.ndarray, np.ndarray]:
    """Select *p_n* samples closest to the maximum value, then subsample.

    Args:
        p_n: Number of candidate samples to draw from each array.
        a1: Non-zero reference pixel values (1-D).
        b1: Non-zero subject pixel values (1-D).
        id_indices: Random indices used to subsample from the candidates.

    Returns:
        Tuple of (sub_samples, ref_samples) after subsampling.
    """
    max_ref = np.max(a1)
    idx_ref = np.argsort(np.abs(a1 - max_ref))
    ref_candidates = a1[idx_ref[:p_n]]

    max_sub = np.max(b1)
    idx_sub = np.argsort(np.abs(b1 - max_sub))
    sub_candidates = b1[idx_sub[:p_n]]

    safe_indices = id_indices[
        id_indices < min(len(sub_candidates), len(ref_candidates))
    ]
    if len(safe_indices) == 0:
        return sub_candidates, ref_candidates
    return sub_candidates[safe_indices], ref_candidates[safe_indices]


def _compute_sample(
    p_n: int,
    a1: np.ndarray,
    b1: np.ndarray,
    id_indices: np.ndarray,
    num_sampling_rounds: int = 3,
) -> Tuple[np.ndarray, np.ndarray]:
    """Multiple rounds of distance-based sampling, concatenated (non-zero only).

    Args:
        p_n: Number of samples per round.
        a1: Non-zero reference pixel values (1-D).
        b1: Non-zero subject pixel values (1-D).
        id_indices: Random subsample indices.
        num_sampling_rounds: Number of sampling rounds.

    Returns:
        Tuple of (sub_combined, ref_combined).
    """
    pairs = [
        _compute_distances(p_n, a1, b1, id_indices) for _ in range(num_sampling_rounds)
    ]
    sub_combined = np.concatenate([s[s != 0] for s, _ in pairs])
    ref_combined = np.concatenate([r[r != 0] for _, r in pairs])
    return sub_combined, ref_combined


def _sample_selection(
    p_n: int,
    a: np.ndarray,
    b: np.ndarray,
    id_indices: np.ndarray,
    num_sampling_rounds: int = 3,
) -> Tuple[np.ndarray, np.ndarray]:
    """Select representative sample pairs from a single quantization level.

    Args:
        p_n: Number of sample points.
        a: Flattened reference pixels (masked by quantization level; 0 = outside).
        b: Flattened subject pixels (masked by quantization level; 0 = outside).
        id_indices: Random sub-sampling indices.
        num_sampling_rounds: Number of sampling rounds.

    Returns:
        Tuple of (sub_samples, ref_samples), each 1-D.
    """
    a1 = a[a != 0]
    b1 = b[b != 0]

    if len(a1) == 0 or len(b1) == 0:
        return np.array([0.0]), np.array([0.0])

    if len(a1) < p_n or len(b1) < p_n:
        min_len = min(len(a1), len(b1))
        return b1[:min_len], a1[:min_len]

    sub_1, ref_1 = _compute_sample(p_n, a1, b1, id_indices, num_sampling_rounds)

    sub = np.concatenate([sub_1[sub_1 != 0], sub_1[sub_1 == 0]])
    ref = np.concatenate([ref_1[ref_1 != 0], ref_1[ref_1 == 0]])
    return sub, ref


def _linear_reg(
    sub_samples: np.ndarray,
    ref_samples: np.ndarray,
    image_band: np.ndarray,
) -> Tuple[np.ndarray, float, float]:
    """OLS linear regression to normalize one image band.

    Fits ``ref = intercept + slope * sub`` and applies the transformation
    to *image_band*.

    Args:
        sub_samples: Subject sample values (1-D).
        ref_samples: Reference sample values (1-D).
        image_band: Full subject band to normalize (H, W).

    Returns:
        Tuple of (normalized_band, adjusted_r_squared, rmse).

    Raises:
        ValueError: If fewer than 2 valid samples are available for regression.
    """
    mask = (sub_samples != 0) & (ref_samples != 0)
    sub_clean = sub_samples[mask]
    ref_clean = ref_samples[mask]

    if len(sub_clean) < 2:
        raise ValueError(
            f"Insufficient samples for regression: got {len(sub_clean)}, need >= 2"
        )

    X = sub_clean.reshape(-1, 1)
    y = ref_clean

    model = LinearRegression().fit(X, y)
    intercept, slope = model.intercept_, model.coef_[0]

    norm_band = intercept + slope * image_band

    y_pred = model.predict(X)
    ss_res = np.sum((y - y_pred) ** 2)
    ss_tot = np.sum((y - np.mean(y)) ** 2)
    n = len(y)
    r2 = 1 - ss_res / ss_tot if ss_tot != 0 else 0.0
    r_adj = 1 - (1 - r2) * (n - 1) / (n - 2) if n > 2 else 0.0
    rmse = float(np.sqrt(np.mean((y - y_pred) ** 2)))

    return norm_band, r_adj, rmse


def _lirrn(
    p_n: int,
    sub_img: np.ndarray,
    ref_img: np.ndarray,
    num_quantisation_classes: int = 3,
    num_sampling_rounds: int = 3,
    subsample_ratio: float = 0.1,
    rng: Optional[np.random.Generator] = None,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Core LIRRN algorithm.

    Implements Location-Independent Relative Radiometric Normalization.
    Processes each band independently via multi-Otsu thresholding,
    stratified sampling and per-band linear regression.

    Args:
        p_n: Number of sample points per quantization level.
        sub_img: Subject image (H, W, B) float64.
        ref_img: Reference image (H, W, B) float64.
        num_quantisation_classes: Number of brightness strata (default 3).
        num_sampling_rounds: Number of sampling rounds (default 3).
        subsample_ratio: Fraction of candidates retained (default 0.1).
        rng: Numpy random Generator for reproducibility.

    Returns:
        Tuple of (normalized_image, rmse_per_band, r_adj_per_band).
    """
    if rng is None:
        rng = np.random.default_rng()

    num_bands = sub_img.shape[2]

    id_indices = rng.integers(0, p_n, size=max(1, round(subsample_ratio * p_n)))

    norm_img = np.zeros_like(sub_img, dtype=np.float64)
    rmse = np.zeros(num_bands)
    r_adj = np.zeros(num_bands)

    # Quantize each band into brightness levels via multi-Otsu
    sub_labels = np.zeros_like(sub_img, dtype=np.int32)
    ref_labels = np.zeros_like(ref_img, dtype=np.int32)

    for j in range(num_bands):
        for img, labels in [(sub_img, sub_labels), (ref_img, ref_labels)]:
            nonzero = img[:, :, j][img[:, :, j] != 0]
            if len(nonzero) > 0:
                try:
                    thresh = threshold_multiotsu(
                        nonzero, classes=num_quantisation_classes
                    )
                    labels[:, :, j] = np.digitize(img[:, :, j], bins=thresh) + 1
                except ValueError:
                    labels[:, :, j] = 1

    # For each band: sample from quantization levels then regress
    for j in range(num_bands):
        sub_list, ref_list = [], []

        for level in range(1, num_quantisation_classes + 1):
            a = np.where(ref_labels[:, :, j] == level, ref_img[:, :, j], 0).ravel()
            b = np.where(sub_labels[:, :, j] == level, sub_img[:, :, j], 0).ravel()
            sub_s, ref_s = _sample_selection(p_n, a, b, id_indices, num_sampling_rounds)
            sub_list.append(sub_s)
            ref_list.append(ref_s)

        all_sub = np.concatenate(sub_list)
        all_ref = np.concatenate(ref_list)

        try:
            norm_img[:, :, j], r_adj[j], rmse[j] = _linear_reg(
                all_sub, all_ref, sub_img[:, :, j]
            )
        except ValueError:
            warnings.warn(
                f"Band {j}: insufficient samples for regression, "
                "returning band unchanged.",
                stacklevel=2,
            )
            norm_img[:, :, j] = sub_img[:, :, j]
            continue

        ref_band = ref_img[:, :, j]
        ref_min = ref_band.min()
        ref_max = ref_band.max()
        if ref_min == ref_max:
            warnings.warn(
                f"Band {j}: reference band has constant value "
                f"{ref_min}; skipping clipping of normalized band.",
                stacklevel=2,
            )
        else:
            norm_img[:, :, j] = np.clip(norm_img[:, :, j], ref_min, ref_max)

    return norm_img, rmse, r_adj


def normalize_radiometric(
    subject_image: Union[str, np.ndarray],
    reference_image: Union[str, np.ndarray],
    output_path: Optional[str] = None,
    method: str = "lirrn",
    p_n: int = 500,
    num_quantisation_classes: int = 3,
    num_sampling_rounds: int = 3,
    subsample_ratio: float = 0.1,
    random_state: Optional[Union[int, np.random.Generator]] = None,
) -> Tuple[np.ndarray, Dict[str, np.ndarray]]:
    """Normalize subject image radiometry to match a reference image.

    Adjusts brightness and contrast of the subject image so that its pixel
    value distribution matches the reference image. This is essential for
    multi-temporal analysis where images are acquired under different
    atmospheric conditions, sensor calibrations, or illumination angles.

    Currently supports the LIRRN (Location-Independent Relative Radiometric
    Normalization) method, which uses multi-Otsu thresholding and linear
    regression to identify pseudo-invariant features and transform pixel
    values band-by-band.

    Reference: doi:10.3390/s24072272

    Args:
        subject_image: Path to the subject GeoTIFF or numpy array with
            shape (H, W, B). The image to be normalized.
        reference_image: Path to the reference GeoTIFF or numpy array with
            shape (H, W, B). The target radiometry to match.
        output_path: Path to save the normalized image as GeoTIFF. Only
            applicable when *subject_image* is a file path (so spatial
            metadata is available). If None, the array is returned without
            saving. Default: None.
        method: Normalization method. Currently only ``"lirrn"`` is
            supported. Default: ``"lirrn"``.
        p_n: Number of pseudo-invariant feature samples per quantization
            level. Higher values increase accuracy but slow computation.
            Default: 500.
        num_quantisation_classes: Number of brightness strata for stratified
            sampling. Default: 3.
        num_sampling_rounds: Number of iterative refinement rounds for
            sample selection. Default: 3.
        subsample_ratio: Fraction of candidates retained for regression.
            Default: 0.1.
        random_state: Seed or numpy Generator for reproducible results.
            Default: None (non-deterministic).

    Returns:
        Tuple of (normalized_image, metrics) where:
            - normalized_image: numpy array (H, W, B) float64.
            - metrics: dict with keys ``"rmse"`` and ``"r_adj"``, each a
              numpy array of length B.

    Raises:
        ValueError: If *method* is not ``"lirrn"``.
        ValueError: If *p_n* < 1 or *num_sampling_rounds* < 1.
        ValueError: If subject and reference have different band counts.
        ValueError: If input arrays are not 3-dimensional.
        ValueError: If *output_path* is set but *subject_image* is an array.
        FileNotFoundError: If file paths do not point to existing files.

    Examples:
        Normalize a satellite image using file paths:

        >>> from geoai import normalize_radiometric
        >>> norm_img, metrics = normalize_radiometric(
        ...     "subject.tif",
        ...     "reference.tif",
        ...     output_path="normalized.tif",
        ... )
        >>> print(f"RMSE per band: {metrics['rmse']}")

        Normalize using numpy arrays:

        >>> import numpy as np
        >>> subject = np.random.rand(100, 100, 4)
        >>> reference = np.random.rand(120, 120, 4)
        >>> norm_img, metrics = normalize_radiometric(subject, reference)
        >>> norm_img.shape
        (100, 100, 4)

    Note:
        The subject and reference images must have the same number of bands
        but may have different spatial dimensions (height and width).
    """
    # --- Validate parameters ---
    if method != "lirrn":
        raise ValueError(
            f"Unsupported normalization method {method!r}. "
            "Currently only 'lirrn' is supported."
        )
    if p_n < 1:
        raise ValueError(f"p_n must be >= 1, got {p_n}")
    if num_sampling_rounds < 1:
        raise ValueError(f"num_sampling_rounds must be >= 1, got {num_sampling_rounds}")
    if subsample_ratio <= 0 or subsample_ratio > 1:
        raise ValueError(f"subsample_ratio must be in (0, 1], got {subsample_ratio}")

    # --- Resolve inputs ---
    profile = None
    if isinstance(subject_image, str):
        sub_arr, profile = _load_raster(subject_image)
    else:
        sub_arr = np.asarray(subject_image, dtype=np.float64)
        if sub_arr.ndim != 3:
            raise ValueError(
                f"subject_image must be 3-D (H, W, B), got {sub_arr.ndim}-D"
            )

    if isinstance(reference_image, str):
        ref_arr, _ = _load_raster(reference_image)
    else:
        ref_arr = np.asarray(reference_image, dtype=np.float64)
        if ref_arr.ndim != 3:
            raise ValueError(
                f"reference_image must be 3-D (H, W, B), got {ref_arr.ndim}-D"
            )

    if output_path is not None and profile is None:
        raise ValueError(
            "output_path requires subject_image to be a file path "
            "(not an array) so that spatial metadata is available."
        )

    # Band count check
    if sub_arr.shape[2] != ref_arr.shape[2]:
        raise ValueError(
            f"Band count mismatch: subject has {sub_arr.shape[2]} bands, "
            f"reference has {ref_arr.shape[2]} bands."
        )

    # Handle NaN / inf
    if np.any(~np.isfinite(sub_arr)):
        warnings.warn(
            "subject_image contains NaN or infinite values; " "replacing with 0.",
            stacklevel=2,
        )
        sub_arr = np.nan_to_num(sub_arr, nan=0.0, posinf=0.0, neginf=0.0)

    if np.any(~np.isfinite(ref_arr)):
        warnings.warn(
            "reference_image contains NaN or infinite values; " "replacing with 0.",
            stacklevel=2,
        )
        ref_arr = np.nan_to_num(ref_arr, nan=0.0, posinf=0.0, neginf=0.0)

    # --- Build RNG ---
    if isinstance(random_state, np.random.Generator):
        rng = random_state
    else:
        rng = np.random.default_rng(random_state)

    # --- Run normalization ---
    norm_img, rmse, r_adj = _lirrn(
        p_n,
        sub_arr,
        ref_arr,
        num_quantisation_classes=num_quantisation_classes,
        num_sampling_rounds=num_sampling_rounds,
        subsample_ratio=subsample_ratio,
        rng=rng,
    )

    metrics = {"rmse": rmse, "r_adj": r_adj}

    if output_path is not None:
        _save_raster(output_path, norm_img, profile)

    return norm_img, metrics
