"""Shadow depth computation for terrain visibility analysis.

Computes how deep each terrain point is in shadow (or how much clearance
it has above the shadow boundary) relative to an observer. This replaces
binary visibility masks with a continuous scalar field.

The shadow depth is measured in meters along the observer's local vertical
at each terrain point:

- Positive values: point is visible with this much vertical clearance
- Negative values: point is in shadow, needs a mast of abs(depth) meters
- Zero: point lies on the shadow boundary

Algorithm: for each grid point, trace a ray in grid space back toward the
observer using DDA (digital differential analyzer) traversal. At each grid
cell the ray crosses, linearly interpolate the terrain elevation angle
between the bounding grid points. The maximum interpolated elevation angle
along the ray determines whether the target is shadowed and by how much.

Complexity: O(M * N * max(M, N)). No coordinate conversions in the inner
loop -- all geometry is pre-computed in a single vectorized pass.

Example::

    from gri_utils import terrain
    import numpy as np

    sheet = terrain.lla_to_xyz_sheet(lats, lons, alts)
    observer = np.array([x, y, z])  # ECEF meters

    # Get shadow depth field
    depth = terrain.get_shadow_depth(sheet, observer)

    # Visibility mask is just the sign
    visible = depth > 0

    # Find shadow boundaries
    contours = terrain.find_contours(depth, level=0.0)
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray

# Grid-space thresholds for DDA ray tracing
_GRID_COINCIDENCE_TOL = 0.01  # Minimum grid distance to consider distinct points
_GRID_DIRECTION_TOL = 1e-10  # Minimum direction component to traverse an axis
_RAY_END_FRACTION = 0.99  # Skip crossings this close to the observer (t >= 0.99)


def get_shadow_depth(
    sheet_xyz: NDArray[np.floating] | Sequence,
    observer_xyz: NDArray[np.floating] | Sequence,
) -> NDArray[np.floating]:
    """Compute shadow depth for each terrain point relative to observer.

    Positive values indicate the point is visible with vertical clearance
    (meters along local up). Negative values indicate the point is in
    shadow and needs a mast of abs(depth) meters to gain line of sight.
    Zero values lie on the shadow boundary.

    Args:
        sheet_xyz: Terrain sheet in ECEF meters. Shape (M, N, 3).
        observer_xyz: Observer position in ECEF meters. Shape (3,).

    Returns:
        Shadow depth in meters. Shape (M, N). Positive = visible with
        clearance, negative = shadowed, zero = boundary.
    """
    sheet_xyz = np.asarray(sheet_xyz, dtype=np.float64)
    observer_xyz = np.asarray(observer_xyz, dtype=np.float64)
    # Observer's local up direction (geocentric approximation)
    observer_up = observer_xyz / np.linalg.norm(observer_xyz)

    # Vector from observer to each grid point
    vec = sheet_xyz - observer_xyz  # (M, N, 3)

    # Height above observer's horizontal plane
    h_point = np.sum(vec * observer_up, axis=-1)  # (M, N)

    # Horizontal distance from observer
    dist_sq = np.sum(vec * vec, axis=-1)  # (M, N)
    r_sq = dist_sq - h_point**2
    r_point = np.sqrt(np.maximum(r_sq, 0.0))  # (M, N)

    # Elevation angle from observer to each grid point
    elev = np.arctan2(h_point, r_point)  # (M, N)

    # Observer's fractional grid position
    obs_row, obs_col = _observer_grid_position(observer_xyz, sheet_xyz)

    # For each target, trace ray back to observer and find max blocking
    # elevation along the way
    return _ray_march_all(elev, h_point, r_point, obs_row, obs_col)


def _observer_grid_position(
    observer_xyz: NDArray[np.floating],
    sheet_xyz: NDArray[np.floating],
) -> tuple[float, float]:
    """Convert observer ECEF position to fractional grid (row, col).

    Projects the observer onto the grid coordinate system using the
    sheet's geographic bounds.

    Args:
        observer_xyz: Observer position in ECEF meters. Shape (3,).
        sheet_xyz: Terrain sheet in ECEF meters. Shape (M, N, 3).

    Returns:
        Tuple of (row_frac, col_frac). May be outside [0, M-1] x [0, N-1]
        if the observer is outside the grid bounds.
    """
    from gri_utils.conversion import xyz_to_lla  # noqa: PLC0415

    from .sheet_bounds import get_sheet_bounds  # noqa: PLC0415

    m, n, _ = sheet_xyz.shape
    bounds = get_sheet_bounds(sheet_xyz)
    min_lat, max_lat = bounds.lat_bounds
    min_lon, max_lon = bounds.lon_bounds

    observer_lla = xyz_to_lla(observer_xyz)
    obs_lat = float(observer_lla[0])
    obs_lon = float(observer_lla[1])

    row_frac = (obs_lat - min_lat) / (max_lat - min_lat) * (m - 1)
    col_frac = (obs_lon - min_lon) / (max_lon - min_lon) * (n - 1)

    return float(row_frac), float(col_frac)


def _ray_march_all(
    elev: NDArray[np.floating],
    h_point: NDArray[np.floating],
    r_point: NDArray[np.floating],
    obs_row: float,
    obs_col: float,
) -> NDArray[np.floating]:
    """Trace rays from each grid point to observer, computing shadow depth.

    For each target grid point, walks the ray back toward the observer in
    grid space using DDA. At each grid cell boundary crossing, interpolates
    the terrain elevation angle between the two bounding grid points along
    that axis. Tracks the maximum elevation angle seen along the ray. If
    the target's elevation is below that maximum, it is in shadow.

    Args:
        elev: Elevation angles from observer to each grid point. Shape (M, N).
        h_point: Height above observer horizontal plane. Shape (M, N).
        r_point: Horizontal distance from observer. Shape (M, N).
        obs_row: Observer fractional row in grid space.
        obs_col: Observer fractional column in grid space.

    Returns:
        Shadow depth in meters. Shape (M, N).
    """
    m, n = elev.shape
    depth = np.empty((m, n), dtype=np.float64)

    for i in range(m):
        for j in range(n):
            max_blocking_elev = _trace_ray_to_observer(
                elev,
                obs_row,
                obs_col,
                i,
                j,
            )
            target_elev = elev[i, j]

            if max_blocking_elev <= -np.pi / 2.0:
                # No blocking terrain along the ray, fully visible
                depth[i, j] = abs(h_point[i, j]) + r_point[i, j] * 0.01
            elif target_elev > max_blocking_elev:
                # Visible with clearance
                h_ray = r_point[i, j] * np.tan(max_blocking_elev)
                depth[i, j] = h_point[i, j] - h_ray
            else:
                # Shadowed
                h_ray = r_point[i, j] * np.tan(max_blocking_elev)
                depth[i, j] = h_point[i, j] - h_ray  # negative

    return depth


def _trace_ray_to_observer(
    elev: NDArray[np.floating],
    obs_row: float,
    obs_col: float,
    target_row: int,
    target_col: int,
) -> float:
    """Trace ray from target to observer, return max blocking elevation.

    Uses DDA traversal in grid space. At each grid cell boundary the ray
    crosses, interpolates the elevation angle between the two adjacent
    grid points that straddle the crossing. Returns the maximum
    interpolated elevation encountered along the entire ray.

    Args:
        elev: Elevation angles at grid points. Shape (M, N).
        obs_row: Observer fractional row.
        obs_col: Observer fractional column.
        target_row: Target grid row index.
        target_col: Target grid column index.

    Returns:
        Maximum blocking elevation angle (radians). Returns -pi/2 if no
        blocking terrain found (nothing between target and observer).
    """
    dr = obs_row - target_row
    dc = obs_col - target_col

    if abs(dr) < _GRID_COINCIDENCE_TOL and abs(dc) < _GRID_COINCIDENCE_TOL:
        return -np.pi / 2.0

    max_blocking = -np.pi / 2.0

    # DDA: step through grid cells along the ray from target toward observer.
    # Parameterize as row(t) = target_row + t*dr, col(t) = target_col + t*dc
    # where t=0 is target and t=1 is observer. At each integer row/col
    # crossing, interpolate elevation from adjacent grid points.

    if abs(dr) > _GRID_DIRECTION_TOL:
        row_max = _scan_row_crossings(elev, target_row, target_col, dr, dc)
        max_blocking = max(max_blocking, row_max)

    if abs(dc) > _GRID_DIRECTION_TOL:
        col_max = _scan_col_crossings(elev, target_row, target_col, dr, dc)
        max_blocking = max(max_blocking, col_max)

    return max_blocking


def _scan_row_crossings(
    elev: NDArray[np.floating],
    target_row: int,
    target_col: int,
    dr: float,
    dc: float,
) -> float:
    """Scan row boundary crossings along the ray, return max elevation."""
    m, n = elev.shape
    max_blocking = -np.pi / 2.0

    if dr > 0:
        row_start = target_row + 2
        row_end = min(int(np.floor(target_row + dr)), m - 1)
        row_crossings = range(row_start, row_end + 1)
    else:
        row_start = target_row - 2
        row_end = max(int(np.ceil(target_row + dr)), 0)
        row_crossings = range(row_start, row_end - 1, -1)

    for r in row_crossings:
        t = (r - target_row) / dr
        if t >= _RAY_END_FRACTION:
            continue
        c_at = max(0.0, min(target_col + t * dc, n - 1.0))
        c0 = int(np.floor(c_at))
        c1 = min(c0 + 1, n - 1)
        frac_c = c_at - c0
        r_below = max(r - 1, 0)
        r_above = min(r, m - 1)
        e_low = elev[r_below, c0] * (1 - frac_c) + elev[r_below, c1] * frac_c
        e_high = elev[r_above, c0] * (1 - frac_c) + elev[r_above, c1] * frac_c
        max_blocking = max(max_blocking, e_low, e_high)

    return max_blocking


def _scan_col_crossings(
    elev: NDArray[np.floating],
    target_row: int,
    target_col: int,
    dr: float,
    dc: float,
) -> float:
    """Scan column boundary crossings along the ray, return max elevation."""
    m, n = elev.shape
    max_blocking = -np.pi / 2.0

    if dc > 0:
        col_start = target_col + 2
        col_end = min(int(np.floor(target_col + dc)), n - 1)
        col_crossings = range(col_start, col_end + 1)
    else:
        col_start = target_col - 2
        col_end = max(int(np.ceil(target_col + dc)), 0)
        col_crossings = range(col_start, col_end - 1, -1)

    for c in col_crossings:
        t = (c - target_col) / dc
        if t >= _RAY_END_FRACTION:
            continue
        r_at = max(0.0, min(target_row + t * dr, m - 1.0))
        r0 = int(np.floor(r_at))
        r1 = min(r0 + 1, m - 1)
        frac_r = r_at - r0
        c_left = max(c - 1, 0)
        c_right = min(c, n - 1)
        e_left = elev[r0, c_left] * (1 - frac_r) + elev[r1, c_left] * frac_r
        e_right = elev[r0, c_right] * (1 - frac_r) + elev[r1, c_right] * frac_r
        max_blocking = max(max_blocking, e_left, e_right)

    return max_blocking
