# ============================================================
# proportional_navigation.py
# Модуль 2: Пропорциональная навигация (Proportional Navigation)
#
# Реализует:
#   - Pure Proportional Navigation (PPN)
#   - Augmented Proportional Navigation (APN) с компенсацией
#     ускорения цели
#   - Screen-space fallback (когда нет оценки дальности)
#   - Вычисление точки перехвата (intercept point)
# ============================================================

import numpy as np
from collections import deque

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

import cv2


class ProportionalNavigator:
    """
    Proportional Navigation — классический закон наведения ракет.

    Принцип: ускорение перехватчика перпендикулярно линии визирования (LOS)
    пропорционально скорости вращения LOS и скорости сближения:

        a_cmd = N * Vc * dλ/dt

    где N — навигационная константа (3-5),
        Vc — closing velocity,
        dλ/dt — скорость изменения угла LOS.

    Augmented PN добавляет компенсацию ускорения цели:

        a_cmd = N * Vc * dλ/dt + (N/2) * a_target_normal

    При отсутствии 3D-информации работает в screen-space
    (2D аналог: стабилизировать положение цели в центре кадра).
    """

    def __init__(self):
        self.enabled = bool(PN_ENABLE)
        self.N = float(PN_NAV_GAIN)
        self.augmented = bool(PN_AUGMENTED)

        # Состояние LOS
        self._prev_los = None          # предыдущий угол LOS (рад)
        self._prev_los_t = None
        self._los_rate_smooth = 0.0    # сглаженная скорость LOS (рад/с)
        self._los_rate_raw = 0.0

        # Состояние LOS по обеим осям (для 2D PN в screen space)
        self._prev_los_x = None
        self._prev_los_y = None

        # Для Augmented PN: оценка ускорения цели
        self._prev_target_vel = None
        self._prev_target_vel_t = None
        self._target_accel = np.zeros(2, dtype=np.float32)

        # История для отладки
        self._cmd_hist = deque(maxlen=60)

        # Выходы
        self.accel_cmd = np.zeros(2, dtype=np.float32)    # [ax, ay] нормализованные
        self.intercept_point = None                        # предсказанная точка перехвата
        self.intercept_time = None                         # время до перехвата (сек)
        self.los_rate = 0.0
        self.active = False

    def reset(self):
        self._prev_los = None
        self._prev_los_t = None
        self._los_rate_smooth = 0.0
        self._los_rate_raw = 0.0
        self._prev_los_x = None
        self._prev_los_y = None
        self._prev_target_vel = None
        self._prev_target_vel_t = None
        self._target_accel = np.zeros(2, dtype=np.float32)
        self.accel_cmd = np.zeros(2, dtype=np.float32)
        self.intercept_point = None
        self.intercept_time = None
        self.los_rate = 0.0
        self.active = False
        self._cmd_hist.clear()

    def update_3d(self, my_pos, my_vel, target_pos, target_vel,
                  closing_vel, timestamp):
        """
        3D Proportional Navigation (используется при наличии оценки дальности).

        Args:
            my_pos:      np.array([x, y]) — позиция перехватчика (метры)
            my_vel:      np.array([vx, vy]) — скорость перехватчика (м/с)
            target_pos:  np.array([x, y]) — позиция цели (метры)
            target_vel:  np.array([vx, vy]) — скорость цели (м/с)
            closing_vel: float — скорость сближения (м/с), > 0 = сближаемся
            timestamp:   float — текущее время

        Returns:
            dict с полями accel_cmd, los_rate, intercept_point, intercept_time
        """
        result = self._idle_result()
        if not self.enabled:
            return result

        r = target_pos - my_pos
        range_val = float(np.linalg.norm(r))
        if range_val < 0.01:
            return result

        # Угол LOS
        los_angle = float(np.arctan2(r[1], r[0]))

        # Скорость LOS
        if self._prev_los is not None and self._prev_los_t is not None:
            dt = timestamp - self._prev_los_t
            if dt > 1e-4:
                # Обработка перехода через ±π
                d_los = los_angle - self._prev_los
                if d_los > np.pi:
                    d_los -= 2.0 * np.pi
                elif d_los < -np.pi:
                    d_los += 2.0 * np.pi

                self._los_rate_raw = d_los / dt

                alpha = float(PN_LOS_RATE_SMOOTH)
                self._los_rate_smooth = (
                    alpha * self._los_rate_smooth
                    + (1.0 - alpha) * self._los_rate_raw
                )

        self._prev_los = los_angle
        self._prev_los_t = timestamp
        self.los_rate = self._los_rate_smooth

        # Closing velocity
        vc = max(float(closing_vel), float(PN_MIN_CLOSING_VEL))

        # ─── Pure PN ─────────────────────────────────────────────
        a_pn = self.N * vc * self._los_rate_smooth

        # ─── Augmented PN ────────────────────────────────────────
        a_aug = 0.0
        if self.augmented:
            target_vel_np = np.array(target_vel, dtype=np.float32)
            self._estimate_target_accel(target_vel_np, timestamp)

            # Нормальная компонента ускорения цели (перп. к LOS)
            if range_val > 0.1:
                los_unit = r / range_val
                los_perp = np.array([-los_unit[1], los_unit[0]], dtype=np.float32)
                a_target_normal = float(np.dot(self._target_accel, los_perp))
                a_aug = float(PN_AUG_GAIN) * self.N * 0.5 * a_target_normal

        # Суммарная команда
        a_total = a_pn + a_aug

        # Нормализация
        max_cmd = float(PN_MAX_ACCEL_CMD)
        a_norm = float(clamp(a_total / max(20.0, vc), -max_cmd, max_cmd))

        # Преобразование в декартовы координаты (перп. к LOS)
        los_perp_dir = np.array([-np.sin(los_angle), np.cos(los_angle)],
                                dtype=np.float32)
        self.accel_cmd = a_norm * los_perp_dir

        # Точка перехвата
        self.intercept_point, self.intercept_time = self._compute_intercept(
            my_pos, my_vel, target_pos, target_vel
        )

        self.active = True
        self._cmd_hist.append(float(a_norm))

        result.update({
            "active": True,
            "accel_x": float(self.accel_cmd[0]),
            "accel_y": float(self.accel_cmd[1]),
            "accel_mag": float(abs(a_norm)),
            "los_rate": float(self._los_rate_smooth),
            "los_angle_deg": float(np.degrees(los_angle)),
            "closing_vel": float(closing_vel),
            "range_m": float(range_val),
            "intercept_x": float(self.intercept_point[0]) if self.intercept_point is not None else None,
            "intercept_y": float(self.intercept_point[1]) if self.intercept_point is not None else None,
            "intercept_time": self.intercept_time,
            "a_pn": float(a_pn),
            "a_aug": float(a_aug),
        })
        return result

    def update_screen(self, target_center, frame_w, frame_h,
                      target_vel_px, timestamp):
        """
        Screen-space PN fallback (без 3D-информации).

        Вместо 3D LOS используем углы от центра кадра к цели.
        LOS rate вычисляется из скорости смещения цели по экрану.

        Args:
            target_center: (cx, cy) — центр bbox цели в пикселях
            frame_w, frame_h: размеры кадра
            target_vel_px: (vx, vy) — скорость цели в пикселях/сек (из Kalman)
            timestamp: текущее время

        Returns:
            dict — команды наведения в screen space
        """
        result = self._idle_result()
        if not self.enabled or not PN_SCREEN_FALLBACK:
            return result
        if target_center is None:
            return result

        cx, cy = float(target_center[0]), float(target_center[1])
        scr_cx = float(frame_w) * 0.5
        scr_cy = float(frame_h) * 0.5

        # Нормализованная ошибка (как LOS angle в screen space)
        los_x = (cx - scr_cx) / max(1.0, scr_cx)
        los_y = (cy - scr_cy) / max(1.0, scr_cy)

        # LOS rate по каждой оси
        los_rate_x = 0.0
        los_rate_y = 0.0
        if self._prev_los_x is not None and self._prev_los_t is not None:
            dt = timestamp - self._prev_los_t
            if dt > 1e-4:
                raw_x = (los_x - self._prev_los_x) / dt
                raw_y = (los_y - self._prev_los_y) / dt
                alpha = float(PN_LOS_RATE_SMOOTH)
                los_rate_x = alpha * los_rate_x + (1.0 - alpha) * raw_x
                los_rate_y = alpha * los_rate_y + (1.0 - alpha) * raw_y

        self._prev_los_x = los_x
        self._prev_los_y = los_y
        self._prev_los_t = timestamp

        self.los_rate = float(np.hypot(los_rate_x, los_rate_y))

        # Screen-space "closing velocity" = 1 (нет 3D)
        # Команда = N * los_rate (упрощённый PN)
        cmd_x = float(clamp(self.N * los_rate_x + 0.8 * los_x,
                             -PN_MAX_ACCEL_CMD, PN_MAX_ACCEL_CMD))
        cmd_y = float(clamp(self.N * los_rate_y + 0.8 * los_y,
                             -PN_MAX_ACCEL_CMD, PN_MAX_ACCEL_CMD))

        # Lead: добавляем упреждение по скорости цели
        if target_vel_px is not None:
            vx_n = float(target_vel_px[0]) / max(1.0, scr_cx)
            vy_n = float(target_vel_px[1]) / max(1.0, scr_cy)
            lead_t = float(PN_LEAD_TIME_SEC)
            cmd_x += 0.3 * vx_n * lead_t
            cmd_y += 0.3 * vy_n * lead_t

        cmd_x = float(clamp(cmd_x, -PN_MAX_ACCEL_CMD, PN_MAX_ACCEL_CMD))
        cmd_y = float(clamp(cmd_y, -PN_MAX_ACCEL_CMD, PN_MAX_ACCEL_CMD))

        self.accel_cmd = np.array([cmd_x, cmd_y], dtype=np.float32)
        self.active = True

        result.update({
            "active": True,
            "accel_x": cmd_x,
            "accel_y": cmd_y,
            "accel_mag": float(np.hypot(cmd_x, cmd_y)),
            "los_rate": self.los_rate,
            "los_angle_deg": float(np.degrees(np.arctan2(los_y, los_x))),
            "closing_vel": 0.0,
            "range_m": None,
            "intercept_x": None,
            "intercept_y": None,
            "intercept_time": None,
            "mode": "screen",
        })
        return result

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

    def _estimate_target_accel(self, target_vel, timestamp):
        """Оценка ускорения цели из разности скоростей."""
        if self._prev_target_vel is not None and self._prev_target_vel_t is not None:
            dt = timestamp - self._prev_target_vel_t
            if dt > 1e-4:
                raw_accel = (target_vel - self._prev_target_vel) / dt
                alpha = 0.6
                self._target_accel = (
                    alpha * self._target_accel
                    + (1.0 - alpha) * raw_accel
                ).astype(np.float32)

        self._prev_target_vel = target_vel.copy()
        self._prev_target_vel_t = timestamp

    def _compute_intercept(self, my_pos, my_vel, target_pos, target_vel,
                           max_iter=12):
        """
        Итеративное вычисление точки перехвата.

        Находит время t, при котором перехватчик, летящий со своей скоростью,
        достигает предсказанной позиции цели (target_pos + target_vel * t).
        """
        my_speed = float(np.linalg.norm(my_vel))
        if my_speed < 0.1:
            my_speed = 10.0  # fallback

        t = float(np.linalg.norm(target_pos - my_pos)) / max(my_speed, 0.1)

        for _ in range(max_iter):
            predicted = target_pos + target_vel * t
            dist = float(np.linalg.norm(predicted - my_pos))
            t_new = dist / max(my_speed, 0.1)

            if abs(t_new - t) < 0.01:
                break
            t = 0.5 * t + 0.5 * t_new  # демпфированная итерация

        intercept = target_pos + target_vel * t
        return intercept.astype(np.float32), float(t)

    @staticmethod
    def _idle_result():
        return {
            "active": False,
            "accel_x": 0.0,
            "accel_y": 0.0,
            "accel_mag": 0.0,
            "los_rate": 0.0,
            "los_angle_deg": 0.0,
            "closing_vel": 0.0,
            "range_m": None,
            "intercept_x": None,
            "intercept_y": None,
            "intercept_time": None,
        }

    def status_line(self):
        if not self.enabled:
            return "Proportional Navigation disabled"
        mode = "APN" if self.augmented else "PPN"
        return f"PN ready: N={self.N:.1f} mode={mode}"

    def draw_overlay(self, frame_bgr, state, sx, sy):
        """Рисует вектор наведения и точку перехвата."""
        if not self.enabled or not DRAW_PN_VECTOR:
            return
        if not state.get("active", False):
            return

        h, w = frame_bgr.shape[:2]
        cx = int(w * 0.5)
        cy = int(h * 0.5)

        # Вектор команды
        ax = float(state.get("accel_x", 0.0))
        ay = float(state.get("accel_y", 0.0))
        scale = 80.0
        ex = int(cx + ax * scale)
        ey = int(cy + ay * scale)

        color = (0, 255, 128) if state.get("active") else (100, 100, 100)
        cv2.arrowedLine(frame_bgr, (cx, cy), (ex, ey), color, 2, cv2.LINE_AA)

        # Точка перехвата
        ix = state.get("intercept_x")
        iy = state.get("intercept_y")
        if ix is not None and iy is not None and DRAW_INTERCEPT_POINT:
            ipx = int(float(ix) / max(1e-6, float(sx)))
            ipy = int(float(iy) / max(1e-6, float(sy)))
            cv2.drawMarker(frame_bgr, (ipx, ipy), (0, 255, 128),
                           cv2.MARKER_TILTED_CROSS, 14, 2)

        # Текст
        los_r = state.get("los_rate", 0.0)
        a_mag = state.get("accel_mag", 0.0)
        it = state.get("intercept_time")
        txt = f"PN: a={a_mag:.2f} LOS_r={los_r:.3f}"
        if it is not None:
            txt += f" t_int={it:.2f}s"
        cv2.putText(frame_bgr, txt, (20, 260),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.55, color, 2)
