"""Monocular depth estimation."""

from __future__ import annotations

import threading
from pathlib import Path

import cv2
import numpy as np


def _pick_device() -> str:
    """Best available inference device."""
    try:
        import torch
    except ImportError:
        return "cpu"
    try:
        if torch.cuda.is_available():
            return "cuda"
        if torch.backends.mps.is_available():
            return "mps"
    except (AttributeError, RuntimeError):
        return "cpu"
    return "cpu"


class DepthEstimator:
    """Monocular depth from a YOLO depth model."""

    def __init__(self, model_path: str | Path, imgsz: int = 384) -> None:
        from ultralytics import YOLO

        model_path = Path(model_path)
        if not model_path.exists():
            raise FileNotFoundError(f"Depth model not found: {model_path}")
        self.device = _pick_device()
        self.model = YOLO(str(model_path))
        self.imgsz = imgsz
        self._lo: float | None = None
        self._hi: float | None = None
        self._bounds_alpha = 0.08

    def __call__(self, frame: np.ndarray) -> np.ndarray:
        """Raw depth map, same size as frame."""
        result = self.model.predict(
            frame, imgsz=self.imgsz, device=self.device, verbose=False
        )[0]
        depth = result.depth.data
        if hasattr(depth, "cpu"):
            depth = depth.cpu().numpy()
        depth = np.asarray(depth, dtype=np.float32)
        if depth.ndim == 3:
            depth = depth[0]
        if depth.shape[:2] != frame.shape[:2]:
            depth = cv2.resize(depth, (frame.shape[1], frame.shape[0]),
                               interpolation=cv2.INTER_LINEAR)
        return depth

    def normalize(self, depth: np.ndarray) -> np.ndarray:
        """Stable 0..1 map, 0 near, 1 far."""
        finite = depth[np.isfinite(depth)]
        if finite.size == 0:
            return np.full(depth.shape, 0.5, dtype=np.float32)
        lo, hi = float(np.percentile(finite, 2.0)), float(np.percentile(finite, 98.0))
        if self._lo is None or self._hi is None:
            self._lo, self._hi = lo, hi
        else:
            self._lo += self._bounds_alpha * (lo - self._lo)
            self._hi += self._bounds_alpha * (hi - self._hi)
        span = max(self._hi - self._lo, 1e-6)
        norm = np.clip((depth - self._lo) / span, 0.0, 1.0)
        return np.nan_to_num(norm, nan=0.5).astype(np.float32)

    @staticmethod
    def colorize(depth_norm: np.ndarray) -> np.ndarray:
        """Turbo colormap of a normalized depth map."""
        u8 = (np.clip(depth_norm, 0.0, 1.0) * 255.0).astype(np.uint8)
        return cv2.applyColorMap(255 - u8, cv2.COLORMAP_TURBO)

    @staticmethod
    def sample(depth_map: np.ndarray, x: float, y: float, radius: int = 9,
               percentile: float = 20.0, default: float = 0.0) -> float:
        """Nearest surface around a point."""
        h, w = depth_map.shape[:2]
        xi = int(np.clip(round(float(x)), 0, w - 1))
        yi = int(np.clip(round(float(y)), 0, h - 1))
        x0, x1 = max(0, xi - radius), min(w, xi + radius + 1)
        y0, y1 = max(0, yi - radius), min(h, yi + radius + 1)
        patch = depth_map[y0:y1, x0:x1]
        patch = patch[np.isfinite(patch)]
        if patch.size == 0:
            return default
        return float(np.percentile(patch, percentile))


class DepthWorker:
    """Background depth estimation thread."""

    def __init__(self, model_path: str | Path, imgsz: int = 384,
                 input_width: int = 640) -> None:
        self.estimator = DepthEstimator(model_path, imgsz)
        self.device = self.estimator.device
        self.input_width = input_width
        self._pending: np.ndarray | None = None
        self._target: tuple[int, int] | None = None
        self._metric: np.ndarray | None = None
        self._norm: np.ndarray | None = None
        self._seq = 0
        self._lock = threading.Lock()
        self._wake = threading.Event()
        self._stop = threading.Event()
        self._thread = threading.Thread(target=self._loop, daemon=True)
        self._thread.start()

    def submit(self, frame: np.ndarray) -> None:
        """Queue the newest frame."""
        h, w = frame.shape[:2]
        if w > self.input_width:
            k = self.input_width / float(w)
            small = cv2.resize(frame, (self.input_width, max(1, int(round(h * k)))),
                               interpolation=cv2.INTER_AREA)
        else:
            small = frame.copy()
        with self._lock:
            self._pending = small
            self._target = (w, h)
        self._wake.set()

    def _loop(self) -> None:
        """Consume frames until stopped."""
        while not self._stop.is_set():
            self._wake.wait(0.1)
            self._wake.clear()
            with self._lock:
                frame, target = self._pending, self._target
                self._pending = None
            if frame is None or target is None:
                continue
            try:
                metric = self.estimator(frame)
                norm = self.estimator.normalize(metric)
            except Exception as exc:  # keep the app alive on backend errors
                print(f"[!] depth failed: {exc}", flush=True)
                self._stop.set()
                return
            if (metric.shape[1], metric.shape[0]) != target:
                metric = cv2.resize(metric, target, interpolation=cv2.INTER_LINEAR)
                norm = cv2.resize(norm, target, interpolation=cv2.INTER_LINEAR)
            with self._lock:
                self._metric = metric
                self._norm = norm
                self._seq += 1

    @property
    def latest(self) -> np.ndarray | None:
        """Newest metric depth map."""
        with self._lock:
            return self._metric

    @property
    def latest_norm(self) -> np.ndarray | None:
        """Newest normalized depth map."""
        with self._lock:
            return self._norm

    @property
    def ready(self) -> bool:
        """A depth map is available."""
        return self.latest is not None

    @property
    def frames(self) -> int:
        """Number of finished estimations."""
        with self._lock:
            return self._seq

    def close(self) -> None:
        """Stop the worker thread."""
        self._stop.set()
        self._wake.set()
        self._thread.join(timeout=1.0)