"""Radial bounding surfaces for fast terrain intersection pre-filtering.

This module provides a SheetBounds structure that captures the radial extent
(min/avg/max radius from Earth center) and geographic bounds of an XYZ sheet.
These bounds enable fast rejection of rays or observables that cannot possibly
intersect a terrain sheet.

Example::

    from gri_utils import terrain
    import numpy as np

    # Create a terrain sheet
    sheet = terrain.lla_to_xyz_sheet(lats, lons, alts)

    # Get bounds for fast pre-filtering
    bounds = terrain.get_sheet_bounds(sheet)

    # Quick check if a ray could intersect
    origin = np.array([...])
    direction = np.array([...])
    could_hit, approx_point = terrain.ray_intersects_bounds(
        origin, direction, bounds
    )
    if not could_hit:
        return None  # Skip detailed intersection
"""

from __future__ import annotations

from typing import TYPE_CHECKING, NamedTuple

import numpy as np

from gri_utils.conversion import xyz_to_lla

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


# Algorithm constants
_MIN_DIRECTION_NORM = 1e-12
_MIN_RAY_DENOM = 1e-12


class SheetBounds(NamedTuple):
    """Radial bounding information for an XYZ sheet.

    Attributes:
        center_xyz: Centroid of sheet in ECEF meters. Shape (3,).
        center_radius: Distance from Earth center to centroid in meters.
        min_radius: Minimum radius (lowest terrain point) in meters.
        max_radius: Maximum radius (highest terrain point) in meters.
        avg_radius: Average radius across all points in meters.
        lat_bounds: (min_lat, max_lat) in degrees.
        lon_bounds: (min_lon, max_lon) in degrees.
    """

    center_xyz: np.ndarray
    center_radius: float
    min_radius: float
    max_radius: float
    avg_radius: float
    lat_bounds: tuple[float, float]
    lon_bounds: tuple[float, float]


def get_sheet_bounds(
    sheet_xyz: NDArray[np.floating] | Sequence,
) -> SheetBounds:
    """Compute radial bounding surfaces from XYZ sheet.

    Calculates min/avg/max radii from Earth center and geographic bounds
    for fast intersection pre-filtering.

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

    Returns:
        SheetBounds with radial and geographic extent information.
    """
    sheet_xyz = np.asarray(sheet_xyz, dtype=np.float64)
    _m, _n, _ = sheet_xyz.shape  # Shape validation only

    # Compute radii from Earth center for all points
    radii = np.linalg.norm(sheet_xyz, axis=-1)  # Shape (M, N)

    min_radius = float(np.min(radii))
    max_radius = float(np.max(radii))
    avg_radius = float(np.mean(radii))

    # Compute centroid
    center_xyz = np.mean(sheet_xyz.reshape(-1, 3), axis=0)
    center_radius = float(np.linalg.norm(center_xyz))

    # Convert to LLA to get geographic bounds
    xyz_flat = sheet_xyz.reshape(-1, 3)
    lla_flat = xyz_to_lla(xyz_flat)

    lats = lla_flat[:, 0]
    lons = lla_flat[:, 1]

    lat_bounds = (float(np.min(lats)), float(np.max(lats)))
    lon_bounds = (float(np.min(lons)), float(np.max(lons)))

    return SheetBounds(
        center_xyz=center_xyz,
        center_radius=center_radius,
        min_radius=min_radius,
        max_radius=max_radius,
        avg_radius=avg_radius,
        lat_bounds=lat_bounds,
        lon_bounds=lon_bounds,
    )


def ray_intersects_bounds(
    origin_xyz: NDArray[np.floating] | Sequence,
    direction_xyz: NDArray[np.floating] | Sequence,
    bounds: SheetBounds,
    *,
    method: str = "sphere",
) -> tuple[bool, NDArray[np.floating] | None]:
    """Fast check if ray could intersect sheet's radial bounds.

    Args:
        origin_xyz: Ray origin in ECEF meters. Shape (3,).
        direction_xyz: Ray direction vector (will be normalized). Shape (3,).
        bounds: SheetBounds from get_sheet_bounds().
        method: Intersection method:
            - "sphere": Spherical shell intersection (more accurate for large sheets)
            - "plane": Tangent plane at center (faster, good for small sheets)

    Returns:
        Tuple of (could_intersect, approx_point):
        - could_intersect: True if ray passes through the radial band
        - approx_point: Approximate intersection point on avg radius surface,
          or None if no intersection. Shape (3,).
    """
    origin_xyz = np.asarray(origin_xyz, dtype=np.float64)
    direction_xyz = np.asarray(direction_xyz, dtype=np.float64)

    # Normalize direction
    dir_norm = np.linalg.norm(direction_xyz)
    if dir_norm < _MIN_DIRECTION_NORM:
        return False, None
    direction_xyz = direction_xyz / dir_norm

    if method == "plane":
        return _ray_intersects_plane(origin_xyz, direction_xyz, bounds)
    # "sphere"
    return _ray_intersects_sphere(origin_xyz, direction_xyz, bounds)


def _ray_intersects_sphere(
    origin_xyz: NDArray[np.floating],
    direction_xyz: NDArray[np.floating],
    bounds: SheetBounds,
) -> tuple[bool, NDArray[np.floating] | None]:
    """Check ray intersection with spherical shell between min and max radius."""
    # Check intersection with outer sphere (max_radius)
    # Ray-sphere intersection: ||origin + t*dir||^2 = r^2
    # a*t^2 + b*t + c = 0 where:
    # a = dir.dir = 1 (normalized)
    # b = 2 * origin.dir
    # c = origin.origin - r^2

    origin_dot_dir = np.dot(origin_xyz, direction_xyz)
    origin_dot_origin = np.dot(origin_xyz, origin_xyz)

    # Check max sphere
    c_max = origin_dot_origin - bounds.max_radius**2
    discriminant_max = origin_dot_dir**2 - c_max

    if discriminant_max < 0:
        # Ray misses outer sphere entirely
        return False, None

    # Check if ray could be within lat/lon bounds
    # Get approximate intersection point on average sphere
    c_avg = origin_dot_origin - bounds.avg_radius**2
    discriminant_avg = origin_dot_dir**2 - c_avg

    if discriminant_avg < 0:
        # Use max sphere intersection for geographic check
        sqrt_disc = np.sqrt(discriminant_max)
        t1 = -origin_dot_dir - sqrt_disc
        t2 = -origin_dot_dir + sqrt_disc

        # Take the closer positive intersection
        if t1 > 0:
            t = t1
        elif t2 > 0:
            t = t2
        else:
            return False, None

        approx_point = origin_xyz + t * direction_xyz
    else:
        sqrt_disc = np.sqrt(discriminant_avg)
        t1 = -origin_dot_dir - sqrt_disc
        t2 = -origin_dot_dir + sqrt_disc

        if t1 > 0:
            t = t1
        elif t2 > 0:
            t = t2
        else:
            # Both intersections behind us - check if we're inside
            # and heading toward the surface
            origin_radius = np.sqrt(origin_dot_origin)
            if bounds.min_radius < origin_radius < bounds.max_radius:
                # We're inside the shell, find where we exit
                c_min = origin_dot_origin - bounds.min_radius**2
                disc_min = origin_dot_dir**2 - c_min
                if disc_min >= 0:
                    t = -origin_dot_dir + np.sqrt(disc_min)
                    if t > 0:
                        approx_point = origin_xyz + t * direction_xyz
                        return _check_geographic_bounds(approx_point, bounds)
            return False, None

        approx_point = origin_xyz + t * direction_xyz

    return _check_geographic_bounds(approx_point, bounds)


def _ray_intersects_plane(
    origin_xyz: NDArray[np.floating],
    direction_xyz: NDArray[np.floating],
    bounds: SheetBounds,
) -> tuple[bool, NDArray[np.floating] | None]:
    """Check ray intersection with tangent plane at sheet center."""
    # Plane normal is the radial direction at center
    center = bounds.center_xyz
    normal = center / bounds.center_radius

    # Ray-plane intersection: t = (center - origin).normal / (dir.normal)
    denom = np.dot(direction_xyz, normal)

    if abs(denom) < _MIN_RAY_DENOM:
        # Ray parallel to plane
        return False, None

    t = np.dot(center - origin_xyz, normal) / denom

    if t < 0:
        # Intersection behind ray origin
        return False, None

    approx_point = origin_xyz + t * direction_xyz

    # Check if point is within radial bounds
    point_radius = np.linalg.norm(approx_point)
    if not (bounds.min_radius <= point_radius <= bounds.max_radius):
        return False, None

    return _check_geographic_bounds(approx_point, bounds)


def _check_geographic_bounds(
    point_xyz: NDArray[np.floating],
    bounds: SheetBounds,
) -> tuple[bool, NDArray[np.floating] | None]:
    """Check if point is within geographic bounds of sheet."""
    lla = xyz_to_lla(point_xyz)
    lat, lon = float(lla[0]), float(lla[1])

    min_lat, max_lat = bounds.lat_bounds
    min_lon, max_lon = bounds.lon_bounds

    # Add small margin for numerical tolerance
    margin = 0.01  # degrees

    if not (min_lat - margin <= lat <= max_lat + margin):
        return False, None

    # Handle longitude wraparound
    if min_lon <= max_lon:
        # Normal case
        if not (min_lon - margin <= lon <= max_lon + margin):
            return False, None
    # Wraps around 180/-180
    elif not (lon >= min_lon - margin or lon <= max_lon + margin):
        return False, None

    return True, point_xyz


def point_within_bounds(
    point_xyz: NDArray[np.floating] | Sequence,
    bounds: SheetBounds,
) -> bool:
    """Check if point is within sheet's radial and geographic bounds.

    Args:
        point_xyz: Point in ECEF meters. Shape (3,).
        bounds: SheetBounds from get_sheet_bounds().

    Returns:
        True if point's radius is between min_radius and max_radius AND
        point's lat/lon is within lat_bounds/lon_bounds.
    """
    point_xyz = np.asarray(point_xyz, dtype=np.float64)

    # Check radial bounds
    radius = float(np.linalg.norm(point_xyz))
    if not (bounds.min_radius <= radius <= bounds.max_radius):
        return False

    # Check geographic bounds
    result, _ = _check_geographic_bounds(point_xyz, bounds)
    return result


def bounds_overlap(
    bounds1: SheetBounds,
    bounds2: SheetBounds,
) -> bool:
    """Check if two sheets' bounds could overlap.

    Useful for sheet-sheet intersection pre-filtering.

    Args:
        bounds1: First SheetBounds.
        bounds2: Second SheetBounds.

    Returns:
        True if the sheets' radial and geographic bounds overlap.
    """
    # Check radial overlap
    if bounds1.max_radius < bounds2.min_radius:
        return False
    if bounds2.max_radius < bounds1.min_radius:
        return False

    # Check latitude overlap
    if bounds1.lat_bounds[1] < bounds2.lat_bounds[0]:
        return False
    if bounds2.lat_bounds[1] < bounds1.lat_bounds[0]:
        return False

    # Check longitude overlap (handle wraparound)
    lon1_min, lon1_max = bounds1.lon_bounds
    lon2_min, lon2_max = bounds2.lon_bounds

    # Simple case: neither wraps around
    if (  # noqa: SIM102
        lon1_min <= lon1_max and lon2_min <= lon2_max
    ):
        if lon1_max < lon2_min or lon2_max < lon1_min:
            return False
    # One or both wrap around - conservatively return True
    # (proper handling of all wraparound cases is complex)

    return True
