"""Extend reliable depth along each visible component's silhouette."""

from dataclasses import replace


def extend_edge_depth(frame, alpha, alpha_threshold=0.9):
    """Return edge-corrected geometry while preserving reliable camera points."""
    import numpy as np
    from scipy import ndimage

    from ..models.onnx_moge2 import depth_to_points

    depth = frame.depth
    alpha = np.asarray(alpha, dtype=np.float32)
    if alpha.shape != depth.shape:
        raise ValueError("Edge Alpha must match the depth dimensions")
    if np.all(alpha == 1.0):
        return frame

    visible = (alpha > 0.0) & (alpha >= alpha_threshold)
    opaque = visible & (alpha >= 0.95)
    candidates = (
        opaque
        & (frame.validity > 0.5)
        & np.isfinite(depth)
        & (depth > 0.0)
    )
    if not candidates.any():
        return frame
    interior_distance = ndimage.distance_transform_edt(
        np.pad(opaque, 3, mode="edge")
    )[3:-3, 3:-3]
    interior = opaque & (interior_distance > 2.0)
    core = candidates & interior
    del interior_distance
    labels, _ = ndimage.label(visible, structure=np.ones((3, 3), dtype=bool))
    corrected = depth.copy()
    background_distance = np.full(depth.shape, np.inf, dtype=np.float32)

    for component_id, bounds in enumerate(ndimage.find_objects(labels), start=1):
        region = tuple(
            slice(max(0, bound.start - 7), min(size, bound.stop + 7))
            for bound, size in zip(bounds, depth.shape)
        )
        component = labels[region] == component_id
        donors = core[region] & component
        if not donors.any():
            donors = candidates[region] & component
        if not donors.any():
            continue

        indices = ndimage.distance_transform_edt(
            ~donors, return_distances=False, return_indices=True
        )
        extended = depth[region][tuple(indices)]
        smoothed = ndimage.gaussian_filter(extended, sigma=0.8, radius=3)
        distance = ndimage.distance_transform_edt(~component)
        background = (
            ~visible[region]
            & (distance <= 4.0)
            & (distance < background_distance[region])
        )
        target = (component & ~interior[region] & ~donors) | background
        corrected[region][target] = smoothed[target]
        background_distance[region][background] = distance[background]

    changed = corrected != depth
    if not changed.any():
        return frame
    points = frame.points
    if points is not None:
        height, width = depth.shape
        intrinsics = frame.intrinsics.copy()
        intrinsics[0] /= width
        intrinsics[1] /= height
        projected = depth_to_points(corrected, intrinsics)
        points = points.copy()
        points[changed] = projected[changed]
    return replace(frame, depth=corrected, points=points)