"""Universal shepherd teacher — Strömbom core + mecanum omega + straggler recovery.

The core collect/drive logic is **identical** to :mod:`strombom` (same
``F_FACTOR``, ``DELTA_COLLECT``, ``DELTA_DRIVE`` thresholds and target
computation) so it inherits the proven ~100 % success rate at n ≤ 8.
Two additions make it useful as a universal teacher:

1. **Omega for mecanum.**  When ``drive_mode="mecanum"``, the teacher
   outputs a non-zero ``omega`` channel so the dog **faces the
   direction of travel**.  During collect the dog faces the target
   sheep; during drive it faces the pen.  This gives the BC student a
   real rotation signal to learn from.

2. **Last-straggler recovery.**  When exactly one sheep remains active
   and it is near the gate, the dog positions itself behind that
   straggler (opposite the gate) and pushes it straight through.  This
   handles the edge case where the last sheep circles the gate posts.

Call signature::

    vx, vy, omega, mode = compute_action(
        dog_xy, dog_heading, sheep_positions, pen_target,
        drive_mode="differential",
    )

For differential drive ``omega`` is always 0.0 and can be ignored.
"""

import math

from herding.world.geometry import (
    PEN_ENTRY, GATE_X, GATE_Y, in_pen,
)

# ---------------------------------------------------------------------------
# Tuning constants — match Strömbom exactly for proven success rates.
# ---------------------------------------------------------------------------

F_FACTOR = 4.0          # collect/drive threshold scaled by √n
DELTA_COLLECT = 1.5      # standoff behind the furthest sheep
DELTA_DRIVE = 2.0        # standoff behind flock CoM

# Omega gain for mecanum (how strongly the dog turns to face target)
OMEGA_GAIN = 0.6

# Recovery: push small flocks (≤ RECOVERY_MAX_N) through the gate one
# sheep at a time. n=1 alone is not enough — at n=2..3 on the round
# field the flock is too small to self-cohere through the 3 m gate but
# the standard collect/drive standoff just orbits them. Push the sheep
# nearest the gate first; once it pens, the rule re-applies to the next.
RECOVERY_MAX_N = 3
RECOVERY_GATE_DIST = 8.0   # only when target sheep is this close to gate
RECOVERY_PUSH_DIST = 1.2   # stand-off behind sheep, away from gate


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _unit(x, y):
    d = math.hypot(x, y)
    if d < 1e-6:
        return 0.0, 0.0
    return x / d, y / d


def _is_active(x, y) -> bool:
    return (not in_pen(x, y)) and y > GATE_Y


def _angle_diff(a, b):
    """Signed shortest angular difference a - b, in [-π, π]."""
    return math.atan2(math.sin(a - b), math.cos(a - b))


def _gate_center():
    """Centre of the gate opening."""
    return (0.5 * (GATE_X[0] + GATE_X[1]), GATE_Y)


# ---------------------------------------------------------------------------
# Core teacher
# ---------------------------------------------------------------------------

def compute_action(dog_xy, dog_heading, sheep_positions,
                   pen_target=PEN_ENTRY, drive_mode="differential"):
    """Return ``(vx, vy, omega, mode)``.

    Parameters
    ----------
    dog_xy : (float, float)
        Dog position in world frame.
    dog_heading : float
        Dog heading in world frame (rad), 0 = +x axis.
    sheep_positions : dict[str, (float, float)]
        Visible sheep positions.
    pen_target : (float, float)
        Centre of the pen gate (defaults to geometry.PEN_ENTRY).
    drive_mode : str
        ``"differential"`` or ``"mecanum"``.

    Returns
    -------
    vx, vy : float
        Velocity intent in [-1, 1].
    omega : float
        Yaw intent in [-1, 1] (0 for differential).
    mode : str
        Phase label: ``"idle"``, ``"collect"``, ``"drive"``, ``"recovery"``.
    """
    active = [(x, y) for (x, y) in sheep_positions.values()
              if _is_active(x, y)]
    if not active:
        return 0.0, 0.0, 0.0, "idle"

    n = len(active)
    com_x = sum(p[0] for p in active) / n
    com_y = sum(p[1] for p in active) / n
    dists = [math.hypot(p[0] - com_x, p[1] - com_y) for p in active]
    radius = max(dists)

    # ---- Small-flock recovery (push sheep through the gate one by one) ----
    # Triggers when the active flock is small (≤ RECOVERY_MAX_N) and the
    # sheep nearest the gate is close enough that direct pushing works.
    # For larger flocks the standard collect/drive logic handles them.
    gc = _gate_center()
    if n <= RECOVERY_MAX_N:
        # Pick the sheep closest to the gate as the recovery target —
        # finishing that one first reduces the active count and lets the
        # remaining sheep get their own recovery turn.
        gate_dists = [math.hypot(p[0] - gc[0], p[1] - gc[1]) for p in active]
        target_idx = min(range(n), key=lambda i: gate_dists[i])
        sx, sy = active[target_idx]
        d_to_gate = gate_dists[target_idx]
        if d_to_gate < RECOVERY_GATE_DIST:
            dx_g = sx - gc[0]
            dy_g = sy - gc[1]
            d_g = math.hypot(dx_g, dy_g)
            if d_g > 0.3:
                ux, uy = dx_g / d_g, dy_g / d_g
            else:
                ux, uy = 0.0, 1.0
            tx = sx + RECOVERY_PUSH_DIST * ux
            ty = sy + RECOVERY_PUSH_DIST * uy
            ax, ay = _unit(tx - dog_xy[0], ty - dog_xy[1])
            mode = "recovery"
            face_target = (sx, sy)
            omega = 0.0
            if drive_mode == "mecanum":
                desired = math.atan2(
                    face_target[1] - dog_xy[1],
                    face_target[0] - dog_xy[0],
                )
                err = _angle_diff(desired, dog_heading)
                omega = max(-1.0, min(1.0, OMEGA_GAIN * err / math.pi))
            return ax, ay, omega, mode

    # ---- Standard Strömbom collect/drive (proven core) ----
    if radius > F_FACTOR * math.sqrt(n):
        # Collect: aim behind the furthest sheep, opposite the CoM.
        idx = max(range(n), key=lambda i: dists[i])
        sx, sy = active[idx]
        ux, uy = _unit(sx - com_x, sy - com_y)
        tx, ty = sx + DELTA_COLLECT * ux, sy + DELTA_COLLECT * uy
        mode = "collect"
        face_target = (sx, sy)
    else:
        # Drive: aim behind the CoM, opposite the pen.
        ux, uy = _unit(com_x - pen_target[0], com_y - pen_target[1])
        tx, ty = com_x + DELTA_DRIVE * ux, com_y + DELTA_DRIVE * uy
        mode = "drive"
        face_target = pen_target

    ax, ay = _unit(tx - dog_xy[0], ty - dog_xy[1])

    # ---- Omega (mecanum only) ----
    omega = 0.0
    if drive_mode == "mecanum" and mode != "idle":
        desired_heading = math.atan2(
            face_target[1] - dog_xy[1],
            face_target[0] - dog_xy[0],
        )
        err = _angle_diff(desired_heading, dog_heading)
        omega = max(-1.0, min(1.0, OMEGA_GAIN * err / math.pi))

    return ax, ay, omega, mode


def compute_action_diff(dog_xy, dog_heading, sheep_positions,
                        pen_target=PEN_ENTRY):
    """Compatibility wrapper returning ``(vx, vy, mode)`` — same as Strömbom.

    Use this when plugging into existing differential-drive code that
    doesn't expect omega.
    """
    vx, vy, _omega, mode = compute_action(
        dog_xy, dog_heading, sheep_positions, pen_target,
        drive_mode="differential",
    )
    return vx, vy, mode