"""
kinematics.py — Tank / tracked-vehicle kinematics for SaltyTank (Issue #122).

Drive modes
───────────
  "2wd"     — standard differential drive, 2 ESC channels (CH1=left, CH3=right).
              Mirrored to 4 channels by the caller (CH1=CH2=left, CH3=CH4=right).

  "4wd"     — skid-steer 4-channel.  Left pair (CH1+CH2) share one speed;
              right pair (CH3+CH4) share one speed.  No slip compensation.

  "tracked" — tank kinematics with track-slip compensation.
              Commanded angular is scaled up by 1/(1-slip_factor) so the
              vehicle achieves the desired turn rate despite track slip.
              Odometry angular is scaled down by (1-slip_factor) to reflect
              the fact that wheel-speed difference overestimates actual yaw rate.

Slip compensation model
───────────────────────
  slip_factor ∈ [0.0, 0.99]

  0.0  → rigid wheel / no slip (identical to standard diff-drive)
  0.1  → light slip (rubber tracks on concrete)
  0.3  → moderate slip (rubber tracks on grass / gravel)  ← default
  0.5  → heavy slip (metal tracks on mud)

  Command path  : angular_cmd = angular_z / (1 - slip_factor)
  Odometry path : angular_odom = (v_right - v_left) / B × (1 - slip_factor)

  The two corrections are complementary:  a vehicle commanded at ω rad/s and
  estimated by odometry will yield consistent trajectories regardless of the
  chosen slip_factor value (as long as it is consistent between command and odom).

Zero-radius turns
─────────────────
  Setting linear_x=0 and angular_z≠0 produces equal-and-opposite track speeds,
  which is a valid zero-radius pivot for tracked vehicles.  The slip factor still
  applies; set slip_factor=0 for a surface with known high traction.

Pure module — no ROS2 dependency; fully unit-testable.
"""

from __future__ import annotations

import math


# ── Constants ────────────────────────────────────────────────────────────────

SLIP_MIN = 0.0
SLIP_MAX = 0.99       # must stay < 1.0 to avoid division-by-zero
DRIVE_MODES = ("2wd", "4wd", "tracked")


# ── Low-level helpers ────────────────────────────────────────────────────────

def _clamp_slip(slip_factor: float) -> float:
    return max(SLIP_MIN, min(SLIP_MAX, slip_factor))


# ── Track-speed computation ──────────────────────────────────────────────────

def tank_speeds(
    linear_x: float,
    angular_z: float,
    track_width: float,
    slip_factor: float = 0.0,
) -> tuple[float, float]:
    """
    Unicycle → tank (tracked) differential drive with optional slip compensation.

    Parameters
    ----------
    linear_x    : forward velocity (m/s)
    angular_z   : desired yaw rate (rad/s, +ve = CCW / left turn)
    track_width : centre-to-centre track separation (m)
    slip_factor : track slip coefficient ∈ [0.0, 0.99]; 0 = no slip

    Returns
    -------
    (v_left, v_right) in m/s
    """
    slip = _clamp_slip(slip_factor)
    # Scale angular command up to achieve desired turn rate despite slip
    angular_cmd = angular_z / (1.0 - slip) if slip > 0.0 else angular_z
    half = angular_cmd * track_width / 2.0
    return linear_x - half, linear_x + half


def skid_steer_speeds(
    linear_x: float,
    angular_z: float,
    track_width: float,
) -> tuple[float, float]:
    """
    Standard unicycle → differential drive (no slip compensation).
    Used for 2WD and 4WD modes.
    """
    half = angular_z * track_width / 2.0
    return linear_x - half, linear_x + half


# ── m/s → ESC PWM ────────────────────────────────────────────────────────────

def speed_to_pwm(
    speed_ms: float,
    max_speed_ms: float,
    pwm_neutral_us: int,
    pwm_range_us: int,
) -> int:
    """
    Convert wheel/track speed (m/s) to ESC PWM pulse width (µs).

    Linear mapping:
        pwm = neutral + (speed / max_speed) * range

    Clamped to [neutral - range, neutral + range].
    """
    if max_speed_ms <= 0.0:
        return pwm_neutral_us
    normalised = max(-1.0, min(1.0, speed_ms / max_speed_ms))
    pwm = pwm_neutral_us + normalised * pwm_range_us
    lo  = pwm_neutral_us - pwm_range_us
    hi  = pwm_neutral_us + pwm_range_us
    return int(round(max(lo, min(hi, pwm))))


# ── Variant-aware speed computation ─────────────────────────────────────────

def compute_track_speeds(
    linear_x: float,
    angular_z: float,
    track_width: float,
    max_speed_ms: float,
    drive_mode: str = "tracked",
    slip_factor: float = 0.3,
) -> tuple[float, float]:
    """
    Compute and clip left/right track/wheel speeds for the selected mode.

    Returns
    -------
    (v_left, v_right) clipped to ±max_speed_ms.

    All modes produce a 2-element tuple; the caller expands to 4 channels
    for 4WD mode by mirroring (vl, vl, vr, vr).
    """
    def clip(v: float) -> float:
        return max(-max_speed_ms, min(max_speed_ms, v))

    if drive_mode == "tracked":
        vl, vr = tank_speeds(linear_x, angular_z, track_width, slip_factor)
    else:
        # "2wd" or "4wd" — standard diff-drive, no slip compensation
        vl, vr = skid_steer_speeds(linear_x, angular_z, track_width)

    return clip(vl), clip(vr)


def expand_to_4ch(v_left: float, v_right: float) -> tuple[float, float, float, float]:
    """
    Expand left/right speeds to 4-channel ESC layout.

    Channel assignment:
        CH1 = left-front   CH2 = left-rear
        CH3 = right-front  CH4 = right-rear
    """
    return v_left, v_left, v_right, v_right


# ── Odometry helper ──────────────────────────────────────────────────────────

def odometry_from_track_speeds(
    v_left: float,
    v_right: float,
    track_width: float,
    slip_factor: float = 0.0,
) -> tuple[float, float]:
    """
    Estimate (linear_x, angular_z) from left/right track speeds.

    Applies slip compensation to the angular estimate:
        angular_z_odom = (v_right - v_left) / B × (1 - slip_factor)

    This is the inverse of tank_speeds for consistent dead-reckoning.

    Parameters
    ----------
    v_left, v_right : measured track speeds (m/s)
    track_width     : centre-to-centre track separation (m)
    slip_factor     : same value used in the command path

    Returns
    -------
    (linear_x, angular_z) in m/s and rad/s
    """
    linear_x = (v_left + v_right) / 2.0
    if track_width > 0.0:
        slip = _clamp_slip(slip_factor)
        angular_z = (v_right - v_left) / track_width * (1.0 - slip)
    else:
        angular_z = 0.0
    return linear_x, angular_z
