"""Terrain sheet operations for XYZ surface grids.

A sheet is a 2D grid of XYZ points representing a terrain surface in ECEF
coordinates. This module provides functions for converting between LLA
heightmaps and XYZ sheets, interpolation, and resolution/bounds utilities.

Example::

    from gri_utils import terrain
    import numpy as np

    # Create a simple 10x10 grid over a small area
    lats = np.linspace(39.0, 39.1, 10)
    lons = np.linspace(-105.0, -104.9, 10)
    alts = np.zeros((10, 10))  # Flat terrain at sea level

    # Convert to XYZ sheet
    sheet = terrain.lla_to_xyz_sheet(lats, lons, alts)

    # Get resolution (average spacing in meters)
    row_res, col_res = terrain.get_sheet_resolution(sheet)
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np
from scipy.interpolate import RegularGridInterpolator

from gri_utils.conversion import lla_to_xyz, xyz_to_lla

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


def lla_to_xyz_sheet(
    lats: NDArray[np.floating] | Sequence,
    lons: NDArray[np.floating] | Sequence,
    altitudes: NDArray[np.floating] | Sequence,
) -> NDArray[np.floating]:
    """Convert LLA heightmap grid to XYZ sheet.

    Args:
        lats: Latitude values in degrees. Shape (M,).
        lons: Longitude values in degrees. Shape (N,).
        altitudes: Altitude values in meters at each grid point. Shape (M, N).

    Returns:
        XYZ sheet in ECEF meters. Shape (M, N, 3).
    """
    lats = np.asarray(lats, dtype=np.float64)
    lons = np.asarray(lons, dtype=np.float64)
    altitudes = np.asarray(altitudes, dtype=np.float64)

    # Create meshgrid of lat/lon
    lon_grid, lat_grid = np.meshgrid(lons, lats)

    # Stack into (M, N, 3) LLA array
    lla_grid = np.stack([lat_grid, lon_grid, altitudes], axis=-1)

    # Reshape to (M*N, 3) for conversion, then reshape back
    m, n = altitudes.shape
    lla_flat = lla_grid.reshape(-1, 3)
    xyz_flat = lla_to_xyz(lla_flat)

    return xyz_flat.reshape(m, n, 3)


def xyz_to_lla_sheet(
    sheet_xyz: NDArray[np.floating] | Sequence,
) -> tuple[NDArray[np.floating], NDArray[np.floating], NDArray[np.floating]]:
    """Convert XYZ sheet back to LLA components.

    Note: This extracts lats/lons from the first row/column. For sheets
    created from regular LLA grids, this recovers the original coordinates.
    For irregular sheets, this provides approximate values.

    Args:
        sheet_xyz: XYZ sheet in ECEF meters. Shape (M, N, 3).

    Returns:
        Tuple of (lats, lons, altitudes) where:
        - lats: Latitude values in degrees. Shape (M,).
        - lons: Longitude values in degrees. Shape (N,).
        - altitudes: Altitude values in meters. Shape (M, N).
    """
    sheet_xyz = np.asarray(sheet_xyz, dtype=np.float64)
    m, n, _ = sheet_xyz.shape

    # Convert all points to LLA
    xyz_flat = sheet_xyz.reshape(-1, 3)
    lla_flat = xyz_to_lla(xyz_flat)
    lla_grid = lla_flat.reshape(m, n, 3)

    # Extract lats from first column, lons from first row
    lats = lla_grid[:, 0, 0]
    lons = lla_grid[0, :, 1]
    altitudes = lla_grid[:, :, 2]

    return lats, lons, altitudes


def interpolate_sheet(
    sheet_xyz: NDArray[np.floating] | Sequence,
    row_frac: float,
    col_frac: float,
    *,
    order: int = 3,
) -> NDArray[np.floating]:
    """Interpolate a point on the XYZ sheet.

    Uses scipy's RegularGridInterpolator for efficient 3D vector interpolation.

    Args:
        sheet_xyz: XYZ sheet in ECEF meters. Shape (M, N, 3).
        row_frac: Fractional row index in range [0, M-1].
        col_frac: Fractional column index in range [0, N-1].
        order: Interpolation order. 0=nearest, 1=linear, 3=cubic (default).
            Note: Cubic requires at least 4 points per dimension; falls back
            to linear for smaller grids.

    Returns:
        Interpolated XYZ point in ECEF meters. Shape (3,).
    """
    sheet_xyz = np.asarray(sheet_xyz, dtype=np.float64)
    m, n, _ = sheet_xyz.shape

    # Create grid coordinate arrays
    rows = np.arange(m, dtype=np.float64)
    cols = np.arange(n, dtype=np.float64)

    # Map order to scipy method, with fallback for small grids
    if order == 0:
        method = "nearest"
    elif order == 1 or m < 4 or n < 4:  # noqa: PLR2004
        # Cubic requires at least 4 points per dimension
        method = "linear"
    else:
        method = "cubic"

    # Create single interpolator for all three XYZ coordinates
    interp = RegularGridInterpolator((rows, cols), sheet_xyz, method=method)

    # Interpolate at the query point
    point = np.array([[row_frac, col_frac]])
    result = interp(point)[0]

    return np.asarray(result, dtype=np.float64)


def interpolate_sheet_batch(
    sheet_xyz: NDArray[np.floating] | Sequence,
    row_fracs: NDArray[np.floating] | Sequence,
    col_fracs: NDArray[np.floating] | Sequence,
    *,
    order: int = 3,
) -> NDArray[np.floating]:
    """Interpolate multiple points on the XYZ sheet.

    Args:
        sheet_xyz: XYZ sheet in ECEF meters. Shape (M, N, 3).
        row_fracs: Fractional row indices. Shape (K,).
        col_fracs: Fractional column indices. Shape (K,).
        order: Interpolation order. 0=nearest, 1=linear, 3=cubic (default).
            Note: Cubic requires at least 4 points per dimension; falls back
            to linear for smaller grids.

    Returns:
        Interpolated XYZ points in ECEF meters. Shape (K, 3).
    """
    sheet_xyz = np.asarray(sheet_xyz, dtype=np.float64)
    row_fracs = np.asarray(row_fracs, dtype=np.float64)
    col_fracs = np.asarray(col_fracs, dtype=np.float64)
    m, n, _ = sheet_xyz.shape

    # Create grid coordinate arrays
    rows = np.arange(m, dtype=np.float64)
    cols = np.arange(n, dtype=np.float64)

    # Map order to scipy method, with fallback for small grids
    if order == 0:
        method = "nearest"
    elif order == 1 or m < 4 or n < 4:  # noqa: PLR2004
        # Cubic requires at least 4 points per dimension
        method = "linear"
    else:
        method = "cubic"

    # Create single interpolator for all three XYZ coordinates
    interp = RegularGridInterpolator((rows, cols), sheet_xyz, method=method)

    # Stack query points as (K, 2) array and interpolate all at once
    points = np.column_stack([row_fracs, col_fracs])
    result = interp(points)

    return np.asarray(result, dtype=np.float64)


def get_sheet_resolution(
    sheet_xyz: NDArray[np.floating] | Sequence,
) -> tuple[float, float]:
    """Compute average grid spacing along each axis in meters.

    Args:
        sheet_xyz: XYZ sheet in ECEF meters. Shape (M, N, 3).

    Returns:
        Tuple of (row_spacing_m, col_spacing_m) representing average
        distance between adjacent points along each axis.
    """
    sheet_xyz = np.asarray(sheet_xyz, dtype=np.float64)

    # Compute distances between adjacent points along rows (axis 0)
    row_diffs = np.diff(sheet_xyz, axis=0)  # Shape (M-1, N, 3)
    row_distances = np.linalg.norm(row_diffs, axis=-1)  # Shape (M-1, N)
    row_spacing = float(np.mean(row_distances))

    # Compute distances between adjacent points along columns (axis 1)
    col_diffs = np.diff(sheet_xyz, axis=1)  # Shape (M, N-1, 3)
    col_distances = np.linalg.norm(col_diffs, axis=-1)  # Shape (M, N-1)
    col_spacing = float(np.mean(col_distances))

    return row_spacing, col_spacing
