# ============================================================
# trackers_hybrid.py  (v2 — affine motion for maneuvering targets)
#
# Ключевые отличия от v1:
#   1. KLT путь оценивает ПОЛНУЮ аффинку (translation + rotation
#      + scale) через cv2.estimateAffinePartial2D, а не только
#      медиану сдвигов. Это необходимо когда цель поворачивается
#      в кадре — при rigid translation inlier-точки разлетаются
#      и квалити падает.
#   2. Применение аффинки к центру и к размеру bbox с clamp'ом
#      на изменение масштаба за кадр (защита от RANSAC-сбоев).
#   3. Более частая переинициализация features (8 вместо 15)
#      для быстрого подстраивания под вращающуюся цель.
#   4. Phase correlation теперь работает на НЕСКОЛЬКИХ масштабах
#      шаблона (0.9x, 1.0x, 1.1x) чтобы устойчиво ловить цель при
#      приближении/удалении.
# ============================================================

import cv2
import numpy as np

from config import *
from helpers import clamp


# ─── Параметры ───────────────────────────────────────────────
HYBRID_CLAHE_CLIP         = 2.5
HYBRID_CLAHE_TILE         = 4
HYBRID_KLT_QUALITY        = 0.005
HYBRID_KLT_QUALITY_RETRY  = 0.002
HYBRID_KLT_MIN_DIST       = 2
HYBRID_KLT_MAX_CORNERS    = 50
HYBRID_KLT_MIN_INLIERS    = 3
HYBRID_KLT_MIN_FOR_AFFINE = 6        # меньше — fallback на median translation
HYBRID_KLT_OUTLIER_THR    = 3.5
HYBRID_PAD_FACTOR         = 0.25
HYBRID_SEARCH_FACTOR      = 2.4      # чуть больше search window для манёвров

HYBRID_PHASE_MIN_RESP     = 0.22     # чуть ослаблен
HYBRID_TM_MIN_SCORE       = 0.42
HYBRID_TM_SCALES          = (0.9, 1.0, 1.1)

# Защита от RANSAC-сбоев (sanity check, размер bbox НЕ обновляется)
HYBRID_MAX_SCALE_PER_FRAME = 1.25    # допустимый диапазон нормы аффинки
HYBRID_MIN_SCALE_PER_FRAME = 0.80    # иначе считаем RANSAC сбоем
HYBRID_MAX_TRANSLATION_PX  = 80.0    # максимальный сдвиг центра за кадр

HYBRID_TEMPLATE_UPDATE_EMA = 0.80    # чуть быстрее обновлять (было 0.85)
HYBRID_REINIT_FEATURES_EVERY = 8     # было 15 — для динамичной цели


class HybridTracker:
    """
    Многоуровневый трекер для мелких манёвренных целей.

    API:
        trk = HybridTracker()
        trk.init(frame_bgr, box)
        new_box = trk.update(frame_bgr)   # None если потеряли
    Поля:
        self.box           — текущий bbox
        self.quality       — 0..1
        self.last_method   — "klt-affine" / "klt-translate" / "phasecorr" / "template" / "lost"
        self.good_count    — число inliers KLT на последнем кадре
    """

    def __init__(self):
        self._clahe = cv2.createCLAHE(
            clipLimit=float(HYBRID_CLAHE_CLIP),
            tileGridSize=(int(HYBRID_CLAHE_TILE), int(HYBRID_CLAHE_TILE)),
        )
        self.reset()

    # ─── Public API ──────────────────────────────────────────
    def reset(self):
        self.prev_gray = None
        self.prev_gray_eq = None
        self.pts = None
        self.box = None
        self.template = None
        self.good_count = 0
        self.quality = 0.0
        self.last_method = "none"
        self._frames_since_feature_init = 0
        self._success_streak = 0

    def init(self, frame_bgr, box):
        if frame_bgr is None or box is None:
            return
        gray = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2GRAY)
        gray_eq = self._clahe.apply(gray)

        self.box = np.array(box, dtype=np.float32)
        self.prev_gray = gray
        self.prev_gray_eq = gray_eq

        self._extract_features(gray_eq, self.box)
        self._update_template(gray_eq, self.box, reset=True)
        self._frames_since_feature_init = 0
        self._success_streak = 0
        self.last_method = "init"

    def update(self, frame_bgr):
        if self.box is None or self.prev_gray is None:
            return None

        gray = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2GRAY)
        gray_eq = self._clahe.apply(gray)

        # ── Попытка 1: KLT (affine или fallback translate) ──
        new_box_klt, q_klt, klt_method = self._try_klt(gray, gray_eq)

        # ── Попытка 2: Phase correlation ────────────────────
        new_box_pc, q_pc = (None, 0.0)
        if new_box_klt is None or q_klt < 0.25:
            new_box_pc, q_pc = self._try_phase_correlation(gray_eq)

        # ── Попытка 3: Template matching (multi-scale) ──────
        new_box_tm, q_tm = (None, 0.0)
        if new_box_klt is None and new_box_pc is None:
            new_box_tm, q_tm = self._try_template_match(gray_eq)

        # ── Выбор лучшего ─────────────────────────────────────
        candidates = []
        if new_box_klt is not None:
            candidates.append((klt_method, new_box_klt, q_klt))
        if new_box_pc is not None:
            candidates.append(("phasecorr", new_box_pc, q_pc))
        if new_box_tm is not None:
            candidates.append(("template", new_box_tm, q_tm))

        if not candidates:
            self.last_method = "lost"
            self.quality = 0.0
            self.prev_gray = gray
            self.prev_gray_eq = gray_eq
            return None

        candidates.sort(key=lambda c: c[2], reverse=True)
        method, new_box, q = candidates[0]

        self.box = new_box
        self.quality = float(q)
        self.last_method = method
        self._success_streak += 1
        self._frames_since_feature_init += 1

        # Переинициализация features и обновление template
        if self._frames_since_feature_init >= int(HYBRID_REINIT_FEATURES_EVERY):
            self._extract_features(gray_eq, self.box)
            self._update_template(gray_eq, self.box, reset=False)
            self._frames_since_feature_init = 0
        elif method in ("phasecorr", "template") and self._success_streak % 4 == 0:
            self._extract_features(gray_eq, self.box)
            self._update_template(gray_eq, self.box, reset=False)

        self.prev_gray = gray
        self.prev_gray_eq = gray_eq
        return self.box

    # ─── Level 1a: KLT with AFFINE motion ────────────────────
    def _try_klt(self, gray, gray_eq):
        """
        Returns (new_box, quality, method_name).
        method_name ∈ {"klt-affine", "klt-translate"} или (None, 0, "none").
        """
        if self.pts is None or len(self.pts) == 0:
            return None, 0.0, "none"

        next_pts, st, err = cv2.calcOpticalFlowPyrLK(
            self.prev_gray_eq, gray_eq, self.pts, None,
            winSize=(21, 21),
            maxLevel=3,
            criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 30, 0.01),
        )
        if next_pts is None or st is None:
            self.good_count = 0
            return None, 0.0, "none"

        st = st.reshape(-1).astype(bool)
        good_old = self.pts[st].reshape(-1, 2)
        good_new = next_pts[st].reshape(-1, 2)
        self.good_count = int(len(good_new))

        if self.good_count < int(HYBRID_KLT_MIN_INLIERS):
            return None, 0.0, "none"

        # ── Путь А: affine через estimateAffinePartial2D (RANSAC) ──
        if self.good_count >= int(HYBRID_KLT_MIN_FOR_AFFINE):
            M, inlier_mask = cv2.estimateAffinePartial2D(
                good_old, good_new,
                method=cv2.RANSAC,
                ransacReprojThreshold=float(HYBRID_KLT_OUTLIER_THR),
                maxIters=200,
                confidence=0.99,
            )
            if M is not None and inlier_mask is not None:
                inlier_cnt = int(inlier_mask.sum())
                if inlier_cnt >= int(HYBRID_KLT_MIN_INLIERS):
                    new_box = self._apply_affine_to_box(M, self.box)
                    if new_box is not None and self._is_sane_box_update(self.box, new_box):
                        # quality: fraction of inliers * абсолютное число
                        frac = float(inlier_cnt) / max(1, self.good_count)
                        quality = float(clamp(
                            frac * (inlier_cnt / max(6.0, HYBRID_KLT_MAX_CORNERS * 0.4)),
                            0.0, 1.0,
                        ))
                        # Обновляем pts только inliers для следующего кадра
                        mask_flat = inlier_mask.reshape(-1).astype(bool)
                        self.pts = good_new[mask_flat].reshape(-1, 1, 2).astype(np.float32)
                        return new_box, quality, "klt-affine"

        # ── Путь Б: fallback median translation ─────────────
        disp = good_new - good_old
        med = np.median(disp, axis=0)
        d = np.linalg.norm(disp - med[None, :], axis=1)
        inliers = d < float(HYBRID_KLT_OUTLIER_THR)
        in_cnt = int(np.count_nonzero(inliers))

        if in_cnt < int(HYBRID_KLT_MIN_INLIERS):
            return None, 0.0, "none"

        disp_in = disp[inliers]
        med = np.median(disp_in, axis=0)
        dx, dy = float(med[0]), float(med[1])

        if abs(dx) > HYBRID_MAX_TRANSLATION_PX or abs(dy) > HYBRID_MAX_TRANSLATION_PX:
            return None, 0.0, "none"

        x1, y1, x2, y2 = self.box
        new_box = np.array([x1 + dx, y1 + dy, x2 + dx, y2 + dy], dtype=np.float32)

        frac = float(in_cnt) / max(1, self.good_count)
        quality = float(clamp(
            frac * (in_cnt / max(6.0, HYBRID_KLT_MAX_CORNERS * 0.5)),
            0.0, 0.85,   # translate даёт меньший макс quality чем affine
        ))

        self.pts = good_new[inliers].reshape(-1, 1, 2).astype(np.float32)
        return new_box, quality, "klt-translate"

    def _apply_affine_to_box(self, M, box):
        """
        Применяет аффинку ТОЛЬКО к центру bbox (translation + rotation).
        РАЗМЕР bbox НЕ ОБНОВЛЯЕТСЯ — остаётся фиксированным.

        Причина: когда KLT-трекер масштабировал бокс своим scale, он
        постепенно расходился по размеру с тем, что детектирует YOLO,
        в результате ByteTrack не ассоциировал детекты с треком
        (hit_streak падал до 0). Размер бокса должен обновляться
        только от YOLO — это единственный надёжный источник истинного
        масштаба цели.

        scale от аффинки используется только для sanity check: если
        он слишком далёк от 1.0, значит RANSAC сбился, возвращаем None.
        """
        x1, y1, x2, y2 = [float(v) for v in box]
        cx = 0.5 * (x1 + x2)
        cy = 0.5 * (y1 + y2)
        w = x2 - x1
        h = y2 - y1

        # Центр через аффинку
        new_cx = M[0, 0] * cx + M[0, 1] * cy + M[0, 2]
        new_cy = M[1, 0] * cx + M[1, 1] * cy + M[1, 2]

        # Sanity check через норму столбца (должна быть ~1.0 для
        # кадра к кадру; большое отклонение = RANSAC сбой)
        scale = float(np.hypot(M[0, 0], M[1, 0]))
        if not (float(HYBRID_MIN_SCALE_PER_FRAME) <= scale <= float(HYBRID_MAX_SCALE_PER_FRAME)):
            return None

        # Размер НЕ меняем — используем старые w, h
        return np.array(
            [new_cx - w * 0.5, new_cy - h * 0.5,
             new_cx + w * 0.5, new_cy + h * 0.5],
            dtype=np.float32,
        )

    def _is_sane_box_update(self, old_box, new_box):
        """Защита от RANSAC-выбросов."""
        ox1, oy1, ox2, oy2 = [float(v) for v in old_box]
        nx1, ny1, nx2, ny2 = [float(v) for v in new_box]
        ocx = 0.5 * (ox1 + ox2)
        ocy = 0.5 * (oy1 + oy2)
        ncx = 0.5 * (nx1 + nx2)
        ncy = 0.5 * (ny1 + ny2)
        shift = float(np.hypot(ncx - ocx, ncy - ocy))
        if shift > float(HYBRID_MAX_TRANSLATION_PX):
            return False
        return True

    # ─── Level 2: Phase correlation ──────────────────────────
    def _try_phase_correlation(self, gray_eq):
        if self.template is None or self.template.size == 0:
            return None, 0.0

        search, (sx, sy) = self._crop_search_window(gray_eq, self.box)
        if search is None:
            return None, 0.0
        h_t, w_t = self.template.shape[:2]
        h_s, w_s = search.shape[:2]
        if h_s < h_t or w_s < w_t:
            return None, 0.0

        cy_s = h_s // 2
        cx_s = w_s // 2
        y0 = max(0, cy_s - h_t // 2)
        x0 = max(0, cx_s - w_t // 2)
        y1 = min(h_s, y0 + h_t)
        x1 = min(w_s, x0 + w_t)
        patch = search[y0:y1, x0:x1]
        if patch.shape != self.template.shape:
            patch = cv2.copyMakeBorder(
                patch, 0, h_t - patch.shape[0], 0, w_t - patch.shape[1],
                cv2.BORDER_REPLICATE,
            )

        try:
            t_f = self.template.astype(np.float32)
            p_f = patch.astype(np.float32)
            (shift_x, shift_y), response = cv2.phaseCorrelate(t_f, p_f)
        except cv2.error:
            return None, 0.0

        if response < float(HYBRID_PHASE_MIN_RESP):
            return None, 0.0
        if abs(shift_x) > HYBRID_MAX_TRANSLATION_PX or abs(shift_y) > HYBRID_MAX_TRANSLATION_PX:
            return None, 0.0

        x1b, y1b, x2b, y2b = self.box
        new_box = np.array(
            [x1b + shift_x, y1b + shift_y, x2b + shift_x, y2b + shift_y],
            dtype=np.float32,
        )
        quality = float(clamp(response * 0.75, 0.0, 0.75))
        return new_box, quality

    # ─── Level 3: Multi-scale template matching ──────────────
    def _try_template_match(self, gray_eq):
        if self.template is None or self.template.size == 0:
            return None, 0.0

        search, (sx, sy) = self._crop_search_window(gray_eq, self.box)
        if search is None:
            return None, 0.0

        best = None   # (score, new_box)

        for scale in HYBRID_TM_SCALES:
            if scale == 1.0:
                tmpl = self.template
            else:
                nh = max(4, int(self.template.shape[0] * scale))
                nw = max(4, int(self.template.shape[1] * scale))
                tmpl = cv2.resize(self.template, (nw, nh), interpolation=cv2.INTER_LINEAR)

            h_t, w_t = tmpl.shape[:2]
            if search.shape[0] < h_t or search.shape[1] < w_t:
                continue

            try:
                result = cv2.matchTemplate(search, tmpl, cv2.TM_CCOEFF_NORMED)
                min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
            except cv2.error:
                continue

            if max_val < float(HYBRID_TM_MIN_SCORE):
                continue

            mx, my = max_loc
            nx1 = sx + mx
            ny1 = sy + my
            new_box = np.array(
                [nx1, ny1, nx1 + w_t, ny1 + h_t], dtype=np.float32
            )
            if best is None or max_val > best[0]:
                best = (float(max_val), new_box)

        if best is None:
            return None, 0.0

        score, new_box = best
        quality = float(clamp(score * 0.72, 0.0, 0.72))
        return new_box, quality

    # ─── Helpers ─────────────────────────────────────────────
    def _extract_features(self, gray_eq, box):
        h, w = gray_eq.shape[:2]
        x1, y1, x2, y2 = box
        bw = max(1.0, x2 - x1)
        bh = max(1.0, y2 - y1)
        pad_x = bw * float(HYBRID_PAD_FACTOR)
        pad_y = bh * float(HYBRID_PAD_FACTOR)

        mx1 = int(clamp(x1 - pad_x, 0, w - 1))
        my1 = int(clamp(y1 - pad_y, 0, h - 1))
        mx2 = int(clamp(x2 + pad_x, 1, w))
        my2 = int(clamp(y2 + pad_y, 1, h))

        mask = np.zeros_like(gray_eq)
        mask[my1:my2, mx1:mx2] = 255

        pts = cv2.goodFeaturesToTrack(
            gray_eq,
            maxCorners=int(HYBRID_KLT_MAX_CORNERS),
            qualityLevel=float(HYBRID_KLT_QUALITY),
            minDistance=int(HYBRID_KLT_MIN_DIST),
            mask=mask,
        )
        if pts is None or len(pts) < 6:
            pts = cv2.goodFeaturesToTrack(
                gray_eq,
                maxCorners=int(HYBRID_KLT_MAX_CORNERS),
                qualityLevel=float(HYBRID_KLT_QUALITY_RETRY),
                minDistance=int(HYBRID_KLT_MIN_DIST),
                mask=mask,
            )
        self.pts = pts.astype(np.float32) if pts is not None else None

    def _update_template(self, gray_eq, box, reset=False):
        h, w = gray_eq.shape[:2]
        x1 = int(clamp(box[0], 0, w - 1))
        y1 = int(clamp(box[1], 0, h - 1))
        x2 = int(clamp(box[2], x1 + 2, w))
        y2 = int(clamp(box[3], y1 + 2, h))
        crop = gray_eq[y1:y2, x1:x2]
        if crop.size == 0:
            return

        if reset or self.template is None or self.template.shape != crop.shape:
            self.template = crop.copy()
        else:
            a = float(HYBRID_TEMPLATE_UPDATE_EMA)
            self.template = (a * self.template.astype(np.float32)
                             + (1.0 - a) * crop.astype(np.float32)).astype(np.uint8)

    def _crop_search_window(self, gray_eq, box):
        h, w = gray_eq.shape[:2]
        cx = 0.5 * (float(box[0]) + float(box[2]))
        cy = 0.5 * (float(box[1]) + float(box[3]))
        bw = max(8.0, float(box[2] - box[0]))
        bh = max(8.0, float(box[3] - box[1]))

        sw = bw * float(HYBRID_SEARCH_FACTOR)
        sh = bh * float(HYBRID_SEARCH_FACTOR)

        sx = int(clamp(cx - sw * 0.5, 0, w - 1))
        sy = int(clamp(cy - sh * 0.5, 0, h - 1))
        ex = int(clamp(cx + sw * 0.5, sx + 2, w))
        ey = int(clamp(cy + sh * 0.5, sy + 2, h))

        search = gray_eq[sy:ey, sx:ex]
        if search.size == 0:
            return None, (0, 0)
        return search, (sx, sy)