"""Differential-drive kinematics matching the Webots robot specs.

The Webots controllers and the training env both use these helpers so the
sim and the real (Webots) physics agree to first order. They do not model
slip, wheel acceleration limits, or contact forces — Webots does that for
us at inference time. The training env has to be close enough that a
policy trained against this kinematic model still works when handed off
to ODE physics.
"""

import math


def kinematics_step(x, y, h, w_left, w_right, wheel_radius, wheel_base, dt):
    """Integrate one step of differential-drive forward kinematics.

    Inputs
    ------
    x, y : robot position (m)
    h    : robot heading (rad), 0 = +x axis
    w_left, w_right : wheel angular velocities (rad/s)
    wheel_radius, wheel_base : robot dimensions (m)
    dt   : timestep (s)

    Returns (new_x, new_y, new_h).
    """
    v = (w_right + w_left) * wheel_radius * 0.5
    omega = (w_right - w_left) * wheel_radius / wheel_base
    new_x = x + v * math.cos(h) * dt
    new_y = y + v * math.sin(h) * dt
    new_h = math.atan2(math.sin(h + omega * dt), math.cos(h + omega * dt))
    return new_x, new_y, new_h


def velocity_to_wheels(vx, vy, h, max_linear, wheel_radius, max_wheel_omega,
                       k_turn=4.0):
    """Convert a desired (vx, vy) intent in [-1, 1]^2 to wheel speeds.

    Mirrors ``drive_action`` in controllers/shepherd_dog/shepherd_dog.py:
    forward speed scales by ``cos(err)`` (clamped to ±90°), and a P
    controller on heading error contributes the wheel-rate differential.
    """
    speed_ms = math.hypot(vx, vy) * max_linear
    if speed_ms < 1e-3:
        return 0.0, 0.0
    target_h = math.atan2(vy, vx)
    err = math.atan2(math.sin(target_h - h), math.cos(target_h - h))
    clamped_err = max(-math.pi / 2, min(math.pi / 2, err))
    fwd_ms = speed_ms * math.cos(clamped_err)
    fwd_rad = fwd_ms / wheel_radius
    turn = k_turn * err
    left = max(-max_wheel_omega, min(max_wheel_omega, fwd_rad - turn))
    right = max(-max_wheel_omega, min(max_wheel_omega, fwd_rad + turn))
    return left, right


def heading_speed_to_wheels(heading, speed_motor, h, max_wheel_omega,
                            k_turn=4.0):
    """Sheep variant: speed already expressed in motor (wheel rad/s) units.

    Matches the existing sheep controller (``controllers/sheep/sheep.py``)
    where ``speed = max(WANDER_SPEED, min(FLEE_SPEED, mag * 3.0))`` and
    these constants are wheel angular velocities, not linear m/s.
    """
    err = math.atan2(math.sin(heading - h), math.cos(heading - h))
    fwd = max(0.0, math.cos(err)) * speed_motor
    turn = k_turn * err
    left = max(-max_wheel_omega, min(max_wheel_omega, fwd - turn))
    right = max(-max_wheel_omega, min(max_wheel_omega, fwd + turn))
    return left, right