# ============================================================
# motion_saliency.py
# Разделение независимого движения и ego-motion камеры.
#
# Поддерживает ОБА типа ego-motion:
#   - affine 2x3  (через cv2.warpAffine)
#   - homography 3x3 (через cv2.warpPerspective)
#
# Кроме того поддерживает режим "dual residual": считается
# residual от обеих моделей одновременно, и в каждой точке
# берётся МИНИМУМ — если хотя бы одна модель объясняет движение
# пикселя, он считается статичным. Устойчиво к вырождению любой
# из моделей на сложных сценах.
# ============================================================

import cv2
import numpy as np

from config import *
from helpers import clamp, clip_box


# ─────────────────────────────────────────────────────────────
# Параметры
# ─────────────────────────────────────────────────────────────
MS_ENABLE               = True
MS_BLUR_KSIZE           = 5
MS_DIFF_THR             = 18
MS_MORPH_OPEN_KSIZE     = 3
MS_MORPH_CLOSE_KSIZE    = 5
MS_EMA_ALPHA            = 0.55
MS_OUTLIER_SIGMA_PX     = 18.0
MS_OUTLIER_WEIGHT       = 0.55
MS_DIFF_WEIGHT          = 0.45
MS_MIN_WARP_SHIFT_PX    = 0.3
MS_MASK_TOP_FRAC        = 0.0
MS_MASK_BOTTOM_FRAC     = 0.0
MS_DRAW_OVERLAY         = True
MS_DRAW_ALPHA           = 0.35

# Dual-residual mode: если True, motion_sal.update_dual() считает
# обе карты и берёт min — нужно если включаете и homography, и
# affine параллельно (устойчивее к вырождению).
MS_DUAL_MODE_DEFAULT    = False


def _warp_prev(prev_gray, M, kind, out_size):
    """Универсальный варп предыдущего кадра под текущий."""
    w, h = out_size
    if M is None or kind == "none":
        return prev_gray

    if kind == "affine":
        tx = float(M[0, 2])
        ty = float(M[1, 2])
        if abs(tx) + abs(ty) < float(MS_MIN_WARP_SHIFT_PX):
            return prev_gray
        return cv2.warpAffine(
            prev_gray, M, (w, h),
            flags=cv2.INTER_LINEAR,
            borderMode=cv2.BORDER_REPLICATE,
        )

    if kind == "homography":
        return cv2.warpPerspective(
            prev_gray, M, (w, h),
            flags=cv2.INTER_LINEAR,
            borderMode=cv2.BORDER_REPLICATE,
        )

    return prev_gray


class MotionSaliency:
    """
    Карта независимого движения на кадре.

    Базовое использование:
        ms = MotionSaliency()
        ms.update(gray, prev_gray, M, kind="homography", outlier_pts=pts)
        score = ms.score_box(box)
        rois = ms.get_rois(thr=0.35)

    Dual-residual (устойчивый на сложных сценах):
        ms.update_dual(gray, prev_gray,
                       H=homography, A=affine,
                       outlier_pts=pts)
    """

    def __init__(self):
        self.enabled = bool(MS_ENABLE)
        self.saliency = None
        self.raw_diff = None
        self.outlier_density = None
        self.last_active_frac = 0.0
        self.last_kind = "none"
        self._ema_alpha = float(MS_EMA_ALPHA)

    def reset(self):
        self.saliency = None
        self.raw_diff = None
        self.outlier_density = None
        self.last_active_frac = 0.0
        self.last_kind = "none"

    # ─── Основной update ─────────────────────────────────────
    def update(self, gray, prev_gray, M, kind="affine", outlier_pts=None):
        """
        M    — матрица ego-motion (2x3 для affine, 3x3 для homography)
        kind — "affine" | "homography" | "none"
        """
        if not self.enabled or gray is None:
            return None

        h, w = gray.shape[:2]
        self.last_kind = kind

        diff_map = self._compute_diff_map(gray, prev_gray, M, kind)
        outlier_map = self._compute_outlier_density(outlier_pts, h, w)

        combined = self._combine(diff_map, outlier_map, h, w)
        self._apply_ema(combined)
        self.last_active_frac = float((self.saliency > 0.25).mean())
        return self.saliency

    # ─── Dual residual update ────────────────────────────────
    def update_dual(self, gray, prev_gray, H=None, A=None, outlier_pts=None):
        """
        Считает residual от homography и affine одновременно, в каждой
        точке берёт min (если любая модель объясняет — пиксель статичен).
        Устойчивее на сценах с несколькими плоскостями.

        H — homography 3x3 или None
        A — affine 2x3 или None
        """
        if not self.enabled or gray is None:
            return None

        h, w = gray.shape[:2]

        diff_h = self._compute_diff_map(gray, prev_gray, H, "homography") if H is not None else None
        diff_a = self._compute_diff_map(gray, prev_gray, A, "affine") if A is not None else None

        if diff_h is not None and diff_a is not None:
            diff_map = np.minimum(diff_h, diff_a)
            self.last_kind = "dual"
        elif diff_h is not None:
            diff_map = diff_h
            self.last_kind = "homography"
        elif diff_a is not None:
            diff_map = diff_a
            self.last_kind = "affine"
        else:
            diff_map = None
            self.last_kind = "none"

        outlier_map = self._compute_outlier_density(outlier_pts, h, w)
        combined = self._combine(diff_map, outlier_map, h, w)
        self._apply_ema(combined)
        self.last_active_frac = float((self.saliency > 0.25).mean())
        return self.saliency

    # ─── Scoring API ─────────────────────────────────────────
    def score_box(self, box):
        if self.saliency is None or box is None:
            return 0.0
        h, w = self.saliency.shape[:2]
        x1 = int(clamp(box[0], 0, w - 1))
        y1 = int(clamp(box[1], 0, h - 1))
        x2 = int(clamp(box[2], x1 + 1, w))
        y2 = int(clamp(box[3], y1 + 1, h))
        patch = self.saliency[y1:y2, x1:x2]
        if patch.size == 0:
            return 0.0
        return float(patch.mean())

    def peak_box(self, box):
        if self.saliency is None or box is None:
            return 0.0
        h, w = self.saliency.shape[:2]
        x1 = int(clamp(box[0], 0, w - 1))
        y1 = int(clamp(box[1], 0, h - 1))
        x2 = int(clamp(box[2], x1 + 1, w))
        y2 = int(clamp(box[3], y1 + 1, h))
        patch = self.saliency[y1:y2, x1:x2]
        if patch.size == 0:
            return 0.0
        return float(patch.max())

    def get_rois(self, thr=0.35, min_area=30, max_rois=6, pad_px=12):
        if self.saliency is None:
            return []
        mask = (self.saliency >= float(thr)).astype(np.uint8)
        if mask.sum() < min_area:
            return []

        n_labels, labels, stats, _ = cv2.connectedComponentsWithStats(mask, connectivity=8)
        rois = []
        h, w = mask.shape[:2]
        for lbl in range(1, int(n_labels)):
            area = int(stats[lbl, cv2.CC_STAT_AREA])
            if area < min_area:
                continue
            x = int(stats[lbl, cv2.CC_STAT_LEFT])
            y = int(stats[lbl, cv2.CC_STAT_TOP])
            ww = int(stats[lbl, cv2.CC_STAT_WIDTH])
            hh = int(stats[lbl, cv2.CC_STAT_HEIGHT])
            box = np.array(
                [x - pad_px, y - pad_px, x + ww + pad_px, y + hh + pad_px],
                dtype=np.float32,
            )
            box = clip_box(box, w, h)
            rois.append((area, box))

        rois.sort(key=lambda x: x[0], reverse=True)
        return [b for _, b in rois[:max_rois]]

    # ─── Internals ───────────────────────────────────────────
    def _compute_diff_map(self, gray, prev_gray, M, kind):
        """Ego-compensated frame differencing → float32 [0..1]."""
        if prev_gray is None or gray.shape != prev_gray.shape:
            return None

        k = int(MS_BLUR_KSIZE)
        if k >= 3 and (k % 2) == 1:
            g_cur = cv2.GaussianBlur(gray, (k, k), 0)
            g_prv = cv2.GaussianBlur(prev_gray, (k, k), 0)
        else:
            g_cur = gray
            g_prv = prev_gray

        h, w = g_cur.shape[:2]
        warped = _warp_prev(g_prv, M, kind, (w, h))

        diff = cv2.absdiff(g_cur, warped)
        _, bin_mask = cv2.threshold(diff, int(MS_DIFF_THR), 255, cv2.THRESH_BINARY)

        ko = int(MS_MORPH_OPEN_KSIZE)
        if ko >= 2:
            kern = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (ko, ko))
            bin_mask = cv2.morphologyEx(bin_mask, cv2.MORPH_OPEN, kern)

        kc = int(MS_MORPH_CLOSE_KSIZE)
        if kc >= 2:
            kern = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kc, kc))
            bin_mask = cv2.morphologyEx(bin_mask, cv2.MORPH_CLOSE, kern)

        amp = diff.astype(np.float32) / 255.0
        gate = (bin_mask > 0).astype(np.float32)
        soft = cv2.GaussianBlur(amp * gate, (9, 9), 0)

        mx = float(soft.max())
        if mx > 1e-6:
            soft = soft / mx
        self.raw_diff = soft
        return soft

    def _compute_outlier_density(self, outlier_pts, h, w):
        """KDE density из outlier точек."""
        if outlier_pts is None or len(outlier_pts) == 0:
            self.outlier_density = None
            return None

        scale = 0.25
        sh = max(1, int(h * scale))
        sw = max(1, int(w * scale))
        acc = np.zeros((sh, sw), dtype=np.float32)

        pts = np.asarray(outlier_pts, dtype=np.float32).reshape(-1, 2)
        for x, y in pts:
            xi = int(clamp(x * scale, 0, sw - 1))
            yi = int(clamp(y * scale, 0, sh - 1))
            acc[yi, xi] += 1.0

        sigma = float(MS_OUTLIER_SIGMA_PX) * scale
        k = int(max(3, round(sigma * 6)) | 1)
        blur = cv2.GaussianBlur(acc, (k, k), sigma)

        density = cv2.resize(blur, (w, h), interpolation=cv2.INTER_LINEAR)
        mx = float(density.max())
        if mx > 1e-6:
            density = density / mx
        self.outlier_density = density
        return density

    def _combine(self, diff_map, outlier_map, h, w):
        combined = np.zeros((h, w), dtype=np.float32)
        if diff_map is not None:
            combined += float(MS_DIFF_WEIGHT) * diff_map
        if outlier_map is not None:
            combined += float(MS_OUTLIER_WEIGHT) * outlier_map
        combined = np.clip(combined, 0.0, 1.0)

        if MS_MASK_TOP_FRAC > 0.0:
            combined[: int(h * MS_MASK_TOP_FRAC), :] = 0.0
        if MS_MASK_BOTTOM_FRAC > 0.0:
            combined[int(h * (1.0 - MS_MASK_BOTTOM_FRAC)) :, :] = 0.0
        return combined

    def _apply_ema(self, combined):
        if self.saliency is None or self.saliency.shape != combined.shape:
            self.saliency = combined
        else:
            a = self._ema_alpha
            self.saliency = a * self.saliency + (1.0 - a) * combined

    # ─── Overlay ─────────────────────────────────────────────
    def draw_overlay(self, frame_bgr, sx=1.0, sy=1.0):
        if not MS_DRAW_OVERLAY or self.saliency is None:
            return
        h, w = frame_bgr.shape[:2]
        sal = cv2.resize(self.saliency, (w, h), interpolation=cv2.INTER_LINEAR)
        sal_u8 = np.clip(sal * 255.0, 0, 255).astype(np.uint8)
        heat = cv2.applyColorMap(sal_u8, cv2.COLORMAP_JET)
        cv2.addWeighted(heat, float(MS_DRAW_ALPHA), frame_bgr, 1.0 - float(MS_DRAW_ALPHA), 0, frame_bgr)

        cv2.putText(
            frame_bgr,
            f"MS[{self.last_kind}] active={self.last_active_frac*100:.1f}%",
            (20, 230),
            cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255, 255, 255), 2,
        )

    def status_line(self):
        if not self.enabled:
            return "Motion saliency disabled"
        return (
            f"Motion saliency ready: diff_w={MS_DIFF_WEIGHT} "
            f"outlier_w={MS_OUTLIER_WEIGHT} ema={MS_EMA_ALPHA}"
        )