"""Shared low-level control helpers used by every dog mode.

Centralised here so the BC student, Strömbom, Sequential, and the DAgger
teacher all apply identical post-processing to their action outputs.
The downstream wheel-velocity layer (``herding.diffdrive``) is unchanged.
"""

from __future__ import annotations

import math


# Speed-modulation: scale action magnitude down when close to the
# nearest sheep. Stops the dog from charging in at full speed and
# scattering the flock. Action norm linearly ramps from MIN_SPEED at
# distance 0 to 1.0 at SLOW_NEAR_SHEEP.
SLOW_NEAR_SHEEP = 2.5
MIN_SPEED = 0.30


def modulate_speed_near_sheep(
    vx: float, vy: float,
    dog_xy: tuple[float, float],
    sheep_positions,
    slow_dist: float = SLOW_NEAR_SHEEP,
    min_scale: float = MIN_SPEED,
) -> tuple[float, float]:
    """Scale (vx, vy) magnitude down when close to the nearest sheep.

    ``sheep_positions`` accepts either a ``{name: (x, y)}`` dict
    (matching what the trackers emit) or an iterable of ``(x, y)``
    tuples. Empty input → action returned unchanged.

    The intent direction is preserved; only magnitude is reduced. With
    ``slow_dist=2.5`` and ``min_scale=0.3``, an action that started at
    norm 1 is multiplied by 0.3 right next to a sheep, by 0.65 at 1 m
    away, and by 1.0 once the nearest sheep is ≥ 2.5 m off.
    """
    if not sheep_positions:
        return vx, vy
    if hasattr(sheep_positions, "values"):
        positions = sheep_positions.values()
    else:
        positions = sheep_positions
    nearest = float("inf")
    for sx, sy in positions:
        d = math.hypot(sx - dog_xy[0], sy - dog_xy[1])
        if d < nearest:
            nearest = d
    if nearest >= slow_dist or nearest == float("inf"):
        return vx, vy
    scale = min_scale + (1.0 - min_scale) * (nearest / slow_dist)
    return vx * scale, vy * scale