"""
hazard_classifier.py — Depth-image-based hazard detection.

Runs independently of YOLO detections.  The scene_detector_node calls
HazardClassifier.analyse() once per frame to check for:

  HAZARD_STAIRS    — stair pattern: repeating horizontal edges at increasing depth
  HAZARD_DROP      — floor cliff below robot's current heading
  HAZARD_WET_FLOOR — specular floor reflection (high-variance narrow depth stripe)
  HAZARD_GLASS_DOOR — depth near-zero where RGB has strong vertical edges

All methods take float32 depth images in metres and optional BGR images.
Returns a list of (hazard_type, confidence, description) tuples.

Tuning constants are conservative to avoid false positives in normal use.
"""

from __future__ import annotations
from typing import List, Tuple

import numpy as np

# Hazard type constants (match SceneObject.msg)
HAZARD_NONE      = 0
HAZARD_STAIRS    = 1
HAZARD_DROP      = 2
HAZARD_WET_FLOOR = 3
HAZARD_GLASS_DOOR = 4
HAZARD_PET       = 5

HazardResult = Tuple[int, float, str]   # (hazard_type, confidence, description)


class HazardClassifier:
    """
    Stateless hazard analyser.  Call analyse() per frame.
    """

    def __init__(
        self,
        drop_threshold_m:   float = 0.25,   # depth discontinuity = drop if > this
        stair_min_steps:    int   = 3,       # minimum horizontal bands to flag stairs
        glass_depth_thr_m:  float = 0.15,   # depth below which we suspect glass (unreturned)
        wet_variance_thr:   float = 0.08,   # depth std-dev in floor stripe = specular
    ):
        self._drop_thr      = drop_threshold_m
        self._stair_steps   = stair_min_steps
        self._glass_thr     = glass_depth_thr_m
        self._wet_var_thr   = wet_variance_thr

    def analyse(
        self,
        depth: np.ndarray,
        bgr:   np.ndarray | None = None,
    ) -> List[HazardResult]:
        """
        Analyse a depth (and optionally BGR) frame for hazards.

        Args:
            depth: float32 H×W depth in metres.  0 = invalid/no return.
            bgr:   uint8 H×W×3 BGR image (same resolution).  None = skip visual checks.

        Returns:
            List of (hazard_type, confidence, description) tuples.
            Empty list = no hazards detected.
        """
        results: List[HazardResult] = []
        h, w = depth.shape

        # ── 1. Floor drop / cliff ──────────────────────────────────────────
        result = self._check_drop(depth, h, w)
        if result:
            results.append(result)

        # ── 2. Stair pattern ───────────────────────────────────────────────
        result = self._check_stairs(depth, h, w)
        if result:
            results.append(result)

        # ── 3. Wet floor (specular) ────────────────────────────────────────
        result = self._check_wet_floor(depth, h, w)
        if result:
            results.append(result)

        # ── 4. Glass door (depth=0, RGB has vertical edges) ───────────────
        if bgr is not None:
            result = self._check_glass(depth, bgr, h, w)
            if result:
                results.append(result)

        return results

    # ── Private checks ──────────────────────────────────────────────────────

    def _check_drop(self, depth: np.ndarray, h: int, w: int) -> HazardResult | None:
        """
        Drop: in the lower 20% of the frame (floor area), check for large depth
        jumps between adjacent rows — indicates floor edge or drop-off.
        """
        floor_start = int(h * 0.70)
        floor       = depth[floor_start:, w // 4: 3 * w // 4]   # central strip
        valid       = (floor > 0.1) & (floor < 8.0)

        if valid.sum() < 50:
            return None

        # Row-mean of valid depths
        row_means = np.array([
            np.nanmean(np.where(valid[r], floor[r], np.nan))
            for r in range(floor.shape[0])
        ])
        row_means = row_means[~np.isnan(row_means)]
        if len(row_means) < 3:
            return None

        # Consecutive row depth differences
        diffs = np.abs(np.diff(row_means))
        max_jump = float(diffs.max())
        if max_jump > self._drop_thr:
            conf = min(1.0, max_jump / (self._drop_thr * 3))
            return HAZARD_DROP, conf, f'floor drop {max_jump:.2f}m detected ahead'
        return None

    def _check_stairs(self, depth: np.ndarray, h: int, w: int) -> HazardResult | None:
        """
        Stairs: horizontal banding pattern in depth — multiple rows at distinct
        increasing depths, each band ~step-height apart (typically 0.15–0.22m).
        """
        # Use central 40% strip (avoids walls on sides)
        col_start, col_end = int(w * 0.3), int(w * 0.7)
        strip = depth[:, col_start:col_end]
        valid = (strip > 0.3) & (strip < 5.0)

        # Row medians
        row_med = np.array([
            np.median(strip[r, valid[r]]) if valid[r].sum() > 5 else np.nan
            for r in range(h)
        ])
        finite_mask = ~np.isnan(row_med)
        if finite_mask.sum() < 20:
            return None

        # Smooth to suppress noise
        kernel = np.ones(5) / 5
        smoothed = np.convolve(row_med[finite_mask], kernel, mode='same')

        # Count direction changes (peaks / valleys) — stairs produce alternating pattern
        diffs   = np.diff(smoothed)
        changes = np.sum(np.diff(np.sign(diffs)) != 0)

        # Also check that depth range is consistent with steps (0.1–0.3m per band)
        range_m = float(smoothed.max() - smoothed.min())
        step_like = 0.3 < range_m < 3.0

        if changes >= self._stair_steps * 2 and step_like:
            conf = min(1.0, changes / 20.0)
            return HAZARD_STAIRS, conf, f'stair pattern detected ({changes} bands)'
        return None

    def _check_wet_floor(self, depth: np.ndarray, h: int, w: int) -> HazardResult | None:
        """
        Wet/glossy floor: depth returns are noisy / highly variable in the
        floor region because specular reflections scatter the IR pattern.
        A large std-dev in the floor stripe with valid-looking mean depth is suspect.
        """
        floor_start = int(h * 0.75)
        floor = depth[floor_start:, w // 4: 3 * w // 4]
        valid = (floor > 0.1) & (floor < 3.0)
        if valid.sum() < 30:
            return None

        valid_vals = floor[valid]
        std = float(valid_vals.std())
        if std > self._wet_var_thr:
            conf = min(1.0, std / (self._wet_var_thr * 4))
            return HAZARD_WET_FLOOR, conf, f'wet/specular floor (depth std={std:.3f}m)'
        return None

    def _check_glass(
        self,
        depth: np.ndarray,
        bgr:   np.ndarray,
        h: int,
        w: int,
    ) -> HazardResult | None:
        """
        Glass door: central column has near-zero depth returns (IR passes through)
        while the RGB image shows strong vertical edges (door frame / objects behind).

        Requires both depth and BGR images.
        """
        import cv2

        central_d = depth[:, w // 3: 2 * w // 3]
        near_zero = (central_d < self._glass_thr) & (central_d >= 0)
        zero_frac = float(near_zero.mean())

        if zero_frac < 0.15:
            return None

        # Check for strong vertical edges in RGB
        gray    = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
        sobel_x = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
        central_sobel = np.abs(sobel_x[:, w // 3: 2 * w // 3])
        edge_strength = float(central_sobel.mean())

        if edge_strength > 15.0:
            conf = min(1.0, zero_frac * 2 * edge_strength / 50)
            return HAZARD_GLASS_DOOR, conf, (
                f'glass door suspected (zero_depth={zero_frac:.0%}, edge={edge_strength:.1f})'
            )
        return None
