"""Depth Feature Extractor — Surface roughness from RealSense depth."""

from __future__ import annotations
import numpy as np
from typing import Optional
import cv2


class DepthExtractor:
    """Extract terrain roughness features from RealSense depth images."""

    def __init__(
        self,
        roi_width_px: int = 320,
        roi_height_px: int = 240,
        min_depth_m: float = 0.1,
        max_depth_m: float = 3.0,
        bilateral_d: int = 5,
    ):
        self._roi_w = roi_width_px
        self._roi_h = roi_height_px
        self._min_depth = min_depth_m
        self._max_depth = max_depth_m
        self._bilateral_d = bilateral_d

    def extract_roughness(self, depth_image: np.ndarray) -> Optional[float]:
        """Extract surface roughness from depth image."""
        if depth_image is None or depth_image.size == 0:
            return None

        if depth_image.dtype == np.uint16:
            depth_m = depth_image.astype(np.float32) / 1000.0
        else:
            depth_m = depth_image.astype(np.float32)

        h, w = depth_m.shape
        x_start = (w - self._roi_w) // 2
        y_start = (h - self._roi_h) // 2
        x_end = x_start + self._roi_w
        y_end = y_start + self._roi_h

        roi = depth_m[max(0, y_start):min(h, y_end), max(0, x_start):min(w, x_end)]

        valid_mask = (roi > self._min_depth) & (roi < self._max_depth)
        valid_depths = roi[valid_mask]

        if len(valid_depths) < 10:
            return None

        roi_smooth = cv2.bilateralFilter(roi.astype(np.float32), self._bilateral_d, 0.1, 1.0)
        depth_variance = np.var(valid_depths)

        grad_x = cv2.Sobel(roi_smooth, cv2.CV_32F, 1, 0, ksize=3)
        grad_y = cv2.Sobel(roi_smooth, cv2.CV_32F, 0, 1, ksize=3)
        grad_magnitude = np.sqrt(grad_x**2 + grad_y**2)

        valid_grad = grad_magnitude[valid_mask]
        grad_mean = np.mean(valid_grad) if len(valid_grad) > 0 else 0.0

        depth_roughness = min(1.0, depth_variance / 0.05)
        grad_roughness = min(1.0, grad_mean / 0.02)

        roughness = 0.6 * depth_roughness + 0.4 * grad_roughness
        return float(roughness)
