frontier_code/solutions/control_lab.py

"""Frontier F01 lab: feedback control for a hovering drone and a path-following car.

Python 3.13 standard library only. SI units throughout: metres, seconds,
kilograms, newtons, radians.

Two plants share one idea: measure, compare with a target, and push the
system toward the target every tick.

* ``simulate_hover`` flies a drone straight up to a target altitude with a PID
  controller whose thrust is limited to what the motors can produce.
* ``simulate_pure_pursuit`` drives a kinematic bicycle-model car along a path
  by steering toward a point one look-ahead distance ahead.
"""

from __future__ import annotations

import math
import random
from dataclasses import dataclass

G = 9.81  # gravitational acceleration, m/s^2


# ---------------------------------------------------------------------------
# PID controller
# ---------------------------------------------------------------------------

@dataclass
class PID:
    """Discrete PID controller.

    The derivative acts on the *measurement*, not the error, so a sudden
    setpoint change does not produce a spike ("derivative kick"). Conditional
    integration stops the integral from growing while the output is pinned at
    a limit and the error would push it further past that limit (anti-windup).
    """

    kp: float
    ki: float
    kd: float
    out_min: float = -math.inf
    out_max: float = math.inf
    anti_windup: bool = True
    integral: float = 0.0
    prev_measurement: float | None = None

    def reset(self) -> None:
        self.integral = 0.0
        self.prev_measurement = None

    def update(self, setpoint: float, measurement: float, dt: float) -> float:
        if dt <= 0.0:
            raise ValueError("dt must be positive")
        error = setpoint - measurement
        if self.prev_measurement is None:
            measurement_rate = 0.0
        else:
            measurement_rate = (measurement - self.prev_measurement) / dt
        self.prev_measurement = measurement

        candidate = self.integral + error * dt
        raw = self.kp * error + self.ki * candidate - self.kd * measurement_rate
        out = min(max(raw, self.out_min), self.out_max)

        pinned_high = out >= self.out_max and error > 0.0
        pinned_low = out <= self.out_min and error < 0.0
        if not self.anti_windup or not (pinned_high or pinned_low):
            self.integral = candidate
        return out


# ---------------------------------------------------------------------------
# Drone altitude hold
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class HoverSample:
    t: float          # time, s
    z: float          # true altitude, m
    vz: float         # vertical velocity, m/s
    thrust: float     # commanded total thrust, N


def simulate_hover(
    kp: float,
    ki: float,
    kd: float,
    *,
    mass: float = 1.5,
    drag: float = 0.3,
    setpoint: float = 10.0,
    duration: float = 10.0,
    dt: float = 0.01,
    feedforward_mass: float | None = None,
    payload: float = 0.0,
    max_thrust_ratio: float = 2.0,
    noise_std: float = 0.0,
    seed: int = 0,
    anti_windup: bool = True,
) -> list[HoverSample]:
    """Fly from z = 0 to ``setpoint`` and return the trajectory.

    Plant: (mass + payload) * z'' = T - (mass + payload) * g - drag * z'.
    The controller believes the vehicle weighs ``feedforward_mass`` and adds
    that hover thrust as feedforward (None means no feedforward). Thrust is
    limited to [0, max_thrust_ratio * mass * g]; motors cannot pull down.
    The altimeter adds Gaussian noise with standard deviation ``noise_std``.
    """
    if dt <= 0.0 or duration <= 0.0:
        raise ValueError("dt and duration must be positive")
    true_mass = mass + payload
    t_max = max_thrust_ratio * mass * G
    ff = 0.0 if feedforward_mass is None else feedforward_mass * G
    pid = PID(kp, ki, kd, out_min=-ff, out_max=t_max - ff, anti_windup=anti_windup)
    rng = random.Random(seed)

    z = vz = 0.0
    samples: list[HoverSample] = []
    steps = round(duration / dt)
    for k in range(steps + 1):
        t = k * dt
        measured = z + (rng.gauss(0.0, noise_std) if noise_std > 0.0 else 0.0)
        thrust = ff + pid.update(setpoint, measured, dt)
        samples.append(HoverSample(t, z, vz, thrust))
        # Semi-implicit Euler: update velocity first, then position.
        accel = (thrust - true_mass * G - drag * vz) / true_mass
        vz += accel * dt
        z += vz * dt
    return samples


@dataclass(frozen=True)
class StepMetrics:
    overshoot_pct: float        # peak above target as % of the step size
    rise_time: float | None     # 10% -> 90% of the step, s
    settling_time: float | None  # last entry into the +/-2% band, s
    steady_state_error: float   # target - mean altitude over the final second, m


def step_metrics(samples: list[HoverSample], setpoint: float, start: float = 0.0) -> StepMetrics:
    step = setpoint - start
    if step == 0.0:
        raise ValueError("setpoint must differ from start")
    peak = max(s.z for s in samples)
    overshoot = max(0.0, (peak - setpoint) / abs(step) * 100.0)

    t10 = next((s.t for s in samples if (s.z - start) >= 0.1 * step), None)
    t90 = next((s.t for s in samples if (s.z - start) >= 0.9 * step), None)
    rise = None if t10 is None or t90 is None else t90 - t10

    band = 0.02 * abs(step)
    settling = None
    for i in range(len(samples) - 1, -1, -1):
        if abs(samples[i].z - setpoint) > band:
            settling = samples[i + 1].t if i + 1 < len(samples) else None
            break
    else:
        settling = samples[0].t

    last_t = samples[-1].t
    tail = [s.z for s in samples if s.t >= last_t - 1.0]
    sse = setpoint - sum(tail) / len(tail)
    return StepMetrics(overshoot, rise, settling, sse)


def finite_difference(values: list[float], dt: float) -> list[float]:
    """Backward-difference derivative: (v[k] - v[k-1]) / dt, first entry 0."""
    return [0.0] + [(values[k] - values[k - 1]) / dt for k in range(1, len(values))]


# ---------------------------------------------------------------------------
# Car path following: kinematic bicycle model + pure pursuit
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class CarState:
    x: float     # rear-axle position, m
    y: float
    yaw: float   # heading, rad
    steer: float  # front-wheel angle, rad


def bicycle_step(state: CarState, speed: float, steer: float, wheelbase: float, dt: float) -> CarState:
    """Advance the kinematic bicycle model (reference point: rear axle)."""
    x = state.x + speed * math.cos(state.yaw) * dt
    y = state.y + speed * math.sin(state.yaw) * dt
    yaw = state.yaw + speed / wheelbase * math.tan(steer) * dt
    return CarState(x, y, yaw, steer)


def pure_pursuit_steer(state: CarState, target: tuple[float, float], wheelbase: float) -> float:
    """Steering angle that puts the rear axle on a circle through ``target``.

    alpha is the angle from the heading to the target; ld is the distance to it.
    Curvature kappa = 2 sin(alpha) / ld, and the bicycle model needs
    tan(delta) = wheelbase * kappa.
    """
    dx, dy = target[0] - state.x, target[1] - state.y
    ld = math.hypot(dx, dy)
    if ld == 0.0:
        return 0.0
    alpha = math.atan2(dy, dx) - state.yaw
    kappa = 2.0 * math.sin(alpha) / ld
    return math.atan(wheelbase * kappa)


def _closest_on_segment(p, a, b):
    ax, ay = a
    bx, by = b
    vx, vy = bx - ax, by - ay
    seg2 = vx * vx + vy * vy
    u = 0.0 if seg2 == 0.0 else max(0.0, min(1.0, ((p[0] - ax) * vx + (p[1] - ay) * vy) / seg2))
    cx, cy = ax + u * vx, ay + u * vy
    return (cx, cy), math.hypot(p[0] - cx, p[1] - cy), u


def cross_track_error(point: tuple[float, float], path: list[tuple[float, float]]) -> float:
    """Unsigned distance from ``point`` to the polyline ``path``."""
    return min(_closest_on_segment(point, path[i], path[i + 1])[1] for i in range(len(path) - 1))


def lookahead_point(point, path, lookahead: float, start_index: int = 0):
    """First point on the path, at or after the closest point, that is ``lookahead`` away.

    Returns (target, closest_segment_index). Falls back to the final path point
    when the path ends inside the look-ahead circle.

    The closest-point search only looks a bounded distance ahead of the
    previous index. A global search breaks on paths that revisit the same
    place (a loop, a figure eight): the "closest" segment can belong to a later
    lap, and the car then chases a point it has not reached yet.
    """
    window = 2.0 * lookahead + 5.0
    best_i, best_d, best_u = start_index, math.inf, 0.0
    travelled = 0.0
    for i in range(start_index, len(path) - 1):
        _, d, u = _closest_on_segment(point, path[i], path[i + 1])
        if d < best_d:
            best_i, best_d, best_u = i, d, u
        travelled += math.dist(path[i], path[i + 1])
        if travelled > window:
            break
    for i in range(best_i, len(path) - 1):
        a, b = path[i], path[i + 1]
        u0 = best_u if i == best_i else 0.0
        # Solve |a + u (b - a) - p| = lookahead for u in [u0, 1] (largest root).
        vx, vy = b[0] - a[0], b[1] - a[1]
        fx, fy = a[0] - point[0], a[1] - point[1]
        qa = vx * vx + vy * vy
        qb = 2.0 * (fx * vx + fy * vy)
        qc = fx * fx + fy * fy - lookahead * lookahead
        disc = qb * qb - 4.0 * qa * qc
        if qa == 0.0 or disc < 0.0:
            continue
        u = (-qb + math.sqrt(disc)) / (2.0 * qa)
        if u0 <= u <= 1.0:
            return (a[0] + u * vx, a[1] + u * vy), best_i
    return path[-1], best_i


def simulate_pure_pursuit(
    path: list[tuple[float, float]],
    *,
    start: CarState,
    speed: float = 10.0,
    lookahead: float = 8.0,
    wheelbase: float = 2.9,
    max_steer: float = math.radians(35.0),
    max_steer_rate: float | None = None,
    dt: float = 0.02,
    duration: float = 20.0,
) -> tuple[list[CarState], list[float]]:
    """Drive along ``path``; return the states and cross-track errors per tick.

    ``max_steer_rate`` (rad/s) models a steering actuator that cannot turn the
    wheels instantly. None means an ideal, instantaneous actuator.
    """
    if lookahead <= 0.0 or dt <= 0.0:
        raise ValueError("lookahead and dt must be positive")
    state = start
    states, errors = [state], [cross_track_error((state.x, state.y), path)]
    index = 0
    for _ in range(round(duration / dt)):
        target, index = lookahead_point((state.x, state.y), path, lookahead, index)
        steer = max(-max_steer, min(max_steer, pure_pursuit_steer(state, target, wheelbase)))
        if max_steer_rate is not None:
            step = max_steer_rate * dt
            steer = state.steer + max(-step, min(step, steer - state.steer))
        state = bicycle_step(state, speed, steer, wheelbase, dt)
        states.append(state)
        errors.append(cross_track_error((state.x, state.y), path))
    return states, errors


def straight_path(length: float = 400.0, step: float = 1.0) -> list[tuple[float, float]]:
    n = round(length / step)
    return [(i * step, 0.0) for i in range(n + 1)]


def circle_path(radius: float, turns: float = 1.0, points: int = 720) -> list[tuple[float, float]]:
    """Counter-clockwise circle centred at (0, radius), starting at the origin heading +x."""
    total = 2.0 * math.pi * turns
    return [
        (radius * math.sin(total * k / points), radius - radius * math.cos(total * k / points))
        for k in range(points + 1)
    ]


def lane_change_path(length: float = 200.0, shift: float = 3.5, start: float = 60.0,
                     width: float = 40.0, step: float = 0.5) -> list[tuple[float, float]]:
    """Straight road with a smooth (cosine) lane change of ``shift`` metres."""
    pts = []
    n = round(length / step)
    for i in range(n + 1):
        x = i * step
        if x <= start:
            y = 0.0
        elif x >= start + width:
            y = shift
        else:
            y = shift * 0.5 * (1.0 - math.cos(math.pi * (x - start) / width))
        pts.append((x, y))
    return pts