hookvision-rk/tests/test_state_machine.py

"""Test hook state evaluation logic (evaluate_hook_state).

The state machine transitions (main_window_v2):
  IDLE -> HOOKING (hook-lug intersecting, debounce 3 frames)
  HOOKING -> READY (calibrated offset match via _side_match, debounce ready_debounce)
  HOOKING/READY -> IDLE (no intersection; READY exits after ready_debounce frames)
  Any -> RELEASING (downward movement, debounce 3)

For RELEASING: both hook and lug centers must move down by > RELEASE_DELTA_THRESHOLD pixels.
"""

import pytest
from hookvision.models import HookState, Config, BucketModel


def make_hook(x, y, w=20, h=20):
    return {"x": x, "y": y, "w": w, "h": h, "label": "hook", "color": (0, 0, 255)}


def make_lug(x, y, w=20, h=20):
    return {"x": x, "y": y, "w": w, "h": h, "label": "lug", "color": (255, 0, 0)}


def make_ready(x, y, w=10, h=10):
    return {"x": x, "y": y, "w": w, "h": h, "label": "ready", "color": (0, 255, 255)}


def calibrated_bucket(x_off=0.0, y_off=0.0):
    b = BucketModel(id=1, description="test", code_x="X")
    b.hook_lug_x_offset_ratio_left = x_off
    b.hook_lug_y_offset_ratio_left = y_off
    b.hook_lug_x_offset_ratio_right = x_off
    b.hook_lug_y_offset_ratio_right = y_off
    return b


class MockMainWindow:
    """Minimal mock of MainWindow_v2 for testing evaluate_hook_state."""

    STATE_THRESHOLD = 3
    RELEASE_DELTA_THRESHOLD = 5

    def __init__(self):
        self.left_objects = []
        self.right_objects = []
        self.hook_state = HookState.IDLE
        self.prev_state = HookState.IDLE
        self.state_counter = 0
        self.prev_left_hook = None
        self.prev_right_hook = None
        self.prev_left_lug = None
        self.prev_right_lug = None
        self.active_bucket = None
        self.config = Config()
        self._calib_warned = set()
        self.last_state = None  # tracks update_ui_for_state calls

    def evaluate_hook_state(self):
        # This is the exact implementation from main_window_v2.py
        left_hooks = [o for o in self.left_objects if 'hook' in o['label'].lower()]
        left_lugs = [o for o in self.left_objects if 'lug' in o['label'].lower()]
        right_hooks = [o for o in self.right_objects if 'hook' in o['label'].lower()]
        right_lugs = [o for o in self.right_objects if 'lug' in o['label'].lower()]

        def find_nearest_pair(hooks, lugs):
            if not hooks or not lugs:
                return None
            best_pair = None
            best_dist = float('inf')
            for h in hooks:
                hx = h['x'] + h['w'] / 2
                hy = h['y'] + h['h'] / 2
                for l in lugs:
                    lx = l['x'] + l['w'] / 2
                    ly = l['y'] + l['h'] / 2
                    dist = (hx - lx) ** 2 + (hy - ly) ** 2
                    if dist < best_dist:
                        best_dist = dist
                        best_pair = (h, l)
            return best_pair

        left_pair = find_nearest_pair(left_hooks, left_lugs)
        right_pair = find_nearest_pair(right_hooks, right_lugs)

        self._push_pair_traj("left", left_pair)
        self._push_pair_traj("right", right_pair)
        self._update_suppress_window()

        # 吊耳贴着画面上沿（框顶被裁）：不再判定状态，状态框不显示
        if (left_pair and self._lug_touches_top(left_pair[1])) or \
                (right_pair and self._lug_touches_top(right_pair[1])):
            self.hook_state = HookState.IDLE
            self.prev_state = HookState.IDLE
            self.state_counter = 0
            self._clear_state_display()
            return

        is_releasing = False

        # 处理左侧配对
        if left_pair:
            h, l = left_pair
            if self.prev_left_hook and self.prev_left_lug:
                dy_h = (h['y'] + h['h'] / 2) - self.prev_left_hook[1]
                dy_l = (l['y'] + l['h'] / 2) - self.prev_left_lug[1]
                if dy_h > self.RELEASE_DELTA_THRESHOLD and dy_l > self.RELEASE_DELTA_THRESHOLD:
                    is_releasing = True
            self.prev_left_hook = (h['x'] + h['w'] / 2, h['y'] + h['h'] / 2)
            self.prev_left_lug = (l['x'] + l['w'] / 2, l['y'] + l['h'] / 2)
        else:
            self.prev_left_hook = None
            self.prev_left_lug = None

        # 处理右侧配对
        if right_pair:
            h, l = right_pair
            if self.prev_right_hook and self.prev_right_lug:
                dy_h = (h['y'] + h['h'] / 2) - self.prev_right_hook[1]
                dy_l = (l['y'] + l['h'] / 2) - self.prev_right_lug[1]
                if dy_h > self.RELEASE_DELTA_THRESHOLD and dy_l > self.RELEASE_DELTA_THRESHOLD:
                    is_releasing = True
            self.prev_right_hook = (h['x'] + h['w'] / 2, h['y'] + h['h'] / 2)
            self.prev_right_lug = (l['x'] + l['w'] / 2, l['y'] + l['h'] / 2)
        else:
            self.prev_right_hook = None
            self.prev_right_lug = None

        # 状态判断
        if is_releasing:
            new_state = HookState.RELEASING
        else:
            can_lift = False
            if self.active_bucket:
                can_lift = (self._side_match("left", left_pair)
                            or self._side_match("right", right_pair))
            if can_lift:
                new_state = HookState.READY
            elif self._has_intersecting_pair(left_pair, right_pair):
                new_state = HookState.HOOKING
            else:
                new_state = HookState.IDLE

        # 防抖：进出 READY 对称用 ready_debounce，非 READY → IDLE 即时，其余用 STATE_THRESHOLD
        if new_state == HookState.IDLE and self.hook_state != HookState.READY:
            threshold = 1
        elif new_state == HookState.READY or self.hook_state == HookState.READY:
            threshold = self.config.ai.ready_debounce
        else:
            threshold = self.STATE_THRESHOLD

        if new_state == self.prev_state:
            self.state_counter += 1
        else:
            self.prev_state = new_state
            self.state_counter = 1

        if self.state_counter >= threshold:
            if new_state != self.hook_state:
                self.hook_state = new_state
                self.last_state = new_state
            self.state_counter = 0

        for obj in self.left_objects + self.right_objects:
            obj.pop('paired', None)
        if left_pair:
            left_pair[0]['paired'] = True
            left_pair[1]['paired'] = True
        if right_pair:
            right_pair[0]['paired'] = True
            right_pair[1]['paired'] = True

    def _push_pair_traj(self, side, pair):
        pass

    def _update_suppress_window(self):
        pass

    def _clear_state_display(self):
        self.last_state = None

    @staticmethod
    def _lug_touches_top(l):
        return l['y'] <= 0

    def rects_intersect(self, a, b):
        ax1, ay1, ax2, ay2 = a['x'], a['y'], a['x']+a['w'], a['y']+a['h']
        bx1, by1, bx2, by2 = b['x'], b['y'], b['x']+b['w'], b['y']+b['h']
        return not (ax2 < bx1 or bx2 < ax1 or ay2 < by1 or by2 < ay1)

    def _has_intersecting_pair(self, left_pair, right_pair):
        if left_pair:
            h, l = left_pair
            if self.rects_intersect(h, l):
                return True
        if right_pair:
            h, l = right_pair
            if self.rects_intersect(h, l):
                return True
        return False

    def _side_match(self, side, pair):
        bucket = self.active_bucket
        if not bucket or not pair:
            return False
        h, l = pair
        if not self.rects_intersect(h, l):
            return False
        if self._lug_touches_top(l):
            return False
        lw, lh = l['w'], l['h']
        if lw <= 0 or lh <= 0:
            return False
        xr = (h['x'] + h['w'] / 2 - (l['x'] + l['w'] / 2)) / lw
        yr = ((h['y'] + h['h']) - (l['y'] + l['h'])) / lh
        cal_x = getattr(bucket, f"hook_lug_x_offset_ratio_{side}", None)
        cal_y = getattr(bucket, f"hook_lug_y_offset_ratio_{side}", None)
        if cal_x is None or cal_y is None:
            return False
        tol = self.config.ai.offset_ratio_tolerance
        return abs(xr - cal_x) <= tol and abs(yr - cal_y) <= tol


class TestStateMachine:
    def test_initial_state_idle(self):
        mw = MockMainWindow()
        mw.evaluate_hook_state()
        assert mw.hook_state == HookState.IDLE

    def test_no_objects_stays_idle(self):
        mw = MockMainWindow()
        for _ in range(5):
            mw.evaluate_hook_state()
        assert mw.hook_state == HookState.IDLE

    def test_hooking_requires_debounce(self):
        mw = MockMainWindow()
        hook = make_hook(0, 0)
        lug = make_lug(5, 5)  # overlapping
        mw.left_objects = [hook, lug]

        # First call: counter starts, no transition
        mw.evaluate_hook_state()
        assert mw.hook_state == HookState.IDLE
        assert mw.state_counter == 1

        # Second call: still not enough
        mw.evaluate_hook_state()
        assert mw.hook_state == HookState.IDLE

        # Third call: threshold reached, transition to HOOKING
        mw.evaluate_hook_state()
        assert mw.hook_state == HookState.HOOKING

    def test_idle_immediate_no_debounce(self):
        mw = MockMainWindow()
        hook = make_hook(0, 0)
        lug = make_lug(5, 5)
        mw.left_objects = [hook, lug]

        # Transition to HOOKING
        for _ in range(3):
            mw.evaluate_hook_state()
        assert mw.hook_state == HookState.HOOKING

        # Remove objects - should immediately go IDLE
        mw.left_objects = []
        mw.evaluate_hook_state()
        assert mw.hook_state == HookState.IDLE

    def test_lug_at_top_skips_judgment(self):
        mw = MockMainWindow()
        mw.active_bucket = calibrated_bucket()
        mw.left_objects = [make_hook(40, 40), make_lug(40, 40)]
        for _ in range(3):
            mw.evaluate_hook_state()
        assert mw.hook_state == HookState.READY
        # 吊耳贴着画面上沿 → 不判定、状态重置为 IDLE、不显示状态框
        mw.left_objects = [make_hook(40, 0), make_lug(40, 0)]
        mw.evaluate_hook_state()
        assert mw.last_state is None
        assert mw.hook_state == HookState.IDLE

    def test_ready_state(self):
        mw = MockMainWindow()
        mw.active_bucket = calibrated_bucket()
        hook = make_hook(40, 40)
        lug = make_lug(40, 40)  # 中心对齐，匹配标定偏移 0,0
        mw.left_objects = [hook, lug]

        for _ in range(3):
            mw.evaluate_hook_state()
        assert mw.hook_state == HookState.READY

    def test_releasing_state(self):
        mw = MockMainWindow()
        hook = make_hook(0, 0)
        lug = make_lug(5, 5)
        mw.left_objects = [hook, lug]

        # First establish HOOKING
        for _ in range(3):
            mw.evaluate_hook_state()
        assert mw.hook_state == HookState.HOOKING

        # Simulate continuous downward movement (>5px each frame) for debounce
        for step in range(1, 4):
            mw.left_objects = [make_hook(0, step * 6), make_lug(5, step * 6 + 5)]
            mw.evaluate_hook_state()
        assert mw.hook_state == HookState.RELEASING

    def test_releasing_upward_movement_no_change(self):
        mw = MockMainWindow()
        hook = make_hook(0, 0)
        lug = make_lug(5, 5)
        mw.left_objects = [hook, lug]

        for _ in range(3):
            mw.evaluate_hook_state()
        assert mw.hook_state == HookState.HOOKING

        # Move upward (negative dy) - should NOT trigger releasing
        mw.left_objects = [make_hook(0, -10), make_lug(5, -10)]

        for _ in range(3):
            mw.evaluate_hook_state()
        assert mw.hook_state != HookState.RELEASING

    def test_releasing_small_movement_no_change(self):
        mw = MockMainWindow()
        hook = make_hook(0, 0)
        lug = make_lug(5, 5)
        mw.left_objects = [hook, lug]

        for _ in range(3):
            mw.evaluate_hook_state()
        assert mw.hook_state == HookState.HOOKING

        # Move down by 3px (less than threshold of 5) - should NOT trigger
        mw.left_objects = [make_hook(0, 3), make_lug(5, 3)]

        for _ in range(3):
            mw.evaluate_hook_state()
        assert mw.hook_state != HookState.RELEASING

    def test_both_sides_for_releasing(self):
        mw = MockMainWindow()
        hook_l = make_hook(0, 0)
        lug_l = make_lug(5, 5)
        hook_r = make_hook(100, 0)
        lug_r = make_lug(105, 5)
        mw.left_objects = [hook_l, lug_l]
        mw.right_objects = [hook_r, lug_r]

        for _ in range(3):
            mw.evaluate_hook_state()
        assert mw.hook_state == HookState.HOOKING

        # Only left side moves down progressively
        for step in range(1, 4):
            mw.left_objects = [make_hook(0, step * 6), make_lug(5, step * 6 + 5)]
            mw.evaluate_hook_state()
        assert mw.hook_state == HookState.RELEASING

    def test_ready_requires_intersecting_pair(self):
        mw = MockMainWindow()
        mw.active_bucket = calibrated_bucket()
        # 即使标定匹配，hook/lug 不相交也不得 READY
        hook = make_hook(0, 0)
        lug = make_lug(200, 200)  # far apart
        mw.left_objects = [hook, lug]

        for _ in range(3):
            mw.evaluate_hook_state()
        # No intersecting pair -> IDLE
        assert mw.hook_state == HookState.IDLE

    def test_paired_flag_set(self):
        mw = MockMainWindow()
        hook1 = make_hook(0, 0)
        lug1 = make_lug(5, 5)
        hook2 = make_hook(100, 100)  # far away, unpaired
        lug2 = make_lug(200, 200)
        mw.left_objects = [hook1, lug1, hook2, lug2]

        mw.evaluate_hook_state()

        paired = [o for o in mw.left_objects if o.get('paired')]
        unpaired = [o for o in mw.left_objects if not o.get('paired')]
        assert len(paired) == 2  # hook1 and lug1
        assert len(unpaired) == 2  # hook2 and lug2

    def test_reset_prev_positions_when_no_pair(self):
        mw = MockMainWindow()
        hook = make_hook(0, 0)
        lug = make_lug(5, 5)
        mw.left_objects = [hook, lug]
        mw.evaluate_hook_state()
        assert mw.prev_left_hook is not None

        mw.left_objects = []
        mw.evaluate_hook_state()
        assert mw.prev_left_hook is None
        assert mw.prev_left_lug is None