# ============================================================
# intercept_fsm.py
# Модуль 3: Finite State Machine фаз перехвата
#
# Состояния:
#   SEARCH    — нет цели, полное сканирование
#   ACQUIRE   — цель обнаружена, набираем подтверждений
#   TRACK     — стабильный трек, удержание / сближение
#   INTERCEPT — активная атака, PN-наведение
#   TERMINAL  — финальная фаза (<1.5 сек до контакта)
#   LOST      — потеря цели в ходе перехвата, попытка реакквизиции
#
# FSM управляет:
#   - какой закон наведения использовать
#   - throttle bias
#   - разрешение на переключение трека
#   - параметры детектора (ROI / fullscan)
# ============================================================

import time
import numpy as np

from config import *
from config_intercept import *
from helpers import clamp

import cv2


class InterceptPhase:
    """Enum-like для фаз."""
    SEARCH = "SEARCH"
    ACQUIRE = "ACQUIRE"
    TRACK = "TRACK"
    INTERCEPT = "INTERCEPT"
    TERMINAL = "TERMINAL"
    LOST = "LOST"


class InterceptFSM:
    """
    Конечный автомат фаз перехвата.

    Определяет текущую фазу на основе:
    - состояния трекера (confirmed, miss_streak, hit_streak)
    - оценки дальности (range_m)
    - tau (time-to-contact)
    - уверенности (confidence)

    Каждая фаза определяет набор параметров для guidance, throttle и трекера.
    """

    def __init__(self):
        self.enabled = bool(INTERCEPT_FSM_ENABLE)
        self.phase = InterceptPhase.SEARCH
        self.prev_phase = InterceptPhase.SEARCH
        self.phase_start_time = time.perf_counter()
        self.phase_frame_count = 0

        # Счётчики для переходов
        self._acquire_frames = 0
        self._lost_frames = 0
        self._terminal_locked = False

        # Замороженные команды для TERMINAL lock
        self._frozen_cmd = None

        # Лог переходов
        self.transitions = []

    def reset(self):
        self._transition_to(InterceptPhase.SEARCH)
        self._acquire_frames = 0
        self._lost_frames = 0
        self._terminal_locked = False
        self._frozen_cmd = None
        self.transitions.clear()

    def update(self, confirmed, hit_streak, miss_streak,
               range_state, guidance_state, frame_id):
        """
        Обновить FSM. Вызывается каждый кадр.

        Args:
            confirmed:      bool — цель подтверждена трекером
            hit_streak:     int — кол-во последовательных hit'ов
            miss_streak:    int — кол-во последовательных miss'ов
            range_state:    dict из RangeEstimator.update()
            guidance_state: dict из guidance.py или PN
            frame_id:       int

        Returns:
            dict — параметры текущей фазы для всех модулей
        """
        if not self.enabled:
            return self._default_params()

        range_m = range_state.get("range_m") if range_state else None
        tau = range_state.get("tau", float('inf')) if range_state else float('inf')
        closing_vel = range_state.get("closing_vel", 0.0) if range_state else 0.0
        range_conf = range_state.get("confidence", 0.0) if range_state else 0.0
        range_phase = range_state.get("phase", "far") if range_state else "far"

        self.phase_frame_count += 1

        # ─── Переходы ───────────────────────────────────────────

        if self.phase == InterceptPhase.SEARCH:
            if confirmed:
                self._transition_to(InterceptPhase.ACQUIRE)

        elif self.phase == InterceptPhase.ACQUIRE:
            if not confirmed and miss_streak > 3:
                self._transition_to(InterceptPhase.SEARCH)
            elif hit_streak >= int(FSM_TRACK_CONFIRM_HITS):
                self._transition_to(InterceptPhase.TRACK)
            elif self.phase_frame_count > int(FSM_ACQUIRE_TIMEOUT):
                self._transition_to(InterceptPhase.SEARCH)

        elif self.phase == InterceptPhase.TRACK:
            if not confirmed and miss_streak > int(FSM_LOST_TIMEOUT):
                self._transition_to(InterceptPhase.LOST)
            elif self._should_intercept(range_m, tau, closing_vel, range_conf):
                self._transition_to(InterceptPhase.INTERCEPT)

        elif self.phase == InterceptPhase.INTERCEPT:
            if not confirmed and miss_streak > int(FSM_LOST_TIMEOUT) // 2:
                self._transition_to(InterceptPhase.LOST)
            elif self._should_terminal(range_m, tau, range_conf):
                self._transition_to(InterceptPhase.TERMINAL)
            elif self._should_back_to_track(range_m, closing_vel):
                self._transition_to(InterceptPhase.TRACK)

        elif self.phase == InterceptPhase.TERMINAL:
            # В TERMINAL мы уже не откатываемся — либо контакт,
            # либо пролёт → LOST
            if not confirmed and miss_streak > 10:
                # Вероятно пролетели мимо или столкнулись
                self._transition_to(InterceptPhase.SEARCH)
            elif tau > float(TAU_WARN_SEC) * 2.0 and miss_streak > 3:
                # Tau вырос обратно — пролёт
                self._transition_to(InterceptPhase.LOST)

        elif self.phase == InterceptPhase.LOST:
            if confirmed and hit_streak >= 2:
                # Реакквизиция — возвращаемся в TRACK или INTERCEPT
                if self._should_intercept(range_m, tau, closing_vel, range_conf):
                    self._transition_to(InterceptPhase.INTERCEPT)
                else:
                    self._transition_to(InterceptPhase.TRACK)
            elif self.phase_frame_count > int(FSM_LOST_RECOVER_TIMEOUT):
                self._transition_to(InterceptPhase.SEARCH)

        # ─── Параметры текущей фазы ─────────────────────────────
        return self._build_params(range_m, tau, closing_vel,
                                  guidance_state, confirmed)

    # ─── Условия переходов ───────────────────────────────────────

    def _should_intercept(self, range_m, tau, closing_vel, conf):
        """TRACK → INTERCEPT"""
        if range_m is not None and range_m <= float(FSM_INTERCEPT_RANGE_M):
            return True
        if closing_vel > float(PN_MIN_CLOSING_VEL) and tau < float(TAU_WARN_SEC) * 2.0:
            return True
        return False

    def _should_terminal(self, range_m, tau, conf):
        """INTERCEPT → TERMINAL"""
        if tau <= float(FSM_TERMINAL_TAU_SEC) and tau > 0:
            return True
        if range_m is not None and range_m <= float(FSM_TERMINAL_RANGE_M):
            return True
        return False

    def _should_back_to_track(self, range_m, closing_vel):
        """INTERCEPT → TRACK (цель удаляется)"""
        if closing_vel < -2.0 and self.phase_frame_count > 30:
            return True
        if range_m is not None and range_m > float(FSM_INTERCEPT_RANGE_M) * 1.5:
            return True
        return False

    # ─── Параметры фаз ──────────────────────────────────────────

    def _build_params(self, range_m, tau, closing_vel,
                      guidance_state, confirmed):
        """Собирает параметры для всех модулей на основе текущей фазы."""

        params = {
            "phase": self.phase,
            "prev_phase": self.prev_phase,
            "phase_frame_count": self.phase_frame_count,
            "throttle_bias": 0.5,
            "use_pn": False,
            "use_screen_guidance": True,
            "allow_target_switch": True,
            "force_fullscan": False,
            "det_every_override": None,
            "yolo_enabled": True,
            "klt_only": False,
            "lock_commands": False,
            "frozen_cmd": None,
            "range_m": range_m,
            "tau": tau,
            "closing_vel": closing_vel,
        }

        if self.phase == InterceptPhase.SEARCH:
            params.update({
                "throttle_bias": 0.4,
                "use_pn": False,
                "use_screen_guidance": False,
                "allow_target_switch": True,
                "force_fullscan": True,
            })

        elif self.phase == InterceptPhase.ACQUIRE:
            params.update({
                "throttle_bias": 0.45,
                "use_pn": False,
                "use_screen_guidance": True,
                "allow_target_switch": True,
                "det_every_override": 2,
            })

        elif self.phase == InterceptPhase.TRACK:
            params.update({
                "throttle_bias": float(FSM_TRACK_THROTTLE_BIAS),
                "use_pn": False,
                "use_screen_guidance": True,
                "allow_target_switch": False,
            })

        elif self.phase == InterceptPhase.INTERCEPT:
            params.update({
                "throttle_bias": float(FSM_INTERCEPT_THROTTLE_BIAS),
                "use_pn": True,
                "use_screen_guidance": False,
                "allow_target_switch": False,
                "det_every_override": 1,  # максимальная частота детекции
            })

        elif self.phase == InterceptPhase.TERMINAL:
            should_lock = (
                bool(FSM_TERMINAL_LOCK_COMMANDS)
                and tau <= float(FSM_TERMINAL_LOCK_TAU_SEC)
            )

            if should_lock and not self._terminal_locked:
                # Замораживаем команды
                self._terminal_locked = True
                if guidance_state is not None:
                    self._frozen_cmd = {
                        "steer_x": guidance_state.get("steer_x", 0.0),
                        "steer_y": guidance_state.get("steer_y", 0.0),
                        "cmd_x": guidance_state.get("cmd_x", 0.0),
                        "cmd_y": guidance_state.get("cmd_y", 0.0),
                    }

            params.update({
                "throttle_bias": float(FSM_TERMINAL_THROTTLE),
                "use_pn": not should_lock,
                "use_screen_guidance": False,
                "allow_target_switch": False,
                "yolo_enabled": not bool(FSM_TERMINAL_USE_KLT_ONLY),
                "klt_only": bool(FSM_TERMINAL_USE_KLT_ONLY),
                "lock_commands": should_lock,
                "frozen_cmd": self._frozen_cmd if should_lock else None,
                "det_every_override": 1,
            })

        elif self.phase == InterceptPhase.LOST:
            params.update({
                "throttle_bias": 0.5,
                "use_pn": False,
                "use_screen_guidance": False,
                "allow_target_switch": True,
                "force_fullscan": True,
            })

        return params

    def _default_params(self):
        """Параметры по умолчанию когда FSM отключена."""
        return {
            "phase": "DISABLED",
            "prev_phase": "DISABLED",
            "phase_frame_count": 0,
            "throttle_bias": 0.5,
            "use_pn": False,
            "use_screen_guidance": True,
            "allow_target_switch": True,
            "force_fullscan": False,
            "det_every_override": None,
            "yolo_enabled": True,
            "klt_only": False,
            "lock_commands": False,
            "frozen_cmd": None,
            "range_m": None,
            "tau": float('inf'),
            "closing_vel": 0.0,
        }

    def _transition_to(self, new_phase):
        """Выполнить переход в новую фазу."""
        if new_phase == self.phase:
            return

        old_phase = self.phase
        self.prev_phase = old_phase
        self.phase = new_phase
        self.phase_start_time = time.perf_counter()
        self.phase_frame_count = 0

        # Сброс фазо-специфичных счётчиков
        if new_phase == InterceptPhase.TERMINAL:
            self._terminal_locked = False
            self._frozen_cmd = None
        elif new_phase == InterceptPhase.SEARCH:
            self._acquire_frames = 0
            self._lost_frames = 0
            self._terminal_locked = False
            self._frozen_cmd = None

        self.transitions.append({
            "time": time.perf_counter(),
            "from": old_phase,
            "to": new_phase,
        })

        # Ограничиваем лог
        if len(self.transitions) > 200:
            self.transitions = self.transitions[-100:]

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

    def status_line(self):
        if not self.enabled:
            return "Intercept FSM disabled"
        return f"Intercept FSM ready: phase={self.phase}"

    def draw_overlay(self, frame_bgr, params):
        """Рисует текущую фазу перехвата."""
        if not self.enabled or not DRAW_INTERCEPT_FSM:
            return

        phase = params.get("phase", "?")
        tau = params.get("tau", float('inf'))
        range_m = params.get("range_m")
        throttle = params.get("throttle_bias", 0.5)

        # Цвета по фазе
        phase_colors = {
            InterceptPhase.SEARCH: (150, 150, 150),
            InterceptPhase.ACQUIRE: (0, 200, 255),
            InterceptPhase.TRACK: (255, 200, 0),
            InterceptPhase.INTERCEPT: (0, 128, 255),
            InterceptPhase.TERMINAL: (0, 0, 255),
            InterceptPhase.LOST: (128, 0, 128),
        }
        color = phase_colors.get(phase, (180, 180, 180))

        # Фон для фазы (полоска сверху)
        h, w = frame_bgr.shape[:2]
        if phase == InterceptPhase.TERMINAL:
            # Мигающая красная полоска
            if (self.phase_frame_count // 4) % 2 == 0:
                cv2.rectangle(frame_bgr, (0, 0), (w, 6), (0, 0, 255), -1)

        # Текст фазы
        range_str = f"{range_m:.1f}m" if range_m is not None else "---"
        tau_str = f"{tau:.2f}s" if tau < 100 else "---"

        txt = (f"FSM: {phase}  R={range_str}  "
               f"tau={tau_str}  thr={throttle:.0%}")
        cv2.putText(frame_bgr, txt, (20, 290),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.60, color, 2)

        # Индикатор use_pn
        if params.get("use_pn"):
            cv2.putText(frame_bgr, "PN", (w - 60, 40),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 128), 2)

        if params.get("lock_commands"):
            cv2.putText(frame_bgr, "CMD LOCK", (w - 160, 70),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
