"""
_depth_hole_fill.py — Depth image hole filling via bilateral interpolation (no ROS2 deps).

Algorithm
---------
A "hole" is any pixel where depth == 0, depth is NaN, or depth is outside the
valid range [d_min, d_max].

Each pass replaces every hole pixel with the spatial-Gaussian-weighted mean of
valid pixels in a (kernel_size × kernel_size) neighbourhood:

    filled[x,y] = Σ G(||p - q||; σ) · d[q]  /  Σ G(||p - q||; σ)
                  q ∈ valid neighbours of (x,y)

The denominator (sum of spatial weights over valid pixels) normalises correctly
even at image borders and around isolated valid pixels.

Multiple passes with geometrically growing kernels are applied so that:
  Pass 1  kernel_size           — fills small holes (≤ kernel_size/2 px radius)
  Pass 2  kernel_size × 2.5    — fills medium holes
  Pass 3  kernel_size × 6.0    — fills large holes / fronto-parallel surfaces

After all passes any remaining zeros are left as-is (no valid neighbourhood data).

Because only the spatial Gaussian (not a depth range term) is used as the weighting
function, this is equivalent to a bilateral filter with σ_range → ∞.  In practice
this produces smooth, physically plausible fills in the depth domain.

Public API
----------
  fill_holes(depth, kernel_size=5, d_min=0.1, d_max=10.0, max_passes=3) → ndarray
  valid_mask(depth, d_min=0.1, d_max=10.0) → bool ndarray
"""

from __future__ import annotations

import math
from typing import Optional

import numpy as np


# Kernel size multipliers for successive passes
_PASS_SCALE = [1.0, 2.5, 6.0]


def valid_mask(
    depth:  np.ndarray,
    d_min:  float = 0.1,
    d_max:  float = 10.0,
) -> np.ndarray:
    """
    Return a boolean mask of valid (non-hole) pixels.

    Parameters
    ----------
    depth : (H, W) float32 ndarray, depth in metres
    d_min : minimum valid depth (m)
    d_max : maximum valid depth (m)

    Returns
    -------
    (H, W) bool ndarray — True where depth is finite and in [d_min, d_max]
    """
    return np.isfinite(depth) & (depth >= d_min) & (depth <= d_max)


def fill_holes(
    depth:       np.ndarray,
    kernel_size: int   = 5,
    d_min:       float = 0.1,
    d_max:       float = 10.0,
    max_passes:  int   = 3,
) -> np.ndarray:
    """
    Fill zero/NaN depth pixels using multi-pass spatial Gaussian interpolation.

    Parameters
    ----------
    depth       : (H, W) float32 ndarray, depth in metres
    kernel_size : initial kernel side length (pixels, forced odd, ≥ 3)
    d_min       : minimum valid depth — pixels below this are treated as holes
    d_max       : maximum valid depth — pixels above this are treated as holes
    max_passes  : number of fill passes (1–3); each uses a larger kernel

    Returns
    -------
    (H, W) float32 ndarray — same as input, with holes filled where possible.
    Pixels with no valid neighbours after all passes remain 0.0.
    Original valid pixels are never modified.
    """
    import cv2

    depth = np.asarray(depth, dtype=np.float32)
    # Replace NaN with 0 so arithmetic is clean
    depth = np.where(np.isfinite(depth), depth, 0.0).astype(np.float32)

    mask = valid_mask(depth, d_min, d_max)    # True where already valid
    result = depth.copy()
    n_passes = max(1, min(max_passes, len(_PASS_SCALE)))

    for i in range(n_passes):
        if mask.all():
            break   # no holes left

        ks   = _odd_kernel_size(kernel_size, _PASS_SCALE[i])
        half = ks // 2
        sigma = max(half / 2.0, 0.5)

        gk     = cv2.getGaussianKernel(ks, sigma).astype(np.float32)
        kernel = (gk @ gk.T)

        # Multiply depth by mask so invalid pixels contribute 0 weight
        d_valid = np.where(mask, result, 0.0).astype(np.float32)
        w_valid = mask.astype(np.float32)

        sum_d = cv2.filter2D(d_valid, ddepth=-1, kernel=kernel,
                             borderType=cv2.BORDER_REFLECT)
        sum_w = cv2.filter2D(w_valid, ddepth=-1, kernel=kernel,
                             borderType=cv2.BORDER_REFLECT)

        # Where we have enough weight, compute the weighted mean
        has_data = sum_w > 1e-6
        interp   = np.where(has_data, sum_d / np.where(has_data, sum_w, 1.0), 0.0)

        # Only fill holes — never overwrite original valid pixels
        result = np.where(mask, result, interp.astype(np.float32))

        # Update mask with newly filled pixels (for the next pass)
        newly_filled = (~mask) & (result > 0.0)
        mask = mask | newly_filled

    return result.astype(np.float32)


# ── Internal helpers ──────────────────────────────────────────────────────────

def _odd_kernel_size(base: int, scale: float) -> int:
    """Return the nearest odd integer to base * scale, minimum 3."""
    raw = max(3, int(round(base * scale)))
    return raw if raw % 2 == 1 else raw + 1
