"""Differential-drive kinematics, shared by the env and Webots controllers.

First-order rigid-body model — no slip, wheel-accel limits, or contact
forces. Webots' ODE physics handles those at inference; the env stays
close enough to first order that a policy trained here transfers.
"""

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]² to wheel speeds.

    Forward speed scales by ``cos(err)`` (clamped to ±90°); 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 in wheel rad/s, target as a heading angle."""
    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