# ============================================================
# trackers_safe.py
# SafeKalman8D — drop-in замена для Kalman8D.
#
# Фиксит критический баг из логов v4: цель уходит вдаль,
# реально уменьшается по высоте, Kalman оценивает отрицательный
# vh, и без YOLO measurement updates экстраполирует h → 0.
# Bbox вырождается в 73×4, прилипает к линии горизонта, держится
# там 12 секунд, весь трек гибнет.
#
# Защита строится на трёх принципах:
#
# 1. ИЗМЕРЕНИЯМ ДОВЕРЯЕМ. Если YOLO говорит "цель 18×9" — бокс
#    такой и становится. Реальная усадка не ограничивается.
#
# 2. ПРЕДСКАЗАНИЯМ — НЕТ. Между measurement'ами Kalman может
#    сжимать bbox только с ограниченной скоростью (3.5% за
#    кадр максимум). Это физически реалистично для удаляющейся
#    цели и не даёт drift'у обрушить размер за 15 кадров.
#
# 3. ЖЁСТКИЕ ФИЗИЧЕСКИЕ ПРЕДЕЛЫ. Aspect ratio всегда 0.40–3.5
#    (нормальные пропорции дрона), минимальная сторона 8px
#    (ниже трек бессмысленен), максимальная сторона 400px.
#
# Использование:
#   было:   from trackers import Kalman8D
#           kf = Kalman8D()
#   стало:  from trackers_safe import SafeKalman8D
#           kf = SafeKalman8D()
#
# Никаких других изменений в main.py не требуется.
# ============================================================

import numpy as np

from trackers import Kalman8D


# ─── Физические пределы bbox ─────────────────────────────────
SAFE_KF_MIN_SIDE     = 6.0       # минимальная сторона (пиксели) — снижено для далёких самолётов
SAFE_KF_MAX_SIDE     = 400.0     # максимальная сторона
SAFE_KF_MAX_ASPECT   = 6.0       # w/h максимум — самолёт в профиль с размахом крыла
SAFE_KF_MIN_ASPECT   = 0.30      # w/h минимум — самолёт хвостом или узкий ракурс

# ─── Ограничение скорости сжатия предсказанием ──────────────
# Между measurement'ами bbox может сжиматься/расти максимум на
# эту долю за кадр. Ослаблено для быстро маневрирующих целей:
#   0.945 = до 5.5% сжатия за кадр (было 3.5%)
#   1.060 = до 6.0% роста за кадр (было 4.0%)
SAFE_KF_MIN_SHRINK_PER_FRAME = 0.945
SAFE_KF_MAX_GROW_PER_FRAME   = 1.060

# ─── Velocity decay для долгих потерь измерений ─────────────
# После N predict'ов без update начинаем гасить velocity
# каждый кадр. Защита от "улёта" Kalman'а далеко от реального
# положения цели (особенно актуально для мелких далёких целей
# где YOLO не детектит регулярно).
SAFE_KF_DECAY_START_FRAMES = 8       # начинать после 8 кадров без update
SAFE_KF_VELOCITY_DECAY     = 0.97    # 3% гашения за кадр после порога


class SafeKalman8D(Kalman8D):
    """
    Kalman8D с sanity constraint'ами. Полностью API-совместим
    с оригиналом.

    Дополнительные поля для диагностики:
        self.sanity_reject_count — общий счётчик срабатываний
        self.sanity_last_reason  — причина последнего срабатывания
        self.last_meas_w, last_meas_h — размер последнего
            принятого measurement (для отладки)
    """

    def __init__(self):
        super().__init__()
        self.sanity_reject_count = 0
        self.sanity_last_reason = ""
        self.last_meas_w = None
        self.last_meas_h = None
        # Сохраняем w, h после последнего update() — от этого
        # значения отсчитывается допустимая скорость сжатия.
        # Это "наше последнее подтверждённое знание о размере".
        self._anchor_w = None
        self._anchor_h = None
        # Сколько predict'ов прошло с последнего measurement
        self._frames_since_update = 0

    # ─── Override: init_from_box ─────────────────────────────
    def init_from_box(self, box):
        if box is None:
            return
        box = np.asarray(box, dtype=np.float32)
        box = self._force_valid_box(box)
        super().init_from_box(box)

        w = float(self.x[2, 0])
        h = float(self.x[3, 0])
        self._anchor_w = w
        self._anchor_h = h
        self.last_meas_w = w
        self.last_meas_h = h
        self._frames_since_update = 0

    # ─── Override: predict ───────────────────────────────────
    def predict(self, dt, q_scale=1.0):
        """
        Predict + shrink-rate limit + aspect ratio constraint.
        Плюс velocity decay: при долгом отсутствии measurements
        постепенно гасим скорость, иначе Kalman "улетает" в
        сторону по инерции (особенно для мелких далёких целей).
        """
        # Считаем старые w, h ДО predict
        prev_w = float(self.x[2, 0])
        prev_h = float(self.x[3, 0])

        # ─── Velocity decay при долгой потере measurements ─
        if self._frames_since_update >= SAFE_KF_DECAY_START_FRAMES:
            decay = float(SAFE_KF_VELOCITY_DECAY)
            self.x[4, 0] *= decay   # vx
            self.x[5, 0] *= decay   # vy
            self.x[6, 0] *= decay   # vw
            self.x[7, 0] *= decay   # vh

        # Стандартный Kalman predict
        super().predict(dt, q_scale=q_scale)
        self._frames_since_update += 1

        new_w = float(self.x[2, 0])
        new_h = float(self.x[3, 0])

        # ─── Ограничение скорости сжатия ────────────────────
        # Если Kalman хочет сжать bbox быстрее допустимого —
        # откатываем к допустимому пределу и обнуляем velocity.
        min_allowed_w = prev_w * SAFE_KF_MIN_SHRINK_PER_FRAME
        min_allowed_h = prev_h * SAFE_KF_MIN_SHRINK_PER_FRAME
        max_allowed_w = prev_w * SAFE_KF_MAX_GROW_PER_FRAME
        max_allowed_h = prev_h * SAFE_KF_MAX_GROW_PER_FRAME

        rate_clamped = False
        if new_w < min_allowed_w:
            new_w = min_allowed_w
            self.x[6, 0] = 0.0   # обнуляем vw
            rate_clamped = True
        elif new_w > max_allowed_w:
            new_w = max_allowed_w
            self.x[6, 0] = 0.0
            rate_clamped = True

        if new_h < min_allowed_h:
            new_h = min_allowed_h
            self.x[7, 0] = 0.0   # обнуляем vh
            rate_clamped = True
        elif new_h > max_allowed_h:
            new_h = max_allowed_h
            self.x[7, 0] = 0.0
            rate_clamped = True

        # ─── Абсолютные пределы + aspect ratio ──────────────
        new_w, new_h, absolute_clamped = self._clamp_absolute(new_w, new_h)

        # Записываем обратно в состояние
        self.x[2, 0] = new_w
        self.x[3, 0] = new_h

        if rate_clamped:
            self.sanity_reject_count += 1
            self.sanity_last_reason = "predict:rate_clamp"
        elif absolute_clamped:
            self.sanity_reject_count += 1
            self.sanity_last_reason = "predict:abs_clamp"

        return self.x[:4].flatten()

    # ─── Override: update ────────────────────────────────────
    def update(self, z):
        """
        Measurement update. Измерениям доверяем, но проверяем
        на явную вырожденность — если aspect ratio измерения
        >6:1, это точно не дрон (мусорный детект), пропускаем.
        """
        z = np.asarray(z, dtype=np.float32).reshape(-1)
        if len(z) < 4:
            return

        cx, cy, w, h = float(z[0]), float(z[1]), float(z[2]), float(z[3])

        # Жёсткая проверка на явно битое измерение
        if w < SAFE_KF_MIN_SIDE * 0.5 or h < SAFE_KF_MIN_SIDE * 0.5:
            self.sanity_reject_count += 1
            self.sanity_last_reason = f"meas:too_small_{w:.0f}x{h:.0f}"
            return

        if w > SAFE_KF_MAX_SIDE or h > SAFE_KF_MAX_SIDE:
            self.sanity_reject_count += 1
            self.sanity_last_reason = f"meas:too_big_{w:.0f}x{h:.0f}"
            return

        ar = w / max(1e-3, h)
        # Для measurement пределы шире (YOLO на ракурсе сбоку может дать до 7:1 для самолёта)
        if ar > 8.0 or ar < 0.18:
            self.sanity_reject_count += 1
            self.sanity_last_reason = f"meas:bad_ar_{ar:.1f}"
            return

        # Применяем measurement
        z_clean = np.array([cx, cy, w, h], dtype=np.float32)
        super().update(z_clean)

        # Обновляем anchor — это наш "последний достоверный размер"
        self._anchor_w = float(self.x[2, 0])
        self._anchor_h = float(self.x[3, 0])
        self.last_meas_w = w
        self.last_meas_h = h
        self._frames_since_update = 0

    # ─── Override: to_box ────────────────────────────────────
    def to_box(self):
        """
        Последняя линия обороны: абсолютные пределы + aspect ratio.
        Синхронизируем state с результатом clamp'а.
        """
        cx = float(self.x[0, 0])
        cy = float(self.x[1, 0])
        w = float(self.x[2, 0])
        h = float(self.x[3, 0])

        w, h, clamped = self._clamp_absolute(w, h)

        self.x[2, 0] = w
        self.x[3, 0] = h

        return np.array(
            [cx - w * 0.5, cy - h * 0.5, cx + w * 0.5, cy + h * 0.5],
            dtype=np.float32,
        )

    # ─── Internal helpers ────────────────────────────────────
    def _clamp_absolute(self, w, h):
        """Абсолютные пределы + aspect ratio. (w, h, clamped_flag)."""
        clamped = False

        if w < SAFE_KF_MIN_SIDE:
            w = SAFE_KF_MIN_SIDE
            clamped = True
        if h < SAFE_KF_MIN_SIDE:
            h = SAFE_KF_MIN_SIDE
            clamped = True
        if w > SAFE_KF_MAX_SIDE:
            w = SAFE_KF_MAX_SIDE
            clamped = True
        if h > SAFE_KF_MAX_SIDE:
            h = SAFE_KF_MAX_SIDE
            clamped = True

        # Aspect ratio clamp
        ar = w / max(1e-3, h)
        if ar > SAFE_KF_MAX_ASPECT:
            # Слишком широкий → поднимаем h
            h = w / SAFE_KF_MAX_ASPECT
            clamped = True
        elif ar < SAFE_KF_MIN_ASPECT:
            # Слишком высокий → поднимаем w
            w = h * SAFE_KF_MIN_ASPECT
            clamped = True

        return w, h, clamped

    def _force_valid_box(self, box):
        """Принудительно делает box валидным — используется в init."""
        x1, y1, x2, y2 = [float(v) for v in box]
        cx = 0.5 * (x1 + x2)
        cy = 0.5 * (y1 + y2)
        w = max(SAFE_KF_MIN_SIDE, x2 - x1)
        h = max(SAFE_KF_MIN_SIDE, y2 - y1)
        w, h, _ = self._clamp_absolute(w, h)
        return np.array(
            [cx - w * 0.5, cy - h * 0.5, cx + w * 0.5, cy + h * 0.5],
            dtype=np.float32,
        )

    # ─── Status ──────────────────────────────────────────────
    def status_line(self):
        anchor_w = self._anchor_w if self._anchor_w is not None else 0.0
        anchor_h = self._anchor_h if self._anchor_h is not None else 0.0
        return (
            f"SafeKalman8D: anchor={anchor_w:.0f}x{anchor_h:.0f} "
            f"frames_no_meas={self._frames_since_update} "
            f"rejects={self.sanity_reject_count} "
            f"last={self.sanity_last_reason}"
        )
