MAI/guidance.py

import json
import time
from pathlib import Path

import cv2
import numpy as np

from config import *
from helpers import box_center, clamp, clip_box


def _curve_command(value, deadzone, expo, gain):
    v = float(np.clip(value, -1.0, 1.0))
    av = abs(v)
    if av <= deadzone:
        return 0.0
    t = (av - deadzone) / max(1e-6, 1.0 - deadzone)
    out = min(1.0, float(gain) * (t ** float(expo)))
    return out if v >= 0.0 else -out


class ScreenGuidanceController:
    def __init__(self):
        self.enabled = bool(GUIDANCE_ENABLE)
        self.export_enable = bool(self.enabled and GUIDANCE_EXPORT_ENABLE)
        self.export_every = int(max(1, GUIDANCE_EXPORT_EVERY))
        self.export_path = Path(GUIDANCE_EXPORT_PATH)
        if self.export_enable:
            self.export_path.parent.mkdir(parents=True, exist_ok=True)

        self.cmd_keep = float(clamp(GUIDANCE_CMD_SMOOTH, 0.0, 0.97))
        self.point_keep = float(clamp(GUIDANCE_POINT_SMOOTH, 0.0, 0.97))
        self.last_export_frame = -10**9
        self.export_retry_at = 0.0
        self.last_export_warn_ts = 0.0
        self.smoothed_cmd = np.zeros(2, dtype=np.float32)
        self.smoothed_point = None

    def status_line(self):
        if not self.enabled:
            return "Guidance disabled"
        if self.export_enable:
            return f"Guidance ready: export={self.export_path}"
        return "Guidance ready"

    def reset_for_target_switch(self, aim_point, cmd_damp=0.0):
        aim = np.array(aim_point, dtype=np.float32)
        self.smoothed_point = aim.copy()
        damp = float(clamp(cmd_damp, 0.0, 1.0))
        self.smoothed_cmd = (self.smoothed_cmd * damp).astype(np.float32)

    def idle_state(self):
        return {
            "active": False,
            "status": "SEARCH",
            "frame_id": -1,
            "target_id": None,
            "frame_w": None,
            "frame_h": None,
            "aim_x": None,
            "aim_y": None,
            "box_w": None,
            "box_h": None,
            "error_x": 0.0,
            "error_y": 0.0,
            "cmd_x": 0.0,
            "cmd_y": 0.0,
            "steer_x": 0.0,
            "steer_y": 0.0,
            "look_dx": 0.0,
            "look_dy": 0.0,
            "confidence": 0.0,
            "on_target": False,
        }

    def update(
        self,
        frame_id,
        frame_w,
        frame_h,
        confirmed,
        locked_box,
        pred_box,
        override_box,
        target_id,
        target_track_score,
        klt_valid,
        klt_quality,
        miss_streak,
        vx,
        vy,
    ):
        state = self.idle_state()
        state["frame_id"] = int(frame_id)
        state["target_id"] = None if target_id is None else int(target_id)
        state["frame_w"] = int(frame_w)
        state["frame_h"] = int(frame_h)

        if not self.enabled:
            return state

        guide_box = None
        if confirmed and override_box is not None:
            guide_box = np.array(override_box, dtype=np.float32)
        elif confirmed and locked_box is not None:
            guide_box = np.array(locked_box, dtype=np.float32)
        elif confirmed and pred_box is not None:
            guide_box = np.array(pred_box, dtype=np.float32)

        if guide_box is None:
            self.smoothed_cmd *= 0.60
            state["cmd_x"] = float(self.smoothed_cmd[0])
            state["cmd_y"] = float(self.smoothed_cmd[1])
            state["steer_x"] = float(self.smoothed_cmd[0])
            state["steer_y"] = float(self.smoothed_cmd[1])
            state["look_dx"] = float(self.smoothed_cmd[0] * float(GUIDANCE_LOOK_MAX_STEP_PX))
            state["look_dy"] = float(self.smoothed_cmd[1] * float(GUIDANCE_LOOK_MAX_STEP_PX))
            self._maybe_export(state)
            return state

        guide_box = clip_box(guide_box, frame_w, frame_h)
        center = box_center(guide_box).astype(np.float32)
        box_w = float(max(0.0, guide_box[2] - guide_box[0]))
        box_h = float(max(0.0, guide_box[3] - guide_box[1]))

        lead = np.array([float(vx), float(vy)], dtype=np.float32) * float(GUIDANCE_LEAD_SEC)
        lead_norm = float(np.linalg.norm(lead))
        if lead_norm > float(GUIDANCE_MAX_LEAD_PX):
            lead *= float(GUIDANCE_MAX_LEAD_PX) / max(1e-6, lead_norm)

        aim = center + lead
        aim[1] += float(GUIDANCE_BOX_BIAS_Y) * max(2.0, float(guide_box[3] - guide_box[1]))
        aim[0] = float(clamp(aim[0], 0.0, float(frame_w - 1)))
        aim[1] = float(clamp(aim[1], 0.0, float(frame_h - 1)))

        if self.smoothed_point is None:
            self.smoothed_point = aim.copy()
        else:
            self.smoothed_point = (
                (self.point_keep * self.smoothed_point)
                + ((1.0 - self.point_keep) * aim)
            ).astype(np.float32)

        screen_center = np.array([0.5 * float(frame_w), 0.5 * float(frame_h)], dtype=np.float32)
        error_px = self.smoothed_point - screen_center
        error_norm = np.array(
            [
                float(error_px[0] / max(1.0, 0.5 * float(frame_w))),
                float(error_px[1] / max(1.0, 0.5 * float(frame_h))),
            ],
            dtype=np.float32,
        )
        error_norm = np.clip(error_norm, -1.0, 1.0)

        raw_cmd = np.array(
            [
                _curve_command(error_norm[0], GUIDANCE_DEADZONE_NORM, GUIDANCE_EXPO, GUIDANCE_GAIN_X),
                _curve_command(error_norm[1], GUIDANCE_DEADZONE_NORM, GUIDANCE_EXPO, GUIDANCE_GAIN_Y),
            ],
            dtype=np.float32,
        )
        self.smoothed_cmd = (
            (self.cmd_keep * self.smoothed_cmd)
            + ((1.0 - self.cmd_keep) * raw_cmd)
        ).astype(np.float32)

        confidence = 0.0
        confidence += 0.40 if confirmed else 0.0
        confidence += 0.15 if locked_box is not None else 0.05

        # KLT теперь главный источник доверия (klt_quality ≈ 0.93)
        if klt_valid:
            confidence += 0.30 * float(clamp(klt_quality, 0.0, 1.0))
        else:
            confidence += 0.05 * float(clamp(klt_quality, 0.0, 1.0))

        # ByteTrack score — второстепенный
        if target_track_score is not None:
            confidence += 0.10 * float(clamp(target_track_score, 0.0, 1.0))

        # Мягче штраф за miss (YOLO слабый, это временно)
        confidence -= 0.025 * float(clamp(miss_streak, 0, 10))

        confidence = float(clamp(confidence, 0.0, 1.0))

        err_mag = float(np.linalg.norm(error_norm))
        on_target = bool(
            confirmed
            and (confidence >= float(GUIDANCE_CONF_MIN))
            and (err_mag <= float(GUIDANCE_ON_TARGET_NORM))
        )

        status = "SEARCH"
        if confirmed and override_box is not None and miss_streak > 0:
            status = "REACQ"
        elif confirmed and miss_streak > 0:
            status = "HOLD"
        elif confirmed:
            status = "LOCK"

        state.update(
            {
                "active": bool(confirmed),
                "status": status,
                "aim_x": float(self.smoothed_point[0]),
                "aim_y": float(self.smoothed_point[1]),
                "box_w": box_w,
                "box_h": box_h,
                "error_x": float(error_norm[0]),
                "error_y": float(error_norm[1]),
                "cmd_x": float(self.smoothed_cmd[0]),
                "cmd_y": float(self.smoothed_cmd[1]),
                "steer_x": float(self.smoothed_cmd[0]),
                "steer_y": float(self.smoothed_cmd[1]),
                "look_dx": float(self.smoothed_cmd[0] * float(GUIDANCE_LOOK_MAX_STEP_PX)),
                "look_dy": float(self.smoothed_cmd[1] * float(GUIDANCE_LOOK_MAX_STEP_PX)),
                "confidence": confidence,
                "on_target": on_target,
            }
        )
        self._maybe_export(state)
        return state

    def draw_overlay(self, frame_bgr, state, sx, sy):
        if (not self.enabled) or (not DRAW_GUIDANCE):
            return

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

        color = (0, 220, 255) if state.get("active", False) else (120, 120, 120)
        cv2.drawMarker(frame_bgr, (cx, cy), color, markerType=cv2.MARKER_CROSS, markerSize=18, thickness=1)

        aim_x = state.get("aim_x", None)
        aim_y = state.get("aim_y", None)
        if aim_x is not None and aim_y is not None:
            ax = int(float(aim_x) / max(1e-6, float(sx)))
            ay = int(float(aim_y) / max(1e-6, float(sy)))
            cv2.circle(frame_bgr, (ax, ay), 8, color, 2)
            cv2.line(frame_bgr, (cx, cy), (ax, ay), color, 1, cv2.LINE_AA)
            cv2.putText(
                frame_bgr,
                f"G {state['status']} steer=({state['steer_x']:+.2f},{state['steer_y']:+.2f}) conf={state['confidence']:.2f}",
                (20, 165),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.60,
                color,
                2,
            )

    def _maybe_export(self, state):
        frame_id = int(state["frame_id"])
        now = time.monotonic()
        if (not self.export_enable) or ((frame_id - self.last_export_frame) < self.export_every) or (now < self.export_retry_at):
            return

        payload = {
            "frame_id": frame_id,
            "active": bool(state["active"]),
            "status": str(state["status"]),
            "target_id": state["target_id"],
            "frame_w": state["frame_w"],
            "frame_h": state["frame_h"],
            "aim_x": state["aim_x"],
            "aim_y": state["aim_y"],
            "box_w": state["box_w"],
            "box_h": state["box_h"],
            "error_x": float(state["error_x"]),
            "error_y": float(state["error_y"]),
            "cmd_x": float(state["cmd_x"]),
            "cmd_y": float(state["cmd_y"]),
            "steer_x": float(state["steer_x"]),
            "steer_y": float(state["steer_y"]),
            "look_dx": float(state["look_dx"]),
            "look_dy": float(state["look_dy"]),
            "confidence": float(state["confidence"]),
            "on_target": bool(state["on_target"]),
        }
        tmp_path = self.export_path.with_suffix(self.export_path.suffix + ".tmp")
        try:
            with tmp_path.open("w", encoding="utf-8") as f:
                json.dump(payload, f, ensure_ascii=True, indent=2)
            self._replace_export_file(tmp_path)
        except PermissionError as exc:
            self.export_retry_at = time.monotonic() + 0.25
            self._warn_export_error(frame_id, exc)
            return
        except OSError as exc:
            self.export_retry_at = time.monotonic() + 0.50
            self._warn_export_error(frame_id, exc)
            return

        self.export_retry_at = 0.0
        self.last_export_frame = frame_id

    def _replace_export_file(self, tmp_path):
        last_err = None
        for delay_sec in (0.0, 0.01, 0.03):
            if delay_sec > 0.0:
                time.sleep(delay_sec)
            try:
                tmp_path.replace(self.export_path)
                return
            except PermissionError as exc:
                last_err = exc

        if last_err is not None:
            raise last_err

    def _warn_export_error(self, frame_id, exc):
        now = time.monotonic()
        if (now - self.last_export_warn_ts) < 2.0:
            return
        print(f"[guidance] export skipped at frame {frame_id} for {self.export_path}: {exc}")
        self.last_export_warn_ts = now