"""
_terrain_roughness.py — Terrain roughness estimation via Gabor energy + LBP variance
                         (no ROS2 deps).

Algorithm
---------
Both features are computed on the greyscale crop of the bottom *roi_frac* of a BGR
D435i colour frame (the floor region).

1. Gabor filter-bank energy (4 orientations × 2 spatial wavelengths, quadrature pairs):

       E = mean over all (θ, λ) of  mean_pixels( |real_resp|² + |imag_resp|² )

   A quadrature pair uses phase ψ=0 (real) and ψ=π/2 (imaginary), giving orientation-
   and phase-invariant energy.  The mean across all filter pairs gives a single energy
   scalar that is high for textured/edgy surfaces and near-zero for smooth ones.

2. LBP variance (Local Binary Pattern, radius=1, 8-point, pure NumPy, no sklearn):

       LBP(x,y) = Σ_k  s(g_k − g_c) · 2^k      (k = 0..7, g_c = centre pixel)

   Computed as vectorised slice comparisons over the 8 cardinal/diagonal neighbours.
   The metric is var(LBP image).  A constant surface gives LBP ≡ 0xFF (all bits set,
   since every neighbour ties with the centre) → variance = 0.  Irregular textures
   produce a dispersed LBP histogram → high variance.

3. Normalised roughness blend:

       roughness = clip( _W_GABOR · E / _GABOR_REF  +  _W_LBP · V / _LBP_REF ,  0, 1 )

   Calibration: _GABOR_REF is tuned so that a coarse gravel / noise image maps to ~1.0;
   _LBP_REF = 5000 (≈ variance of a fully random 8-bit LBP image).

Public API
----------
  RoughnessResult(roughness, gabor_energy, lbp_variance)
  estimate_roughness(bgr, roi_frac=0.40) -> RoughnessResult
"""

from __future__ import annotations

from typing import NamedTuple, Optional

import numpy as np


# ── Tuning constants ──────────────────────────────────────────────────────────

# Gabor filter-bank parameters
_GABOR_N_ORIENTATIONS = 4          # θ ∈ {0°, 45°, 90°, 135°}
_GABOR_WAVELENGTHS    = [5.0, 10.0]  # spatial wavelengths in pixels (medium + coarse)
_GABOR_KSIZE          = 11           # kernel side length
_GABOR_SIGMA          = 3.0          # Gaussian envelope std (pixels)
_GABOR_GAMMA          = 0.5          # spatial aspect ratio

# Normalization references
_GABOR_REF = 500.0    # Gabor mean power at which we consider maximum roughness
_LBP_REF   = 5000.0  # LBP variance at which we consider maximum roughness

# Blend weights (must sum to 1)
_W_GABOR = 0.5
_W_LBP   = 0.5

# Module-level cache for the kernel bank (built once per process)
_gabor_kernels: Optional[list] = None


# ── Data types ────────────────────────────────────────────────────────────────

class RoughnessResult(NamedTuple):
    """Terrain roughness estimate for a single frame."""
    roughness:    float   # blended score in [0, 1]; 0=smooth, 1=rough
    gabor_energy: float   # raw Gabor mean energy (≥ 0)
    lbp_variance: float   # raw LBP image variance (≥ 0)


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

def _build_gabor_kernels() -> list:
    """Build (real_kernel, imag_kernel) pairs for all (orientation, wavelength) combos."""
    import cv2

    kernels = []
    for i in range(_GABOR_N_ORIENTATIONS):
        theta = float(i) * np.pi / _GABOR_N_ORIENTATIONS
        for lam in _GABOR_WAVELENGTHS:
            real_k = cv2.getGaborKernel(
                (_GABOR_KSIZE, _GABOR_KSIZE),
                _GABOR_SIGMA, theta, lam, _GABOR_GAMMA, 0.0,
                ktype=cv2.CV_64F,
            ).astype(np.float32)
            imag_k = cv2.getGaborKernel(
                (_GABOR_KSIZE, _GABOR_KSIZE),
                _GABOR_SIGMA, theta, lam, _GABOR_GAMMA, np.pi / 2.0,
                ktype=cv2.CV_64F,
            ).astype(np.float32)
            kernels.append((real_k, imag_k))
    return kernels


def _get_gabor_kernels() -> list:
    global _gabor_kernels
    if _gabor_kernels is None:
        _gabor_kernels = _build_gabor_kernels()
    return _gabor_kernels


def gabor_energy(grey: np.ndarray) -> float:
    """
    Mean quadrature Gabor energy across the filter bank.

    Parameters
    ----------
    grey : (H, W) uint8 or float ndarray

    Returns
    -------
    float — mean power (≥ 0); higher for textured surfaces
    """
    import cv2

    img = np.asarray(grey, dtype=np.float32)
    if img.size == 0:
        return 0.0

    # Subtract the mean so constant regions (DC) contribute zero energy.
    # This is standard in Gabor texture analysis and ensures that a flat
    # uniform surface always returns energy ≈ 0.
    img = img - float(img.mean())

    kernels = _get_gabor_kernels()
    total   = 0.0
    for real_k, imag_k in kernels:
        r = cv2.filter2D(img, cv2.CV_32F, real_k)
        i = cv2.filter2D(img, cv2.CV_32F, imag_k)
        total += float(np.mean(r * r + i * i))

    return total / len(kernels) if kernels else 0.0


def lbp_variance(grey: np.ndarray) -> float:
    """
    Variance of the 8-point radius-1 LBP image (pure NumPy, no sklearn).

    Parameters
    ----------
    grey : (H, W) uint8 ndarray

    Returns
    -------
    float — variance of LBP values; 0 for flat surfaces, high for irregular textures
    """
    g = np.asarray(grey, dtype=np.int32)
    h, w = g.shape
    if h < 3 or w < 3:
        return 0.0

    # Centre pixel patch (interior only, 1-px border excluded)
    c = g[1:h-1, 1:w-1]

    # 8-neighbourhood comparisons using array slices (vectorised, no loops)
    lbp = (
        ((g[0:h-2, 0:w-2] >= c).view(np.uint8)      ) |
        ((g[0:h-2, 1:w-1] >= c).view(np.uint8) << 1 ) |
        ((g[0:h-2, 2:w  ] >= c).view(np.uint8) << 2 ) |
        ((g[1:h-1, 0:w-2] >= c).view(np.uint8) << 3 ) |
        ((g[1:h-1, 2:w  ] >= c).view(np.uint8) << 4 ) |
        ((g[2:h  , 0:w-2] >= c).view(np.uint8) << 5 ) |
        ((g[2:h  , 1:w-1] >= c).view(np.uint8) << 6 ) |
        ((g[2:h  , 2:w  ] >= c).view(np.uint8) << 7 )
    )

    return float(lbp.var())


# ── Public API ────────────────────────────────────────────────────────────────

def estimate_roughness(
    bgr:      np.ndarray,
    roi_frac: float = 0.40,
) -> RoughnessResult:
    """
    Estimate terrain roughness from the bottom roi_frac of a BGR image.

    Parameters
    ----------
    bgr      : (H, W, 3) uint8 BGR ndarray (or greyscale (H, W))
    roi_frac : fraction of image height to use as the floor ROI (from bottom)

    Returns
    -------
    RoughnessResult(roughness, gabor_energy, lbp_variance)
    roughness is in [0, 1] where 0=smooth, 1=rough.
    """
    import cv2

    # Crop to bottom roi_frac
    h = bgr.shape[0]
    r0 = int(h * max(0.0, 1.0 - min(float(roi_frac), 1.0)))
    roi = bgr[r0:, :]

    if roi.size == 0:
        return RoughnessResult(roughness=0.0, gabor_energy=0.0, lbp_variance=0.0)

    # Convert to greyscale
    if roi.ndim == 3:
        grey = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
    else:
        grey = np.asarray(roi, dtype=np.uint8)

    ge = gabor_energy(grey)
    lv = lbp_variance(grey)

    gabor_norm = min(ge / _GABOR_REF, 1.0)
    lbp_norm   = min(lv / _LBP_REF,   1.0)

    roughness = float(np.clip(_W_GABOR * gabor_norm + _W_LBP * lbp_norm, 0.0, 1.0))

    return RoughnessResult(roughness=roughness, gabor_energy=ge, lbp_variance=lv)
