"""Sheet stitching utilities for mosaicking multiple terrain tiles.

When terrain data spans multiple tiles (e.g., multiple 1x1 degree DTED files),
this module provides functions to combine them into a single sheet.

Example::

    from gri_utils import terrain

    # Get tiles from source
    sheets_with_bounds = source.get_covering_sheets(lats, lons)

    # Stitch into single sheet
    combined_sheet, combined_bounds = terrain.stitch_sheets(sheets_with_bounds)
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

from .source import TerrainBounds

_RESOLUTION_TOLERANCE = 0.01  # 1% tolerance for resolution matching
_MIN_GRID_POINTS = 2  # Minimum points per dimension

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


def stitch_sheets(
    sheets_with_bounds: Sequence[tuple[NDArray[np.floating], TerrainBounds]],
) -> tuple[NDArray[np.floating], TerrainBounds]:
    """Mosaic multiple terrain sheets into a single sheet.

    Combines multiple terrain tiles into a single contiguous array. Tiles
    are placed according to their geographic bounds, with gaps filled by NaN.

    Args:
        sheets_with_bounds: List of (sheet, bounds) tuples. Each sheet must
            be an XYZ array with shape (M, N, 3). All sheets must have
            compatible resolution (similar grid spacing).

    Returns:
        Tuple of (combined_sheet, combined_bounds) where:
        - combined_sheet: XYZ array (M, N, 3) covering all input tiles
        - combined_bounds: TerrainBounds for the combined sheet

    Raises:
        ValueError: If sheets_with_bounds is empty or sheets have incompatible
            resolutions.

    Notes:
        - Sheets must have compatible resolution (grid spacing within 1%)
        - Overlapping regions use data from the first sheet in the list
        - Gaps between tiles are filled with NaN
        - Single-sheet input is returned directly (no copy)

    Example:
        Stitch two adjacent DTED tiles::

            sheets = [
                (sheet1, TerrainBounds(38.0, 39.0, -106.0, -105.0)),
                (sheet2, TerrainBounds(38.0, 39.0, -105.0, -104.0)),
            ]
            combined, bounds = stitch_sheets(sheets)
            # bounds = TerrainBounds(38.0, 39.0, -106.0, -104.0)
    """
    if not sheets_with_bounds:
        msg = "sheets_with_bounds cannot be empty"
        raise ValueError(msg)

    # Single sheet - return directly
    if len(sheets_with_bounds) == 1:
        return sheets_with_bounds[0]

    # Compute combined bounds
    all_bounds = [b for _, b in sheets_with_bounds]
    combined_bounds = TerrainBounds(
        lat_min=min(b.lat_min for b in all_bounds),
        lat_max=max(b.lat_max for b in all_bounds),
        lon_min=min(b.lon_min for b in all_bounds),
        lon_max=max(b.lon_max for b in all_bounds),
    )

    # Compute resolution from first sheet
    first_sheet, first_bounds = sheets_with_bounds[0]
    m, n, _ = first_sheet.shape
    lat_res = (first_bounds.lat_max - first_bounds.lat_min) / (m - 1) if m > 1 else 0.0
    lon_res = (first_bounds.lon_max - first_bounds.lon_min) / (n - 1) if n > 1 else 0.0

    if lat_res == 0.0 or lon_res == 0.0:
        msg = "Sheets must have at least 2 points in each dimension"
        raise ValueError(msg)

    # Verify compatible resolution across all sheets
    for sheet, bounds in sheets_with_bounds[1:]:
        sm, sn, _ = sheet.shape
        s_lat_res = (bounds.lat_max - bounds.lat_min) / (sm - 1) if sm > 1 else 0.0
        s_lon_res = (bounds.lon_max - bounds.lon_min) / (sn - 1) if sn > 1 else 0.0

        if abs(s_lat_res - lat_res) / lat_res > _RESOLUTION_TOLERANCE:
            msg = f"Incompatible latitude resolution: {s_lat_res} vs {lat_res}"
            raise ValueError(msg)
        if abs(s_lon_res - lon_res) / lon_res > _RESOLUTION_TOLERANCE:
            msg = f"Incompatible longitude resolution: {s_lon_res} vs {lon_res}"
            raise ValueError(msg)

    # Compute output grid size
    total_lat_span = combined_bounds.lat_max - combined_bounds.lat_min
    total_lon_span = combined_bounds.lon_max - combined_bounds.lon_min
    out_m = round(total_lat_span / lat_res) + 1
    out_n = round(total_lon_span / lon_res) + 1

    # Initialize output with NaN
    combined_sheet = np.full((out_m, out_n, 3), np.nan, dtype=np.float64)

    # Place each sheet into the combined array
    for sheet, bounds in sheets_with_bounds:
        # Compute where this sheet goes in the output
        row_start = round((bounds.lat_min - combined_bounds.lat_min) / lat_res)
        col_start = round((bounds.lon_min - combined_bounds.lon_min) / lon_res)

        sm, sn, _ = sheet.shape
        row_end = row_start + sm
        col_end = col_start + sn

        # Clip to output bounds (shouldn't be needed, but safety)
        row_start = max(0, row_start)
        col_start = max(0, col_start)
        row_end = min(out_m, row_end)
        col_end = min(out_n, col_end)

        # Compute corresponding slice in source sheet
        src_row_start = max(0, -row_start)
        src_col_start = max(0, -col_start)
        src_row_end = sm - max(0, row_start + sm - out_m)
        src_col_end = sn - max(0, col_start + sn - out_n)

        # Only copy where output is NaN (first sheet wins in overlaps)
        target_slice = combined_sheet[row_start:row_end, col_start:col_end]
        source_slice = sheet[src_row_start:src_row_end, src_col_start:src_col_end]

        # Create mask where target is NaN (any component)
        nan_mask = np.isnan(target_slice[:, :, 0])

        # Copy source data where target is NaN
        target_slice[nan_mask] = source_slice[nan_mask]

    return combined_sheet, combined_bounds


def get_stitch_resolution(
    sheets_with_bounds: Sequence[tuple[NDArray[np.floating], TerrainBounds]],
) -> tuple[float, float]:
    """Get the grid resolution of sheets for stitching.

    Args:
        sheets_with_bounds: List of (sheet, bounds) tuples.

    Returns:
        Tuple of (lat_resolution, lon_resolution) in degrees.

    Raises:
        ValueError: If sheets_with_bounds is empty or first sheet is too small.
    """
    if not sheets_with_bounds:
        msg = "sheets_with_bounds cannot be empty"
        raise ValueError(msg)

    first_sheet, first_bounds = sheets_with_bounds[0]
    m, n, _ = first_sheet.shape

    if m < _MIN_GRID_POINTS or n < _MIN_GRID_POINTS:
        msg = "Sheets must have at least 2 points in each dimension"
        raise ValueError(msg)

    lat_res = (first_bounds.lat_max - first_bounds.lat_min) / (m - 1)
    lon_res = (first_bounds.lon_max - first_bounds.lon_min) / (n - 1)

    return lat_res, lon_res
