# ============================================================
# imm_filter.py
# Модуль 5: IMM-фильтр (Interacting Multiple Model)
#
# Три модели движения цели:
#   1) CV  — Constant Velocity (прямолинейное)
#   2) CA  — Constant Acceleration (разгон / торможение)
#   3) CT  — Coordinated Turn (вираж)
#
# IMM автоматически определяет, какую модель использовать,
# через вероятности Маркова. Это критично для маневрирующей цели,
# т.к. обычный Kalman (CV) запаздывает при манёврах.
# ============================================================

import numpy as np
from config import *
from config_intercept import *
from helpers import clamp, box_center, box_wh, clip_box


class _KalmanCV:
    """
    Constant Velocity модель.
    Состояние: [cx, cy, w, h, vx, vy, vw, vh]  (8D)
    """

    def __init__(self, q_pos, q_vel):
        self.dim_x = 8
        self.dim_z = 4
        self.x = np.zeros((8, 1), dtype=np.float64)
        self.P = np.eye(8, dtype=np.float64) * 50.0

        self.H = np.zeros((4, 8), dtype=np.float64)
        for i in range(4):
            self.H[i, i] = 1.0

        self.R = np.diag([IMM_R_POS, IMM_R_POS, IMM_R_SIZE, IMM_R_SIZE]).astype(np.float64)
        self.q_pos = float(q_pos)
        self.q_vel = float(q_vel)

    def F(self, dt):
        F = np.eye(8, dtype=np.float64)
        F[0, 4] = dt
        F[1, 5] = dt
        F[2, 6] = dt
        F[3, 7] = dt
        return F

    def Q(self, dt):
        q = np.zeros((8, 8), dtype=np.float64)
        dt2 = dt * dt
        q[0, 0] = self.q_pos * dt2
        q[1, 1] = self.q_pos * dt2
        q[2, 2] = self.q_pos * dt2
        q[3, 3] = self.q_pos * dt2
        q[4, 4] = self.q_vel * dt
        q[5, 5] = self.q_vel * dt
        q[6, 6] = self.q_vel * dt
        q[7, 7] = self.q_vel * dt
        return q

    def predict(self, dt):
        F = self.F(dt)
        self.x = F @ self.x
        self.P = F @ self.P @ F.T + self.Q(dt)

    def update(self, z):
        z = np.array(z, dtype=np.float64).reshape(4, 1)
        y = z - self.H @ self.x
        S = self.H @ self.P @ self.H.T + self.R
        K = self.P @ self.H.T @ np.linalg.inv(S)
        self.x = self.x + K @ y
        I = np.eye(8, dtype=np.float64)
        self.P = (I - K @ self.H) @ self.P
        return y, S

    def likelihood(self, z):
        """Правдоподобие измерения (для IMM mixing)."""
        z = np.array(z, dtype=np.float64).reshape(4, 1)
        y = z - self.H @ self.x
        S = self.H @ self.P @ self.H.T + self.R
        try:
            S_inv = np.linalg.inv(S)
            det_S = max(np.linalg.det(S), 1e-300)
            n = self.dim_z
            exp_val = float(-0.5 * (y.T @ S_inv @ y).item())
            exp_val = max(exp_val, -500.0)  # предотвращение underflow
            norm = (2.0 * np.pi) ** (-n / 2.0) * det_S ** (-0.5)
            return max(norm * np.exp(exp_val), 1e-300)
        except np.linalg.LinAlgError:
            return 1e-300


class _KalmanCA:
    """
    Constant Acceleration модель.
    Состояние: [cx, cy, w, h, vx, vy, ax, ay, vw, vh]  (10D)
    Измерение: [cx, cy, w, h]
    """

    def __init__(self, q_pos, q_vel, q_acc):
        self.dim_x = 10
        self.dim_z = 4
        self.x = np.zeros((10, 1), dtype=np.float64)
        self.P = np.eye(10, dtype=np.float64) * 50.0

        self.H = np.zeros((4, 10), dtype=np.float64)
        self.H[0, 0] = 1.0  # cx
        self.H[1, 1] = 1.0  # cy
        self.H[2, 2] = 1.0  # w
        self.H[3, 3] = 1.0  # h

        self.R = np.diag([IMM_R_POS, IMM_R_POS, IMM_R_SIZE, IMM_R_SIZE]).astype(np.float64)
        self.q_pos = float(q_pos)
        self.q_vel = float(q_vel)
        self.q_acc = float(q_acc)

    def F(self, dt):
        F = np.eye(10, dtype=np.float64)
        dt2 = 0.5 * dt * dt
        F[0, 4] = dt      # cx += vx*dt
        F[0, 6] = dt2     # cx += 0.5*ax*dt²
        F[1, 5] = dt      # cy += vy*dt
        F[1, 7] = dt2     # cy += 0.5*ay*dt²
        F[4, 6] = dt      # vx += ax*dt
        F[5, 7] = dt      # vy += ay*dt
        F[2, 8] = dt      # w += vw*dt
        F[3, 9] = dt      # h += vh*dt
        return F

    def Q(self, dt):
        q = np.zeros((10, 10), dtype=np.float64)
        dt2 = dt * dt
        q[0, 0] = self.q_pos * dt2
        q[1, 1] = self.q_pos * dt2
        q[2, 2] = self.q_pos * dt2
        q[3, 3] = self.q_pos * dt2
        q[4, 4] = self.q_vel * dt
        q[5, 5] = self.q_vel * dt
        q[6, 6] = self.q_acc * dt
        q[7, 7] = self.q_acc * dt
        q[8, 8] = self.q_vel * dt
        q[9, 9] = self.q_vel * dt
        return q

    def predict(self, dt):
        F = self.F(dt)
        self.x = F @ self.x
        self.P = F @ self.P @ F.T + self.Q(dt)

    def update(self, z):
        z = np.array(z, dtype=np.float64).reshape(4, 1)
        y = z - self.H @ self.x
        S = self.H @ self.P @ self.H.T + self.R
        K = self.P @ self.H.T @ np.linalg.inv(S)
        self.x = self.x + K @ y
        I = np.eye(10, dtype=np.float64)
        self.P = (I - K @ self.H) @ self.P
        return y, S

    def likelihood(self, z):
        z = np.array(z, dtype=np.float64).reshape(4, 1)
        y = z - self.H @ self.x
        S = self.H @ self.P @ self.H.T + self.R
        try:
            S_inv = np.linalg.inv(S)
            det_S = max(np.linalg.det(S), 1e-300)
            n = self.dim_z
            exp_val = float(-0.5 * (y.T @ S_inv @ y).item())
            exp_val = max(exp_val, -500.0)
            norm = (2.0 * np.pi) ** (-n / 2.0) * det_S ** (-0.5)
            return max(norm * np.exp(exp_val), 1e-300)
        except np.linalg.LinAlgError:
            return 1e-300


class _KalmanCT:
    """
    Coordinated Turn модель.
    Состояние: [cx, cy, w, h, vx, vy, omega, vw, vh]  (9D)
    omega = угловая скорость виража (рад/с)
    """

    def __init__(self, q_pos, q_vel, q_omega):
        self.dim_x = 9
        self.dim_z = 4
        self.x = np.zeros((9, 1), dtype=np.float64)
        self.P = np.eye(9, dtype=np.float64) * 50.0

        self.H = np.zeros((4, 9), dtype=np.float64)
        self.H[0, 0] = 1.0
        self.H[1, 1] = 1.0
        self.H[2, 2] = 1.0
        self.H[3, 3] = 1.0

        self.R = np.diag([IMM_R_POS, IMM_R_POS, IMM_R_SIZE, IMM_R_SIZE]).astype(np.float64)
        self.q_pos = float(q_pos)
        self.q_vel = float(q_vel)
        self.q_omega = float(q_omega)

    def F(self, dt):
        omega = float(self.x[6, 0])
        F = np.eye(9, dtype=np.float64)

        if abs(omega) < 1e-4:
            # При нулевой omega — линейная модель
            F[0, 4] = dt
            F[1, 5] = dt
        else:
            # Нелинейная CT-модель
            so = np.sin(omega * dt)
            co = np.cos(omega * dt)
            F[0, 4] = so / omega
            F[0, 5] = -(1.0 - co) / omega
            F[1, 4] = (1.0 - co) / omega
            F[1, 5] = so / omega
            F[4, 4] = co
            F[4, 5] = -so
            F[5, 4] = so
            F[5, 5] = co

        F[2, 7] = dt
        F[3, 8] = dt
        return F

    def Q(self, dt):
        q = np.zeros((9, 9), dtype=np.float64)
        dt2 = dt * dt
        q[0, 0] = self.q_pos * dt2
        q[1, 1] = self.q_pos * dt2
        q[2, 2] = self.q_pos * dt2
        q[3, 3] = self.q_pos * dt2
        q[4, 4] = self.q_vel * dt
        q[5, 5] = self.q_vel * dt
        q[6, 6] = self.q_omega * dt
        q[7, 7] = self.q_vel * dt
        q[8, 8] = self.q_vel * dt
        return q

    def predict(self, dt):
        F = self.F(dt)
        self.x = F @ self.x
        self.P = F @ self.P @ F.T + self.Q(dt)

    def update(self, z):
        z = np.array(z, dtype=np.float64).reshape(4, 1)
        y = z - self.H @ self.x
        S = self.H @ self.P @ self.H.T + self.R
        K = self.P @ self.H.T @ np.linalg.inv(S)
        self.x = self.x + K @ y
        I = np.eye(9, dtype=np.float64)
        self.P = (I - K @ self.H) @ self.P
        return y, S

    def likelihood(self, z):
        z = np.array(z, dtype=np.float64).reshape(4, 1)
        y = z - self.H @ self.x
        S = self.H @ self.P @ self.H.T + self.R
        try:
            S_inv = np.linalg.inv(S)
            det_S = max(np.linalg.det(S), 1e-300)
            n = self.dim_z
            exp_val = float(-0.5 * (y.T @ S_inv @ y).item())
            exp_val = max(exp_val, -500.0)
            norm = (2.0 * np.pi) ** (-n / 2.0) * det_S ** (-0.5)
            return max(norm * np.exp(exp_val), 1e-300)
        except np.linalg.LinAlgError:
            return 1e-300


# ─────────────────────────────────────────────────────────────
# IMM FILTER
# ─────────────────────────────────────────────────────────────

class IMMFilter:
    """
    Interacting Multiple Model фильтр.

    Управляет тремя фильтрами Калмана с разными моделями движения
    и переключается между ними через байесовское взвешивание.

    Интерфейс совместим с Kalman8D из trackers.py:
    - init_from_box(box)
    - predict(dt)
    - update(z)  где z = [cx, cy, w, h]
    - to_box() → [x1, y1, x2, y2]
    - uncertainty() → float
    """

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

        # Три модели
        self.models = [
            _KalmanCV(IMM_CV_Q_POS, IMM_CV_Q_VEL),
            _KalmanCA(IMM_CA_Q_POS, IMM_CA_Q_VEL, IMM_CA_Q_ACC),
            _KalmanCT(IMM_CT_Q_POS, IMM_CT_Q_VEL, IMM_CT_Q_OMEGA),
        ]
        self.n_models = len(self.models)
        self.model_names = ["CV", "CA", "CT"]

        # Вероятности моделей
        self.mu = np.array(IMM_INIT_PROBS, dtype=np.float64)
        self.mu /= self.mu.sum()

        # Матрица переходов Маркова
        self.TPM = np.array(IMM_TRANSITION_MATRIX, dtype=np.float64)
        # Нормализация строк
        for i in range(self.n_models):
            self.TPM[i] /= max(self.TPM[i].sum(), 1e-12)

        self.initialized = False

        # Кэш для API-совместимости с Kalman8D
        self._merged_x = np.zeros(8, dtype=np.float64)  # [cx, cy, w, h, vx, vy, vw, vh]
        self._merged_P = np.eye(8, dtype=np.float64) * 50.0

    def init_from_box(self, box):
        """Инициализация из bounding box."""
        cx, cy = box_center(box)
        w, h = box_wh(box)

        for m in self.models:
            m.x[:] = 0
            m.x[0, 0] = float(cx)
            m.x[1, 0] = float(cy)
            m.x[2, 0] = float(w)
            m.x[3, 0] = float(h)
            m.P = np.eye(m.dim_x, dtype=np.float64) * 50.0

        self.mu = np.array(IMM_INIT_PROBS, dtype=np.float64)
        self.mu /= self.mu.sum()
        self.initialized = True
        self._update_merged()

    def predict(self, dt, q_scale=1.0):
        """
        IMM Predict: interaction → predict каждой модели.

        Returns:
            np.array [cx, cy, w, h] — merged prediction
        """
        if not self.initialized:
            return np.zeros(4, dtype=np.float32)

        dt = float(max(1e-3, dt))

        # ─── 1. Interaction (mixing) ─────────────────────────────
        # Вычисляем mixing probabilities
        c_bar = self.TPM.T @ self.mu  # predicted model probs
        c_bar = np.maximum(c_bar, 1e-12)

        mixing_probs = np.zeros((self.n_models, self.n_models), dtype=np.float64)
        for j in range(self.n_models):
            for i in range(self.n_models):
                mixing_probs[i, j] = self.TPM[i, j] * self.mu[i] / c_bar[j]

        # Mixed states for each model
        for j in range(self.n_models):
            mj = self.models[j]
            dim = mj.dim_x

            # Mixed state
            x_mixed = np.zeros((dim, 1), dtype=np.float64)
            for i in range(self.n_models):
                mi = self.models[i]
                # Проецируем состояние mi на размерность mj
                xi_proj = self._project_state(mi.x, mi.dim_x, dim)
                x_mixed += mixing_probs[i, j] * xi_proj

            # Mixed covariance
            P_mixed = np.zeros((dim, dim), dtype=np.float64)
            for i in range(self.n_models):
                mi = self.models[i]
                xi_proj = self._project_state(mi.x, mi.dim_x, dim)
                Pi_proj = self._project_cov(mi.P, mi.dim_x, dim)
                diff = xi_proj - x_mixed
                P_mixed += mixing_probs[i, j] * (Pi_proj + diff @ diff.T)

            mj.x = x_mixed
            mj.P = P_mixed

        # ─── 2. Predict each model ──────────────────────────────
        for m in self.models:
            m.predict(dt)

        self._update_merged()
        return self._merged_x[:4].astype(np.float32)

    def update(self, z):
        """
        IMM Update: update каждой модели → пересчёт вероятностей.

        Args:
            z: [cx, cy, w, h]
        """
        if not self.initialized:
            return

        z = np.array(z, dtype=np.float64).reshape(4)

        # ─── 1. Likelihood каждой модели ─────────────────────────
        likelihoods = np.array(
            [m.likelihood(z) for m in self.models],
            dtype=np.float64
        )

        # ─── 2. Update каждой модели ─────────────────────────────
        for m in self.models:
            m.update(z)

        # ─── 3. Обновление вероятностей моделей ──────────────────
        c_bar = self.TPM.T @ self.mu
        c_bar = np.maximum(c_bar, 1e-12)

        self.mu = c_bar * likelihoods
        total = self.mu.sum()
        if total > 1e-300:
            self.mu /= total
        else:
            self.mu = np.array(IMM_INIT_PROBS, dtype=np.float64)
            self.mu /= self.mu.sum()

        self._update_merged()

    def to_box(self):
        """Возвращает merged bounding box [x1, y1, x2, y2]."""
        cx = float(self._merged_x[0])
        cy = float(self._merged_x[1])
        w = float(max(2.0, self._merged_x[2]))
        h = float(max(2.0, self._merged_x[3]))
        return np.array([cx - w * 0.5, cy - h * 0.5,
                         cx + w * 0.5, cy + h * 0.5], dtype=np.float32)

    def uncertainty(self):
        """Суммарная неопределённость (для совместимости с Kalman8D)."""
        return float(
            self._merged_P[0, 0] + self._merged_P[1, 1]
            + self._merged_P[2, 2] + self._merged_P[3, 3]
        )

    def get_velocity(self):
        """Возвращает (vx, vy) — merged скорость в пикс/сек."""
        return float(self._merged_x[4]), float(self._merged_x[5])

    def get_acceleration(self):
        """Возвращает (ax, ay) — оценка ускорения из CA-модели."""
        ca = self.models[1]  # CA модель
        if ca.dim_x >= 8:
            return float(ca.x[6, 0]), float(ca.x[7, 0])
        return 0.0, 0.0

    def get_turn_rate(self):
        """Возвращает omega — угловая скорость из CT-модели."""
        ct = self.models[2]  # CT модель
        if ct.dim_x >= 7:
            return float(ct.x[6, 0])
        return 0.0

    def get_model_probs(self):
        """Возвращает вероятности моделей [p_CV, p_CA, p_CT]."""
        return self.mu.copy()

    def get_dominant_model(self):
        """Возвращает название наиболее вероятной модели."""
        idx = int(np.argmax(self.mu))
        return self.model_names[idx]

    @property
    def x(self):
        """Совместимость с Kalman8D.x — возвращает 8x1 вектор."""
        return self._merged_x.reshape(8, 1).astype(np.float32)

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

    def _project_state(self, x, from_dim, to_dim):
        """Проецирует состояние между моделями разной размерности."""
        out = np.zeros((to_dim, 1), dtype=np.float64)
        # Первые 4 компонента (cx, cy, w, h) всегда совпадают
        n_copy = min(4, from_dim, to_dim)
        out[:n_copy] = x[:n_copy]

        # Скорости vx, vy (индексы 4,5 в CV/CT, 4,5 в CA)
        if from_dim >= 6 and to_dim >= 6:
            out[4] = x[4]
            out[5] = x[5]

        # Скорости размера: зависит от модели
        # CV: vw=x[6], vh=x[7]
        # CA: vw=x[8], vh=x[9]
        # CT: vw=x[7], vh=x[8]
        # Для простоты: копируем что можем
        return out

    def _project_cov(self, P, from_dim, to_dim):
        """Проецирует ковариацию между моделями."""
        out = np.eye(to_dim, dtype=np.float64) * 50.0
        n = min(from_dim, to_dim)
        out[:n, :n] = P[:n, :n]
        return out

    def _update_merged(self):
        """Обновляет merged state как взвешенную сумму моделей."""
        self._merged_x[:] = 0
        self._merged_P[:] = 0

        for i, m in enumerate(self.models):
            xi = self._project_state(m.x, m.dim_x, 8).flatten()
            self._merged_x += self.mu[i] * xi

        for i, m in enumerate(self.models):
            xi = self._project_state(m.x, m.dim_x, 8).flatten()
            Pi = self._project_cov(m.P, m.dim_x, 8)
            diff = (xi - self._merged_x).reshape(8, 1)
            self._merged_P += self.mu[i] * (Pi + diff @ diff.T)

    # ─── Status / Draw ───────────────────────────────────────────

    def status_line(self):
        if not self.enabled:
            return "IMM filter disabled"
        return f"IMM filter ready: models={self.model_names}"

    def draw_overlay(self, frame_bgr):
        """Рисует вероятности моделей."""
        if not self.enabled or not self.initialized or not DRAW_IMM_PROBS:
            return

        import cv2

        h, w = frame_bgr.shape[:2]
        probs = self.get_model_probs()
        dominant = self.get_dominant_model()

        # Барграф вероятностей
        bar_x = w - 150
        bar_y = 100
        bar_w = 120
        bar_h = 16

        colors = [
            (200, 200, 200),  # CV — серый
            (0, 165, 255),    # CA — оранжевый
            (0, 0, 255),      # CT — красный
        ]

        for i, (name, prob) in enumerate(zip(self.model_names, probs)):
            y = bar_y + i * (bar_h + 6)
            # Фон
            cv2.rectangle(frame_bgr, (bar_x, y), (bar_x + bar_w, y + bar_h),
                          (60, 60, 60), -1)
            # Заполнение
            fill_w = int(bar_w * prob)
            cv2.rectangle(frame_bgr, (bar_x, y), (bar_x + fill_w, y + bar_h),
                          colors[i], -1)
            # Текст
            cv2.putText(frame_bgr, f"{name} {prob:.0%}",
                        (bar_x + 4, y + bar_h - 3),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.40,
                        (255, 255, 255), 1)

        cv2.putText(frame_bgr, f"IMM: {dominant}",
                    (bar_x, bar_y - 8),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.50,
                    (0, 255, 255), 1)