"""Ray-terrain sheet intersection algorithm.

Finds the first intersection point between a ray and a terrain sheet using
adaptive step sizing and binary search refinement.

Example::

    from gri_utils import terrain
    import numpy as np

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

    # Define ray from satellite looking down
    origin = np.array([0, 0, 7_000_000])  # 7000 km above Earth
    direction = np.array([0, 0, -1])  # Straight down

    # Find intersection
    hit_point = terrain.ray_sheet_intersect(sheet, origin, direction)
    if hit_point is not None:
        print(f"Hit at {hit_point}")
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

from gri_utils.conversion import xyz_to_lla

from .sheet import interpolate_sheet
from .sheet_bounds import SheetBounds, get_sheet_bounds, ray_intersects_bounds

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


# Algorithm constants
_MIN_DIRECTION_NORM = 1e-12
_MIN_STEP_M = 1.0  # Minimum step size
_MAX_STEP_M = 500.0  # Maximum step size
_REFINEMENT_TOLERANCE_M = 1.0  # Binary search tolerance
_MAX_TERRAIN_SLOPE = 1.0  # tan(45 degrees), max assumed terrain slope
_MAX_ITERATIONS = 100_000
_MIN_ALT_THRESHOLD = -1000.0  # Well below sea level (meters)
_MAX_ALT_THRESHOLD = 20000.0  # Above Everest (meters)


def ray_sheet_intersect(
    sheet_xyz: NDArray[np.floating] | Sequence,
    origin_xyz: NDArray[np.floating] | Sequence,
    direction_xyz: NDArray[np.floating] | Sequence,
    *,
    altitude_epsilon_m: float = 0.0,
    order: int = 1,
) -> NDArray[np.floating] | None:
    """Find first intersection of ray with terrain sheet.

    Uses adaptive step sizing based on clearance above terrain. Step size
    is proportional to clearance, ensuring terrain rising at 45 degrees
    cannot be "stepped over".

    Args:
        sheet_xyz: Terrain sheet in ECEF meters. Shape (M, N, 3).
        origin_xyz: Ray origin in ECEF meters. Shape (3,).
        direction_xyz: Ray direction vector (will be normalized). Shape (3,).
        altitude_epsilon_m: Offset above terrain surface to find intersection.
            0.0 = bare terrain (default).
            10.0 = find where ray passes within 10m of terrain.
        order: Interpolation order. 0=nearest, 1=bilinear (default), 3=bicubic.

    Returns:
        ECEF coordinates of first intersection point (meters), shape (3,),
        or None if no intersection found.
    """
    sheet_xyz = np.asarray(sheet_xyz, dtype=np.float64)
    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 None
    direction_xyz = direction_xyz / dir_norm

    # Get sheet bounds for fast rejection
    bounds = get_sheet_bounds(sheet_xyz)

    # Quick check if ray could intersect
    could_hit, _approx_point = ray_intersects_bounds(
        origin_xyz,
        direction_xyz,
        bounds,
        method="sphere",
    )
    if not could_hit:
        return None

    # Get LLA bounds for the sheet
    min_lat, max_lat = bounds.lat_bounds
    min_lon, max_lon = bounds.lon_bounds

    # Compute starting distance (skip to where ray enters terrain altitude band)
    start_dist = _get_start_dist(origin_xyz, direction_xyz, bounds)
    if start_dist is None:
        return None

    # Adaptive ray march
    return _march(
        sheet_xyz,
        origin_xyz,
        direction_xyz,
        start_dist,
        altitude_epsilon_m,
        min_lat,
        max_lat,
        min_lon,
        max_lon,
        order,
    )


def _get_start_dist(
    origin_xyz: NDArray[np.floating],
    direction_xyz: NDArray[np.floating],
    bounds: SheetBounds,
) -> float | None:
    """Compute starting distance along ray for terrain intersection search."""
    origin_radius = float(np.linalg.norm(origin_xyz))

    # If origin is above max radius, skip to where we enter the altitude band
    if origin_radius > bounds.max_radius:
        # Check if descending
        radial = origin_xyz / origin_radius
        if np.dot(direction_xyz, radial) >= 0:
            return None  # Ascending, won't hit terrain

        # Compute distance to max radius sphere
        # ||origin + t*dir||^2 = r^2
        # t^2 + 2*t*(origin.dir) + origin.origin - r^2 = 0
        origin_dot_dir = np.dot(origin_xyz, direction_xyz)
        origin_dot_origin = np.dot(origin_xyz, origin_xyz)
        c = origin_dot_origin - bounds.max_radius**2
        discriminant = origin_dot_dir**2 - c

        if discriminant < 0:
            return None

        t = -origin_dot_dir - np.sqrt(discriminant)
        return max(0.0, t)

    # If origin is below min radius, check if ascending
    if origin_radius < bounds.min_radius:
        radial = origin_xyz / origin_radius
        if np.dot(direction_xyz, radial) <= 0:
            return None  # Descending, moving away from terrain

    return 0.0


def _march(  # noqa: PLR0913
    sheet_xyz: NDArray[np.floating],
    origin_xyz: NDArray[np.floating],
    direction_xyz: NDArray[np.floating],
    start_dist: float,
    altitude_epsilon_m: float,
    min_lat: float,
    max_lat: float,
    min_lon: float,
    max_lon: float,
    order: int,
) -> NDArray[np.floating] | None:
    """Ray marching with slope-based adaptive stepping."""
    curr_dist = start_dist
    m, n, _ = sheet_xyz.shape

    # Get initial position
    curr_point = origin_xyz + direction_xyz * curr_dist
    curr_lla = xyz_to_lla(curr_point)
    curr_alt = float(curr_lla[2])

    # Get terrain altitude at current position
    terrain_alt = _get_sheet_altitude_at_lla(
        sheet_xyz,
        float(curr_lla[0]),
        float(curr_lla[1]),
        min_lat,
        max_lat,
        min_lon,
        max_lon,
        m,
        n,
        order,
    )
    if terrain_alt is None:
        # Start outside sheet bounds, try to find entry point
        terrain_alt = 0.0  # Assume sea level outside sheet

    clearance = curr_alt - (terrain_alt + altitude_epsilon_m)
    prev_above = clearance > 0
    prev_dist = curr_dist
    prev_alt = curr_alt

    for _ in range(_MAX_ITERATIONS):
        lat, lon = float(curr_lla[0]), float(curr_lla[1])

        # Check if we've left the geographic bounds
        if lat < min_lat - 0.1 or lat > max_lat + 0.1:
            return None
        if lon < min_lon - 0.1 or lon > max_lon + 0.1:
            return None

        # Get terrain altitude
        terrain_alt = _get_sheet_altitude_at_lla(
            sheet_xyz,
            lat,
            lon,
            min_lat,
            max_lat,
            min_lon,
            max_lon,
            m,
            n,
            order,
        )
        if terrain_alt is None:
            # Outside sheet, use large step
            step = _MAX_STEP_M
        else:
            clearance = curr_alt - (terrain_alt + altitude_epsilon_m)
            curr_above = clearance > 0

            # Check for crossing
            if curr_dist > prev_dist and curr_above != prev_above:
                return _refine(
                    sheet_xyz,
                    origin_xyz,
                    direction_xyz,
                    prev_dist,
                    curr_dist,
                    altitude_epsilon_m,
                    min_lat,
                    max_lat,
                    min_lon,
                    max_lon,
                    m,
                    n,
                    order,
                    low_above=prev_above,
                )

            prev_above = curr_above

            # Compute step based on clearance
            step = _slope_skip_dist(clearance)

        # Advance
        prev_dist = curr_dist
        prev_alt = curr_alt
        curr_dist += step

        curr_point = origin_xyz + direction_xyz * curr_dist
        curr_lla = xyz_to_lla(curr_point)
        curr_alt = float(curr_lla[2])

        # Check if we've gone too far (below min terrain or above max)
        if curr_alt < _MIN_ALT_THRESHOLD:
            return None
        if curr_alt > _MAX_ALT_THRESHOLD and curr_alt > prev_alt:
            return None

    return None  # Max iterations exceeded


def _slope_skip_dist(clearance: float) -> float:
    """Compute safe skip distance based on 45-degree max terrain slope."""
    if clearance <= 0:
        return _MIN_STEP_M

    # Skip distance = clearance / max_slope (tan 45 = 1)
    step = clearance / _MAX_TERRAIN_SLOPE
    return float(np.clip(step, _MIN_STEP_M, _MAX_STEP_M))


def _get_sheet_altitude_at_lla(  # noqa: PLR0913
    sheet_xyz: NDArray[np.floating],
    lat: float,
    lon: float,
    min_lat: float,
    max_lat: float,
    min_lon: float,
    max_lon: float,
    m: int,
    n: int,
    order: int,
) -> float | None:
    """Get interpolated altitude from sheet at given lat/lon.

    Returns None if lat/lon is outside sheet bounds.
    """
    # Check bounds
    if lat < min_lat or lat > max_lat:
        return None
    if lon < min_lon or lon > max_lon:
        return None

    # Convert lat/lon to fractional grid indices
    row_frac = (lat - min_lat) / (max_lat - min_lat) * (m - 1)
    col_frac = (lon - min_lon) / (max_lon - min_lon) * (n - 1)

    # Clamp to valid range
    row_frac = float(np.clip(row_frac, 0, m - 1))
    col_frac = float(np.clip(col_frac, 0, n - 1))

    # Interpolate to get XYZ, then extract altitude
    point_xyz = interpolate_sheet(sheet_xyz, row_frac, col_frac, order=order)
    point_lla = xyz_to_lla(point_xyz)
    return float(point_lla[2])


def _refine(  # noqa: PLR0913
    sheet_xyz: NDArray[np.floating],
    origin_xyz: NDArray[np.floating],
    direction_xyz: NDArray[np.floating],
    low_dist: float,
    high_dist: float,
    altitude_epsilon_m: float,
    min_lat: float,
    max_lat: float,
    min_lon: float,
    max_lon: float,
    m: int,
    n: int,
    order: int,
    *,
    low_above: bool,
) -> NDArray[np.floating]:
    """Refine intersection using binary search with configurable interpolation."""
    while high_dist - low_dist > _REFINEMENT_TOLERANCE_M:
        mid_dist = (low_dist + high_dist) / 2
        mid_point = origin_xyz + direction_xyz * mid_dist
        mid_lla = xyz_to_lla(mid_point)

        terrain_alt = _get_sheet_altitude_at_lla(
            sheet_xyz,
            float(mid_lla[0]),
            float(mid_lla[1]),
            min_lat,
            max_lat,
            min_lon,
            max_lon,
            m,
            n,
            order,
        )
        if terrain_alt is None:
            terrain_alt = 0.0

        mid_above = mid_lla[2] > terrain_alt + altitude_epsilon_m

        if mid_above == low_above:
            low_dist = mid_dist
        else:
            high_dist = mid_dist

    final_dist = (low_dist + high_dist) / 2
    return origin_xyz + direction_xyz * final_dist


def ray_intersects_sheet_bounds(
    sheet_xyz: NDArray[np.floating] | Sequence,
    origin_xyz: NDArray[np.floating] | Sequence,
    direction_xyz: NDArray[np.floating] | Sequence,
) -> bool:
    """Quick check if ray could intersect sheet's bounding box.

    This is a convenience wrapper around get_sheet_bounds and ray_intersects_bounds.

    Args:
        sheet_xyz: Terrain sheet in ECEF meters. Shape (M, N, 3).
        origin_xyz: Ray origin in ECEF meters. Shape (3,).
        direction_xyz: Ray direction vector. Shape (3,).

    Returns:
        True if ray could potentially intersect the sheet.
    """
    sheet_xyz = np.asarray(sheet_xyz, dtype=np.float64)
    origin_xyz = np.asarray(origin_xyz, dtype=np.float64)
    direction_xyz = np.asarray(direction_xyz, dtype=np.float64)

    bounds = get_sheet_bounds(sheet_xyz)
    could_hit, _ = ray_intersects_bounds(origin_xyz, direction_xyz, bounds)
    return could_hit
