# ============================================================
# range_estimation.py
# Модуль 1: Оценка дальности до цели и time-to-contact (tau)
#
# Методы:
#   1) Монокулярная оценка по известному размеру цели
#   2) Tau — время до столкновения по скорости роста bbox
#   3) Скорость сближения (closing velocity) из производной range
# ============================================================

import numpy as np
from collections import deque

from config import *
from config_intercept import *
from helpers import box_wh, box_area, box_center, clamp


class RangeEstimator:
    """
    Оценка дальности до цели по размеру bounding box в монокулярной камере.

    Принцип: range = (known_size * focal_length) / apparent_size
    Дополнительно вычисляет tau (time-to-contact) по скорости изменения
    площади bbox, что не требует калибровки и даёт прямую оценку
    времени до столкновения.
    """

    def __init__(self):
        self.enabled = bool(RANGE_ENABLE)

        # Калибровка
        self.focal_px = float(CAMERA_FOCAL_LENGTH_PX)
        self.target_size_m = float(TARGET_KNOWN_SIZE_M)
        self.min_apparent = float(RANGE_MIN_APPARENT_SIZE_PX)

        # Сглаженные оценки
        self.range_m = None          # текущая оценка дальности (метры)
        self.range_raw = None        # без сглаживания
        self.closing_vel = 0.0       # скорость сближения (м/с), >0 = сближаемся
        self.tau = float('inf')      # time-to-contact (секунды)
        self.tau_raw = float('inf')

        # Внутреннее состояние
        self._prev_range = None
        self._prev_area = None
        self._prev_t = None
        self._range_alpha = float(RANGE_SMOOTH_ALPHA)
        self._tau_alpha = float(TAU_SMOOTH_ALPHA)

        # История для производных (скользящее окно)
        self._area_hist = deque(maxlen=8)    # (timestamp, area)
        self._range_hist = deque(maxlen=8)   # (timestamp, range)

        # Confidence: насколько доверяем оценке
        self.confidence = 0.0

    def reset(self):
        """Полный сброс при потере цели."""
        self.range_m = None
        self.range_raw = None
        self.closing_vel = 0.0
        self.tau = float('inf')
        self.tau_raw = float('inf')
        self._prev_range = None
        self._prev_area = None
        self._prev_t = None
        self._area_hist.clear()
        self._range_hist.clear()
        self.confidence = 0.0

    def update(self, box, timestamp, dt, confirmed=False):
        """
        Обновить оценку дальности и tau.

        Args:
            box: bounding box цели [x1, y1, x2, y2] в effective pixels
            timestamp: текущее время (секунды)
            dt: дельта времени с предыдущего кадра
            confirmed: цель подтверждена трекером

        Returns:
            dict с полями:
                range_m     — дальность (метры) или None
                range_raw   — дальность без сглаживания
                closing_vel — скорость сближения (м/с), >0 = сближаемся
                tau         — time-to-contact (сек), inf = не сближаемся
                tau_raw     — без сглаживания
                confidence  — 0..1
                phase       — "far" | "approach" | "warn" | "critical"
        """
        result = self._idle_result()

        if not self.enabled or box is None:
            return result

        w, h = box_wh(box)
        apparent = float(max(w, h))
        area = float(w * h)

        if apparent < self.min_apparent:
            return result

        # ─── 1. Дальность по размеру ────────────────────────────
        range_raw = self._range_from_size(apparent)
        range_raw = float(clamp(range_raw, RANGE_MIN_M, RANGE_MAX_M))
        self.range_raw = range_raw

        if self.range_m is None:
            self.range_m = range_raw
        else:
            alpha = self._range_alpha
            self.range_m = alpha * self.range_m + (1.0 - alpha) * range_raw

        # ─── 2. Closing velocity из dRange/dt ────────────────────
        self._range_hist.append((timestamp, range_raw))
        self.closing_vel = self._compute_closing_vel()

        # ─── 3. Tau из скорости изменения площади ────────────────
        self._area_hist.append((timestamp, area))
        tau_raw = self._compute_tau_from_area()
        self.tau_raw = tau_raw

        if self.tau == float('inf'):
            self.tau = tau_raw
        else:
            if tau_raw == float('inf'):
                self.tau = self.tau * 1.1  # медленный fade к inf
            else:
                alpha = self._tau_alpha
                self.tau = alpha * self.tau + (1.0 - alpha) * tau_raw
                self.tau = max(0.0, self.tau)

        # ─── 4. Confidence ───────────────────────────────────────
        self.confidence = self._compute_confidence(apparent, confirmed)

        # ─── 5. Phase ────────────────────────────────────────────
        phase = self._classify_phase()

        self._prev_range = range_raw
        self._prev_area = area
        self._prev_t = timestamp

        result.update({
            "range_m": float(self.range_m),
            "range_raw": float(self.range_raw),
            "closing_vel": float(self.closing_vel),
            "tau": float(self.tau),
            "tau_raw": float(self.tau_raw),
            "confidence": float(self.confidence),
            "phase": phase,
            "apparent_size_px": float(apparent),
            "area_px": float(area),
        })
        return result

    # ─── Private ─────────────────────────────────────────────────

    def _range_from_size(self, apparent_px):
        """range = (real_size * focal_length) / apparent_size"""
        if apparent_px < 1e-3:
            return float(RANGE_MAX_M)
        return (self.target_size_m * self.focal_px) / apparent_px

    def _compute_closing_vel(self):
        """
        Скорость сближения из изменения дальности.
        Положительная = сближаемся, отрицательная = удаляемся.
        Использует линейную регрессию по последним точкам.
        """
        if len(self._range_hist) < 3:
            return 0.0

        times = np.array([t for t, _ in self._range_hist], dtype=np.float64)
        ranges = np.array([r for _, r in self._range_hist], dtype=np.float64)

        # Нормализация времени для числовой устойчивости
        t0 = times[0]
        times = times - t0
        dt_span = times[-1] - times[0]
        if dt_span < 0.02:
            return 0.0

        # Линейная регрессия: range = a*t + b → closing_vel = -a
        n = len(times)
        sum_t = np.sum(times)
        sum_r = np.sum(ranges)
        sum_tt = np.sum(times * times)
        sum_tr = np.sum(times * ranges)

        denom = n * sum_tt - sum_t * sum_t
        if abs(denom) < 1e-12:
            return 0.0

        a = (n * sum_tr - sum_t * sum_r) / denom
        return float(-a)  # отрицательный наклон = сближение

    def _compute_tau_from_area(self):
        """
        Time-to-contact из скорости роста площади bbox.
        tau = area / (2 * dArea/dt)

        Физика: для объекта приближающегося со скоростью V на расстоянии R,
        площадь изображения ∝ 1/R², поэтому dA/dt ∝ 2V/R³,
        и tau ≈ A / (2 * dA/dt).
        """
        if len(self._area_hist) < 3:
            return float('inf')

        times = np.array([t for t, _ in self._area_hist], dtype=np.float64)
        areas = np.array([a for _, a in self._area_hist], dtype=np.float64)

        t0 = times[0]
        times = times - t0
        dt_span = times[-1] - times[0]
        if dt_span < 0.02:
            return float('inf')

        # Линейная регрессия для dArea/dt
        n = len(times)
        sum_t = np.sum(times)
        sum_a = np.sum(areas)
        sum_tt = np.sum(times * times)
        sum_ta = np.sum(times * areas)

        denom = n * sum_tt - sum_t * sum_t
        if abs(denom) < 1e-12:
            return float('inf')

        area_rate = (n * sum_ta - sum_t * sum_a) / denom  # dA/dt

        if area_rate < float(TAU_MIN_AREA_RATE):
            return float('inf')  # не сближаемся или слишком медленно

        current_area = float(areas[-1])
        if current_area < 1.0:
            return float('inf')

        tau = current_area / (2.0 * area_rate)
        return max(0.0, float(tau))

    def _compute_confidence(self, apparent_px, confirmed):
        """Confidence оценки дальности: 0..1"""
        conf = 0.0

        # Чем больше цель на экране — тем точнее оценка
        size_conf = float(clamp((apparent_px - 4.0) / 50.0, 0.0, 0.4))
        conf += size_conf

        # Если трек подтверждён
        if confirmed:
            conf += 0.3

        # Если есть достаточная история
        hist_conf = float(clamp(len(self._range_hist) / 6.0, 0.0, 0.2))
        conf += hist_conf

        # Стабильность оценки (мало шума)
        if len(self._range_hist) >= 4:
            ranges = [r for _, r in self._range_hist]
            std = float(np.std(ranges))
            mean = float(np.mean(ranges))
            if mean > 0.1:
                cv = std / mean  # coefficient of variation
                stability = float(clamp(1.0 - cv * 5.0, 0.0, 0.1))
                conf += stability

        return float(clamp(conf, 0.0, 1.0))

    def _classify_phase(self):
        """Фаза сближения по дальности и tau."""
        tau = self.tau if self.tau != float('inf') else 999.0
        range_m = self.range_m if self.range_m is not None else 999.0

        if tau <= float(TAU_CRITICAL_SEC) or range_m <= float(FSM_TERMINAL_RANGE_M):
            return "critical"
        if tau <= float(TAU_WARN_SEC) or range_m <= float(FSM_INTERCEPT_RANGE_M):
            return "warn"
        if range_m <= float(FSM_ACQUIRE_RANGE_M):
            return "approach"
        return "far"

    @staticmethod
    def _idle_result():
        return {
            "range_m": None,
            "range_raw": None,
            "closing_vel": 0.0,
            "tau": float('inf'),
            "tau_raw": float('inf'),
            "confidence": 0.0,
            "phase": "far",
            "apparent_size_px": 0.0,
            "area_px": 0.0,
        }

    def status_line(self):
        if not self.enabled:
            return "Range estimation disabled"
        return (
            f"Range estimation ready: focal={self.focal_px:.0f}px "
            f"target_size={self.target_size_m:.2f}m"
        )

    def draw_overlay(self, frame_bgr, state, sx, sy):
        """Рисует информацию о дальности на кадре."""
        if not self.enabled or not DRAW_RANGE_INFO:
            return

        import cv2

        range_m = state.get("range_m")
        tau = state.get("tau", float('inf'))
        closing = state.get("closing_vel", 0.0)
        phase = state.get("phase", "far")
        conf = state.get("confidence", 0.0)

        # Цвет по фазе
        colors = {
            "far": (180, 180, 180),
            "approach": (0, 220, 255),
            "warn": (0, 165, 255),
            "critical": (0, 0, 255),
        }
        color = colors.get(phase, (180, 180, 180))

        y0 = 200
        lines = []

        if range_m is not None:
            lines.append(f"RNG: {range_m:.1f}m  Vc={closing:+.1f}m/s")

        if tau < 100.0:
            lines.append(f"TAU: {tau:.2f}s  phase={phase}  conf={conf:.2f}")
        else:
            lines.append(f"TAU: ---  phase={phase}  conf={conf:.2f}")

        for i, line in enumerate(lines):
            cv2.putText(
                frame_bgr, line,
                (20, y0 + i * 28),
                cv2.FONT_HERSHEY_SIMPLEX, 0.60, color, 2,
            )
