"""Monocular depth estimation with model selection, resolution scaling, and temporal smoothing.

Supported models (set via NAV_DEPTH_MODEL env or config.py):
  "small"    — Depth Anything V1 Small  (fastest)
  "base"     — Depth Anything V1 Base
  "v2-small" — Depth Anything V2 Small  (recommended for MVP)
  "v2-base"  — Depth Anything V2 Base   (best accuracy)

Performance knobs:
  NAV_DEPTH_INPUT_SIZE — resize input before inference (default 384, try 256 for speed)
  NAV_DEVICE           — "cpu" or "cuda"
"""
import cv2
import numpy as np
import math
import time
from config import (
    DEPTH_MODEL, DEPTH_INPUT_SIZE, DEVICE, GROUND_RATIO, SLOPE_MULTIPLIER,
    SLOPE_CLAMP, SLOPE_DEADZONE,
    TERRAIN_ROCKY_THRESH, TERRAIN_ROUGH_THRESH,
)

_MODEL_MAP = {
    "small":    "LiheYoung/depth-anything-small-hf",
    "base":     "LiheYoung/depth-anything-base-hf",
    "v2-small": "depth-anything/Depth-Anything-V2-Small-hf",
    "v2-base":  "depth-anything/Depth-Anything-V2-Base-hf",
}

_depth_pipe = None
_onnx_session = None
_backend = None  # "hf" or "onnx"


def _load():
    global _depth_pipe, _backend
    if _depth_pipe is not None:
        return

    model_id = _MODEL_MAP.get(DEPTH_MODEL, DEPTH_MODEL)

    # Try ONNX first if available
    if _try_load_onnx(model_id):
        return

    # Fall back to HuggingFace pipeline
    from transformers import pipeline as hf_pipeline
    t0 = time.time()
    _depth_pipe = hf_pipeline("depth-estimation", model=model_id, device=DEVICE)
    _backend = "hf"
    print(f"[Depth] Loaded {model_id} on {DEVICE} ({time.time()-t0:.1f}s)", flush=True)


def _try_load_onnx(model_id):
    """Try loading an ONNX-exported model for faster CPU inference."""
    global _onnx_session, _backend
    try:
        import onnxruntime as ort
        import os
        # Look for local ONNX file
        onnx_path = os.environ.get("NAV_DEPTH_ONNX")
        if not onnx_path:
            return False
        if not os.path.exists(onnx_path):
            print(f"[Depth] ONNX path not found: {onnx_path}", flush=True)
            return False

        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if DEVICE == "cuda" else ['CPUExecutionProvider']
        t0 = time.time()
        _onnx_session = ort.InferenceSession(onnx_path, providers=providers)
        _backend = "onnx"
        print(f"[Depth] Loaded ONNX model ({time.time()-t0:.1f}s) providers={_onnx_session.get_providers()}", flush=True)
        return True
    except ImportError:
        return False


def _infer_hf(rgb_small):
    """Run HuggingFace pipeline on a (possibly resized) RGB image."""
    from PIL import Image as PILImage
    return np.array(_depth_pipe(PILImage.fromarray(rgb_small))["depth"]).astype(np.float32)


def _infer_onnx(rgb_small):
    """Run ONNX session on preprocessed input."""
    # Standard normalization for Depth Anything
    img = rgb_small.astype(np.float32) / 255.0
    mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
    std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
    img = (img - mean) / std
    img = np.transpose(img, (2, 0, 1))[np.newaxis]  # NCHW

    input_name = _onnx_session.get_inputs()[0].name
    result = _onnx_session.run(None, {input_name: img})
    return result[0].squeeze().astype(np.float32)


def estimate_depth(frame_rgb, h, w):
    """Run depth model. Returns normalized depth map (0-1) at original frame resolution."""
    _load()

    # Resize for faster inference
    inp_size = DEPTH_INPUT_SIZE
    if h > inp_size or w > inp_size:
        scale = inp_size / max(h, w)
        sh, sw = int(h * scale), int(w * scale)
        rgb_small = cv2.resize(frame_rgb, (sw, sh))
    else:
        rgb_small = frame_rgb

    # Inference
    if _backend == "onnx":
        # ONNX needs exact square input for some exports
        sq = cv2.resize(rgb_small, (inp_size, inp_size))
        dm = _infer_onnx(sq)
    else:
        dm = _infer_hf(rgb_small)

    # Resize back to original resolution
    dm = cv2.resize(dm, (w, h))

    dmin, dmax = dm.min(), dm.max()
    if dmax - dmin < 1e-8:
        return np.zeros((h, w), dtype=np.float32)
    return (dm - dmin) / (dmax - dmin)


def estimate_slope(depth_norm, h, w, foot_y=None):
    """Compute slope angle and direction from depth map ground region.

    Returns (slope_angle, slope_dir, terrain_type, ground_start_y).
    """
    if foot_y and foot_y > h * 0.3:
        gs = max(0, int(foot_y) - int(h * 0.1))
    else:
        gs = int(h * GROUND_RATIO)

    ground = depth_norm[gs:, :]
    gh, gw = ground.shape
    if gh < 10 or gw < 10:
        return 0.0, "FLAT", "SMOOTH", gs

    gy = cv2.Sobel(ground, cv2.CV_64F, 0, 1, ksize=5)
    cl, cr = gw // 3, 2 * gw // 3

    rows = []
    sh = max(1, gh // 6)
    for i in range(6):
        y0, y1 = i * sh, min((i + 1) * sh, gh)
        c = np.concatenate([gy[y0:y1, :cl].flatten(), gy[y0:y1, cr:].flatten()])
        if len(c):
            rows.append(float(np.median(c)))

    if len(rows) < 3:
        return 0.0, "FLAT", "SMOOTH", gs

    trend = np.polyfit(np.arange(len(rows)), np.array(rows), 1)[0]
    sa = float(np.clip(np.arctan(trend * SLOPE_MULTIPLIER) * 180 / math.pi,
                        -SLOPE_CLAMP, SLOPE_CLAMP))
    if abs(sa) < SLOPE_DEADZONE:
        sa = 0.0

    slope_dir = "FLAT" if abs(sa) < SLOPE_DEADZONE else ("UPHILL" if sa > 0 else "DOWNHILL")

    gvar = float(np.std(ground))
    if gvar > TERRAIN_ROCKY_THRESH:
        terrain = "ROCKY"
    elif gvar > TERRAIN_ROUGH_THRESH:
        terrain = "ROUGH"
    else:
        terrain = "SMOOTH"

    return sa, slope_dir, terrain, gs


class SlopeSmoother:
    """Temporal smoothing for slope estimates across video frames."""

    def __init__(self, alpha=0.7, outlier_thresh=15.0):
        self.alpha = alpha
        self.outlier_thresh = outlier_thresh
        self._angle = 0.0
        self._dir = "FLAT"
        self._terrain = "SMOOTH"
        self._history = []
        self._max_history = 10

    def update(self, raw_angle, raw_dir, raw_terrain):
        """Feed a new raw slope measurement. Returns smoothed (angle, dir, terrain)."""
        self._history.append(raw_angle)
        if len(self._history) > self._max_history:
            self._history.pop(0)

        # Outlier rejection
        if len(self._history) >= 3:
            median = float(np.median(self._history))
            if abs(raw_angle - median) > self.outlier_thresh:
                raw_angle = median

        self._angle = self.alpha * self._angle + (1 - self.alpha) * raw_angle

        if abs(self._angle) < SLOPE_DEADZONE:
            self._angle = 0.0
            self._dir = "FLAT"
        else:
            self._dir = "UPHILL" if self._angle > 0 else "DOWNHILL"

        self._terrain = raw_terrain
        return self._angle, self._dir, self._terrain

    @property
    def trend(self):
        if len(self._history) < 4:
            return "STABLE"
        recent = self._history[-4:]
        avg_diff = sum(recent[i+1] - recent[i] for i in range(len(recent)-1)) / (len(recent)-1)
        if avg_diff > 2.0:
            return "STEEPENING"
        elif avg_diff < -2.0:
            return "FLATTENING"
        return "STABLE"

    @property
    def angle(self):
        return self._angle

    @property
    def direction(self):
        return self._dir

    @property
    def terrain(self):
        return self._terrain


def detect_walls(depth_norm, h, w, close_thresh=0.6, min_area_ratio=0.05, max_std=0.08):
    """Detect wall/door-like surfaces from depth map.

    Looks for large, flat, close vertical regions.
    Returns list of dicts with box, label, proximity, direction.
    """
    # Focus on close regions
    close_mask = (depth_norm > close_thresh).astype(np.uint8) * 255

    # Morphological cleanup
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (15, 15))
    close_mask = cv2.morphologyEx(close_mask, cv2.MORPH_CLOSE, kernel)
    close_mask = cv2.morphologyEx(close_mask, cv2.MORPH_OPEN, kernel)

    contours, _ = cv2.findContours(close_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    walls = []
    min_area = min_area_ratio * h * w

    for cnt in contours:
        area = cv2.contourArea(cnt)
        if area < min_area:
            continue

        x, y, bw, bh = cv2.boundingRect(cnt)

        # Must be taller than wide (vertical surface) or very wide (wall ahead)
        if bh < h * 0.25 and bw < w * 0.4:
            continue

        # Check flatness (low std = flat surface like wall/door)
        region = depth_norm[y:y+bh, x:x+bw]
        if region.std() > max_std:
            continue

        proximity = float(region.mean())
        cx = x + bw // 2
        direction = "LEFT" if cx < w * 0.33 else "RIGHT" if cx > w * 0.66 else "CENTER"
        label = "door" if bw < w * 0.4 and bh > h * 0.5 else "wall"

        walls.append({
            'label': label,
            'box': (x, y, x + bw, y + bh),
            'center': (cx, y + bh // 2),
            'proximity': proximity,
            'dist': "NEAR" if proximity > 0.7 else "MID",
            'direction': direction,
        })

    return walls