TIR_PROJ/tests/test_control.py

"""Control primitives: speed modulation, Strömbom, Sequential, ActiveScan."""

import math

import pytest

from herding.control.active_scan import (
    EMPTY_DEBOUNCE_STEPS, INITIAL_SCAN_STEPS, ActiveScanTeacher,
)
from herding.control.modulation import (
    MIN_SPEED, SLOW_NEAR_SHEEP, modulate_speed_near_sheep,
)
from herding.control.sequential import compute_action as sequential_action
from herding.control.strombom import (
    DELTA_DRIVE, F_FACTOR, compute_action as strombom_action,
)
from herding.control.universal import compute_action as universal_action
from herding.world.geometry import PEN_ENTRY


# ---------------------------------------------------------------------------
# Modulation
# ---------------------------------------------------------------------------

def test_modulation_empty_input_passthrough():
    assert modulate_speed_near_sheep(1.0, 0.0, (0.0, 0.0), []) == (1.0, 0.0)
    assert modulate_speed_near_sheep(1.0, 0.0, (0.0, 0.0), {}) == (1.0, 0.0)


def test_modulation_far_sheep_passthrough():
    vx, vy = modulate_speed_near_sheep(1.0, 0.0, (0.0, 0.0), [(100.0, 0.0)])
    assert (vx, vy) == (1.0, 0.0)


def test_modulation_close_sheep_min_speed():
    vx, vy = modulate_speed_near_sheep(1.0, 0.0, (0.0, 0.0), [(0.0, 0.0)])
    assert math.isclose(vx, MIN_SPEED)
    assert vy == 0.0


def test_modulation_preserves_direction():
    vx, vy = modulate_speed_near_sheep(0.6, 0.8, (0.0, 0.0), [(1.0, 0.0)])
    ratio = math.hypot(vx, vy)
    # Direction preserved.
    assert math.isclose(vx / ratio, 0.6, abs_tol=1e-6)
    assert math.isclose(vy / ratio, 0.8, abs_tol=1e-6)


def test_modulation_linear_ramp_midpoint():
    vx, _ = modulate_speed_near_sheep(1.0, 0.0, (0.0, 0.0),
                                      [(SLOW_NEAR_SHEEP / 2, 0.0)])
    expected = MIN_SPEED + (1.0 - MIN_SPEED) * 0.5
    assert math.isclose(vx, expected, abs_tol=1e-6)


def test_modulation_accepts_dict_input():
    vx_list, _ = modulate_speed_near_sheep(1.0, 0.0, (0.0, 0.0),
                                           [(1.0, 0.0)])
    vx_dict, _ = modulate_speed_near_sheep(1.0, 0.0, (0.0, 0.0),
                                           {"t0": (1.0, 0.0)})
    assert math.isclose(vx_list, vx_dict)


# ---------------------------------------------------------------------------
# Strömbom
# ---------------------------------------------------------------------------

def test_strombom_empty_input_idle():
    vx, vy, mode = strombom_action((0.0, 0.0), {}, PEN_ENTRY)
    assert (vx, vy, mode) == (0.0, 0.0, "idle")


def test_strombom_tight_flock_drives():
    # A tight 3-sheep cluster centred at (0, 8): radius < F_FACTOR·√3.
    sheep = {"s0": (0.0, 8.0), "s1": (0.5, 8.5), "s2": (-0.5, 8.0)}
    vx, vy, mode = strombom_action((0.0, 0.0), sheep, PEN_ENTRY)
    assert mode == "drive"
    assert math.isclose(math.hypot(vx, vy), 1.0, abs_tol=1e-3)


def test_strombom_scattered_flock_collects():
    # Sparse, max radius > F_FACTOR·√n.
    sheep = {"s0": (10.0, 10.0), "s1": (-10.0, -10.0), "s2": (0.0, 0.0)}
    _vx, _vy, mode = strombom_action((0.0, 0.0), sheep, PEN_ENTRY)
    assert mode == "collect"


def test_strombom_ignores_already_penned_sheep():
    """Sheep south of the gate plane are excluded from the active set."""
    sheep = {
        "s_active": (5.0, 5.0),
        "s_penned": (11.5, -20.0),
    }
    # With one active sheep, Strömbom drives (radius = 0 < threshold).
    _vx, _vy, mode = strombom_action((0.0, 0.0), sheep, PEN_ENTRY)
    assert mode == "drive"


# ---------------------------------------------------------------------------
# Sequential
# ---------------------------------------------------------------------------

def test_sequential_empty_input_idle():
    vx, vy, mode = sequential_action((0.0, 0.0), {}, PEN_ENTRY)
    assert (vx, vy, mode) == (0.0, 0.0, "idle")


def test_sequential_targets_closest_to_pen():
    # With 2 sheep (≤ STRAGGLER_THRESHOLD), sequential goes straight to
    # "targeted" phase and pushes the sheep nearest to the pen.
    near = (10.0, -5.0)       # closer to pen entry (11.5, -15)
    far = (-10.0, 10.0)
    sheep = {"near": near, "far": far}
    vx, vy, mode = sequential_action((0.0, 0.0), sheep, PEN_ENTRY)
    assert mode == "targeted"
    # Dog should be directed toward near sheep (south-east), not far (north-west).
    assert vx > 0 and vy < 0


def test_sequential_collects_when_scattered():
    # With >STRAGGLER_THRESHOLD sheep and radius > F_FACTOR*sqrt(n):
    # should use collect (Strombom) not targeted.
    sheep = {f"s{i}": pos for i, pos in enumerate([
        (12.0, 10.0), (-12.0, 10.0), (0.0, 12.0),
        (12.0, -12.0), (-10.0, -8.0),
    ])}
    _vx, _vy, mode = sequential_action((0.0, 0.0), sheep, PEN_ENTRY)
    assert mode in ("collect", "drive")


def test_sequential_drives_when_compact():
    # Compact flock of 5 sheep near centre — should drive, not collect.
    sheep = {f"s{i}": (float(i) * 0.3, float(i) * 0.3)
             for i in range(5)}
    _vx, _vy, mode = sequential_action((0.0, 5.0), sheep, PEN_ENTRY)
    assert mode == "drive"


# ---------------------------------------------------------------------------
# ActiveScan wrapper
# ---------------------------------------------------------------------------

def test_active_scan_initial_phase_rotates():
    teacher = ActiveScanTeacher(strombom_action)
    # First call → opening rotation regardless of input.
    vx, vy, omega, mode = teacher(
        (0.0, 0.0), 0.0, {"s0": (5.0, 0.0)}, PEN_ENTRY)
    assert mode == "scan_initial"
    assert omega == 0.0
    assert math.isclose(math.hypot(vx, vy), 1.0, abs_tol=1e-6)


def test_active_scan_hands_off_to_base_after_opener():
    teacher = ActiveScanTeacher(strombom_action, initial_scan_steps=2)
    # Burn through the opener.
    for _ in range(2):
        teacher((0.0, 0.0), 0.0, {"s0": (0.0, 8.0)}, PEN_ENTRY)
    _vx, _vy, _omega, mode = teacher(
        (0.0, 0.0), 0.0, {"s0": (0.0, 8.0)}, PEN_ENTRY)
    # Either drive (Strömbom mode label) or collect; not scan_initial.
    assert "scan" not in mode


def test_active_scan_holds_last_action_on_brief_empty():
    teacher = ActiveScanTeacher(strombom_action, initial_scan_steps=1)
    # Step once (opening), then once with a visible sheep — sets last_action.
    teacher((0.0, 0.0), 0.0, {}, PEN_ENTRY)
    teacher((0.0, 0.0), 0.0, {"s0": (0.0, 8.0)}, PEN_ENTRY)
    last = teacher.last_action
    # Now a single empty frame → hold.
    vx, vy, _omega, mode = teacher((0.0, 0.0), 0.0, {}, PEN_ENTRY)
    assert mode == "hold"
    assert (vx, vy) == last


def test_active_scan_explores_after_sustained_empty():
    teacher = ActiveScanTeacher(strombom_action, initial_scan_steps=1)
    teacher((0.0, 0.0), 0.0, {}, PEN_ENTRY)  # opener
    for _ in range(EMPTY_DEBOUNCE_STEPS):
        last_vx, last_vy, _omega, mode = teacher(
            (5.0, 5.0), 0.0, {}, PEN_ENTRY)
    assert mode in ("explore", "scan_at_centre")


def test_active_scan_preserves_mecanum_omega():
    """Regression: ActiveScanTeacher must propagate omega from a mecanum
    base teacher, not silently drop it. Without this, BC mecanum demos
    have omega=0 everywhere and the policy never learns to rotate.
    """
    teacher = ActiveScanTeacher(universal_action, initial_scan_steps=1)
    # Burn the opener so we exit phase 1.
    teacher((0.0, 0.0), 0.0, {"s0": (8.0, 8.0)}, PEN_ENTRY,
            drive_mode="mecanum")
    # Place a sheep off to the side so the dog needs to face it.
    # Dog at origin facing +x (heading=0); target at (0, 8) → desired
    # heading +π/2, so omega should be positive.
    vx, vy, omega, mode = teacher(
        (0.0, 0.0), 0.0, {"s0": (0.0, 8.0)}, PEN_ENTRY,
        drive_mode="mecanum")
    assert mode in ("collect", "drive", "recovery")
    assert abs(omega) > 0.05, f"omega should be non-zero on mecanum, got {omega}"


def test_active_scan_reset_clears_state():
    teacher = ActiveScanTeacher(strombom_action, initial_scan_steps=5)
    for _ in range(3):
        teacher((0.0, 0.0), 0.0, {}, PEN_ENTRY)
    assert teacher.step == 3
    teacher.reset()
    assert teacher.step == 0
    assert teacher.empty_streak == 0