Bug Algorithms: Planning with a Contact Sensor

use nalgebra::{Point2, Vector2};

/// One ray of ρ_R(x, ·): the raw distance, or +∞ when the ray saturates at R.
/// A scan is sensor data and nothing else — no scene access anywhere in this file.
pub struct Scan {
    pub x: Point2<f64>,
    pub angles: Vec<f64>,
    pub ranges: Vec<f64>, // f64::INFINITY where ρ ≥ R
    pub max_range: f64,
}

/// A maximal run of rays over which ρ_R is finite and continuous (C §2.2).
/// `to < from` means the run wraps past bearing 0.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Interval { pub from: usize, pub to: usize, pub count: usize }

impl Scan {
    /// D(x) = min_s ρ(x, s): the shortest ray (C eq. 2.4, problem 2.1). With finitely many
    /// rays this never undershoots D and overshoots by D(1 − cos(π/n)) on a flat wall.
    pub fn distance(&self) -> f64 {
        self.ranges.iter().cloned().fold(f64::INFINITY, f64::min)
    }

    /// ∇D(x): the unit vector along the shortest ray, pointing *away* from the obstacle.
    pub fn gradient(&self) -> Vector2<f64> {
        let (k, _) = self.argmin();
        let s = self.angles[k];
        -Vector2::new(s.cos(), s.sin())
    }

    /// (∇D)⊥ with its sign fixed by the side the obstacle is kept on. `Turn::Left`
    /// keeps it on the robot's right (clockwise around a convex obstacle).
    pub fn tangent(&self, turn: Turn) -> Vector2<f64> {
        let g = self.gradient();
        let t = Vector2::new(-g.y, g.x);
        match turn { Turn::Left => -t, Turn::Right => t }
    }

    /// Two consecutive hits belong to one surface when their endpoints are closer than
    /// the spacing a wall at steep incidence would produce: κ·ρ·Δθ, with κ = 8 good to
    /// an incidence of ~83°, floored at `jump_abs` so a hairline step is not a corner.
    fn continuous(&self, k: usize, l: usize, kappa: f64, jump_abs: f64) -> bool {
        let (rk, rl) = (self.ranges[k], self.ranges[l]);
        if !rk.is_finite() || !rl.is_finite() { return false; }
        let dtheta = std::f64::consts::TAU / self.ranges.len() as f64;
        let gap = (self.point(k) - self.point(l)).norm();
        gap <= jump_abs.max(kappa * rk.min(rl) * dtheta)
    }

    /// The intervals of continuity. O(n_rays). Runs straddling the 0/2π seam are merged;
    /// an all-finite, all-continuous scan (inside a closed room) is one interval with no
    /// endpoints — the robot sees no corner to aim for.
    pub fn continuity_intervals(&self, kappa: f64, jump_abs: f64) -> Vec<Interval> {
        let n = self.ranges.len();
        let wrap = |k: isize| ((k % n as isize + n as isize) % n as isize) as usize;
        // Start at a break so wrapped runs come out whole.
        let Some(start) = (0..n).find(|&k| !self.continuous(wrap(k as isize - 1), k, kappa, jump_abs)) else {
            return if self.ranges[0].is_finite() { vec![Interval { from: 0, to: n - 1, count: n }] } else { vec![] };
        };
        let (mut out, mut k, mut visited) = (Vec::new(), start, 0);
        while visited < n {
            if !self.ranges[k].is_finite() { k = wrap(k as isize + 1); visited += 1; continue; }
            let mut count = 1;
            while visited + count < n
                && self.continuous(wrap(k as isize + count as isize - 1), wrap(k as isize + count as isize), kappa, jump_abs)
            { count += 1; }
            out.push(Interval { from: k, to: wrap(k as isize + count as isize - 1), count });
            k = wrap(k as isize + count as isize);
            visited += count;
        }
        out
    }

    fn point(&self, k: usize) -> Point2<f64> {
        let (r, a) = (self.ranges[k], self.angles[k]);
        self.x + r * Vector2::new(a.cos(), a.sin())
    }
}

use nalgebra::{Point2, Vector2};

#[derive(Clone, Copy, Debug)]
pub enum Turn { Left, Right }

pub struct TraceParams {
    /// Δλ, the predictor step along the tangent.
    pub step: f64,
    /// Newton iterations per correction; three is usually one too many.
    pub newton_iters: usize,
    /// Stop correcting once |G| falls below this.
    pub tol: f64,
    /// |D_Y G| below this means the correcting line is tangent to the level set.
    pub singular_tol: f64,
    pub max_steps: usize,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum TraceStop { Closed, Reached, Singular, Budget, Diverged }

/// One predictor–corrector step, kept whole so a figure can draw every hop.
pub struct TraceRecord {
    pub from: Point2<f64>,
    pub tangent: Vector2<f64>,
    pub predicted: Point2<f64>,
    pub corrections: Vec<Point2<f64>>,
    /// |G(y_h)| for h = 0, 1, 2, … — the squares Theorem D.2.1 promises.
    pub residuals: Vec<f64>,
    pub dy_g: f64,
}

pub struct Trace {
    pub points: Vec<Point2<f64>>,
    pub records: Vec<TraceRecord>,
    pub stop: TraceStop,
    pub length: f64,
    pub max_residual: f64,
}

/// Newton along a fixed direction n — the corrector of eq. 2.5 with D_Y G = ∇G·n.
/// `singular` is set when the directional derivative vanished before the tolerance was met.
fn correct_along(
    g: &impl Fn(Point2<f64>) -> f64,
    dg: &impl Fn(Point2<f64>) -> Vector2<f64>,
    y0: Point2<f64>,
    n: Vector2<f64>,
    p: &TraceParams,
) -> (Point2<f64>, Vec<Point2<f64>>, Vec<f64>, f64, bool) {
    let (mut y, mut corr, mut res) = (y0, Vec::new(), vec![g(y0).abs()]);
    let dy_g0 = dg(y0).dot(&n);
    for _ in 0..p.newton_iters {
        let gy = g(y);
        if gy.abs() <= p.tol { break; }
        let dy_g = dg(y).dot(&n);
        if dy_g.abs() < p.singular_tol {
            return (y, corr, res, dy_g0, true); // Theorem D.1.1's hypothesis has failed: say so
        }
        y -= n * (gy / dy_g);                   // eq. 2.5
        corr.push(y);
        res.push(g(y).abs());
    }
    (y, corr, res, dy_g0, false)
}

/// Trace the zero set of G from x₀ until it closes, `stop(x)` fires, the corrector goes
/// singular, or the budget runs out. Generic over G: Chapters 7–10 pass their own.
pub fn trace_curve(
    g: impl Fn(Point2<f64>) -> f64,
    dg: impl Fn(Point2<f64>) -> Vector2<f64>,
    x0: Point2<f64>,
    direction: Turn,
    p: TraceParams,
    stop: impl Fn(Point2<f64>) -> bool,
) -> Trace {
    // x₀ need not be on the curve: Newton along the full gradient first.
    let mut x = x0;
    for _ in 0..8 {
        let (gx, grad) = (g(x), dg(x));
        if gx.abs() <= p.tol || grad.norm_squared() < 1e-18 { break; }
        x -= grad * (gx / grad.norm_squared());
    }
    let start = x;
    let (mut points, mut records, mut length) = (vec![start], Vec::new(), 0.0);
    let mut max_residual = g(start).abs();
    let mut outcome = TraceStop::Budget;

    for _ in 0..p.max_steps {
        let grad = dg(x);
        let n = grad / grad.norm();
        let perp = Vector2::new(-n.y, n.x);
        let t = match direction { Turn::Left => -perp, Turn::Right => perp };
        let predicted = x + t * p.step;                                   // predict
        let (y, corr, res, dy_g, singular) = correct_along(&g, &dg, predicted, n, &p); // correct
        records.push(TraceRecord { from: x, tangent: t, predicted, corrections: corr, residuals: res, dy_g });
        if singular { outcome = TraceStop::Singular; break; }
        let hop = (y - x).norm();
        if !hop.is_finite() || hop > 4.0 * p.step { outcome = TraceStop::Diverged; break; }
        length += hop;
        x = y;
        points.push(x);
        max_residual = max_residual.max(g(x).abs());
        if stop(x) { outcome = TraceStop::Reached; break; }
        if length > 3.0 * p.step && (x - start).norm() < 0.75 * p.step {
            length += (x - start).norm();
            points.push(start);
            outcome = TraceStop::Closed;
            break;
        }
    }
    Trace { points, records, stop: outcome, length, max_residual }
}

/// Convenience: the offset curve of a scene at W*, with G = D − W* and ∇G = ∇D from
/// `Collision::witness` — the nearest wall point and the unit vector away from it.
pub fn offset_curve(scene: &Scene, w_star: f64, x0: Point2<f64>, dir: Turn, p: TraceParams) -> Trace {
    trace_curve(
        |x| scene.distance(x) - w_star,
        |x| scene.witness(x).gradient,
        x0, dir, p, |_| false,
    )
}

use nalgebra::Point2;
use crate::follow::{Contact, Follower, Piece};

#[derive(Clone, Copy, PartialEq, Eq)]
enum Phase { ToGoal, Circumnavigate, Return }

pub struct BugState {
    pub x: Point2<f64>,
    pub path: Vec<Point2<f64>>,
    pub length: f64,
    pub hits: Vec<Point2<f64>>,
    pub leaves: Vec<Point2<f64>>,
    pub done: bool,
    pub failed: bool,
}

/// C Alg. 1. `step` is the motion resolution; the leave point and the perimeter are
/// exact regardless of it, because the bookkeeping is done on the follower's pieces.
pub struct Bug1<'c, C: Contact> {
    contact: &'c C,
    goal: Point2<f64>,
    step: f64,
    turn: Turn,
    pub state: BugState,
    phase: Phase,
    follower: Option<C::Follower>,
    /// Best leave candidate on this obstacle: point, arc length from the hit, goal distance.
    best: (Point2<f64>, f64, f64),
    remaining: f64,
}

impl<'c, C: Contact> Bug1<'c, C> {
    pub fn step(&mut self) {
        if self.state.done || self.state.failed { return; }
        match self.phase {
            Phase::ToGoal => {
                // Lines 2–3: toward the goal until we touch something.
                let heading = (self.goal - self.state.x).normalize();
                let (x, hit) = self.contact.move_toward(self.state.x, self.goal, self.step);
                self.push(x);
                if (x - self.goal).norm() < 1e-9 { self.state.done = true; return; }
                if hit {
                    self.state.hits.push(x);
                    self.follower = Some(self.contact.follow(x, heading, self.turn));
                    self.best = (x, 0.0, (x - self.goal).norm());
                    self.phase = Phase::Circumnavigate;
                }
            }
            Phase::Circumnavigate => {
                // Line 4: a full loop, remembering the closest perimeter point (line 5).
                let f = self.follower.as_mut().unwrap();
                for Piece { a, b, s0 } in f.advance(self.step) {
                    let (q, t, d) = closest_on_segment(self.goal, a, b);
                    if d < self.best.2 - 1e-12 { self.best = (q, s0 + t * (b - a).norm(), d); }
                    self.push(b);
                }
                if f.closed() {
                    // Line 6: the shorter way round.
                    let p = f.perimeter().unwrap();
                    let (fwd, back) = (self.best.1, p - self.best.1);
                    if back < fwd { f.reverse(); }
                    self.remaining = fwd.min(back);
                    self.phase = Phase::Return;
                }
            }
            Phase::Return => {
                let f = self.follower.as_mut().unwrap();
                let ds = self.step.min(self.remaining);
                for Piece { b, .. } in f.advance(ds) { self.push(b); }
                self.remaining -= ds;
                if self.remaining <= 1e-12 {
                    self.push(self.best.0);
                    self.state.leaves.push(self.best.0);
                    // Line 7: would heading for the goal re-enter the obstacle?
                    if self.contact.blocked(self.best.0, self.goal) { self.state.failed = true; return; }
                    self.follower = None;
                    self.phase = Phase::ToGoal;
                }
            }
        }
    }

    fn push(&mut self, p: Point2<f64>) {
        self.state.length += (p - self.state.x).norm();
        self.state.path.push(p);
        self.state.x = p;
    }
}

fn main() {
    // WO₁ = [2,4]×[−1,1] (p₁ = 8), WO₂ = [6,7]×[−2,2] (p₂ = 10); q_start = (0,0), q_goal = (10,0).
    let scene = Scene::from_polygons(&[rect(2.0, -1.0, 4.0, 1.0), rect(6.0, -2.0, 7.0, 2.0)]);
    let (start, goal) = (Point2::new(0.0, 0.0), Point2::new(10.0, 0.0));
    let contact = PolygonContact::new(&scene);

    let b1 = bug1(&contact, start, goal, 0.05, Turn::Left);
    let b2 = bug2(&contact, start, goal, 0.05, Turn::Left);
    let d = (goal - start).norm();
    let p: Vec<f64> = scene.polygons().iter().map(perimeter).collect();
    let n: Vec<usize> = scene.polygons().iter().map(|o| line_crossings(o, start, goal)).collect();

    println!("Bug1   length {:.4}  bound d + 1.5·Σp = {:.4}   hits {}  leaves {}",
             b1.length(), d + 1.5 * p.iter().sum::<f64>(), fmt(&b1.hits), fmt(&b1.leaves));
    println!("Bug2   length {:.4}  bound d + ½·Σ nᵢpᵢ = {:.4}   hits {}  leaves {}",
             b2.length(), d + 0.5 * n.iter().zip(&p).map(|(n, p)| *n as f64 * p).sum::<f64>(),
             fmt(&b2.hits), fmt(&b2.leaves));

    // Tangent Bug's first decision with infinite range: the √5 + √65 tie.
    let tb = TangentBug::new(&scene, start, goal, TangentOptions { range: f64::INFINITY, clearance: 0.1, ..Default::default() });
    let e = tb.evaluate(start);
    println!("TBug∞  first subgoal: tie {:.4} between O=(2,1) and O=(2,-1) → {:?}; length {:.4}",
             e.candidates[0].h, Turn::Left, tangent_bug(&scene, start, goal, 0.05, f64::INFINITY).length());

    // The tracer against Derivation 6: p + 2πW* = 8 + π.
    let tr = offset_curve(&scene.only(0), 0.5, Point2::new(1.4, 0.0), Turn::Left,
                          TraceParams { step: 0.02, newton_iters: 3, tol: 1e-12, singular_tol: 0.2, max_steps: 5000 });
    println!("offset curve of WO1 at W*=0.5: {:?} after {} steps of Δλ=0.02, length {:.4} (8+π = {:.4}), max |G| {:.1e}",
             tr.stop, tr.points.len() - 1, tr.length, 8.0 + std::f64::consts::PI, tr.max_residual);
}

Bug1   length 34.0000  bound d + 1.5·Σp = 37.0000   hits (2,0) (6,0)  leaves (4,0) (7,0)
Bug2   length 16.0000  bound d + ½·Σ nᵢpᵢ = 28.0000   hits (2,0) (6,0)  leaves (4,0) (7,0)
TBug∞  first subgoal: tie 10.2983 between O=(2,1) and O=(2,-1) → Left; length 11.2128
offset curve of WO1 at W*=0.5: Closed after 558 steps of Δλ=0.02, length 11.1414 (8+π = 11.1416), max |G| 3.0e-14

use proptest::prelude::*;

#[test]
fn bug1_two_rectangles() {
    let (scene, start, goal) = two_rectangles();
    let s = bug1(&PolygonContact::new(&scene), start, goal, 0.05, Turn::Left);
    assert_relative_eq!(s.length(), 34.0, epsilon = 1e-6);
    assert_eq!(s.hits, vec![Point2::new(2.0, 0.0), Point2::new(6.0, 0.0)]);
    assert_eq!(s.leaves, vec![Point2::new(4.0, 0.0), Point2::new(7.0, 0.0)]);
}

#[test]
fn bug2_two_rectangles() {
    let (scene, start, goal) = two_rectangles();
    assert_relative_eq!(bug2(&PolygonContact::new(&scene), start, goal, 0.05, Turn::Left).length(), 16.0, epsilon = 1e-6);
}

proptest! {
    /// Both bounds, on seeded fields of 1–6 disjoint convex polygons with random queries.
    /// `Unreachable` must come back iff the goal's component differs from the start's.
    #[test]
    fn bounds_hold(seed in 0u64..500) {
        let field = random_convex_field(SmallRng::seed_from_u64(seed));
        let contact = PolygonContact::new(&field.scene);
        let d = (field.goal - field.start).norm();
        let p: f64 = field.scene.polygons().iter().map(perimeter).sum();
        let np: f64 = field.scene.polygons().iter()
            .map(|o| line_crossings(o, field.start, field.goal) as f64 * perimeter(o)).sum();
        for turn in [Turn::Left, Turn::Right] {
            match bug1(&contact, field.start, field.goal, 0.05, turn) {
                BugOutcome::Reached(path) => prop_assert!(path.length() <= d + 1.5 * p + 1e-9),
                BugOutcome::Unreachable(_) => prop_assert!(!field.connected()),
            }
            match bug2(&contact, field.start, field.goal, 0.05, turn) {
                BugOutcome::Reached(path) => prop_assert!(path.length() <= d + 0.5 * np + 1e-9),
                BugOutcome::Unreachable(_) => prop_assert!(!field.connected()),
            }
        }
    }
}

#[test]
fn offset_square_length() {
    let tr = offset_curve(&Scene::from_polygons(&[rect(2.0, -1.0, 4.0, 1.0)]), 0.5,
                          Point2::new(1.4, 0.0), Turn::Left, TraceParams::fine(0.02));
    assert_eq!(tr.stop, TraceStop::Closed);
    assert_relative_eq!(tr.length, 8.0 + std::f64::consts::PI, epsilon = 1e-3);
    assert!(tr.max_residual < 1e-8);
}

#[test]
fn tracer_reports_singular() {
    // Into the kidney's dent at W* = 0.25: the offset set has a cusp, D_Y G → 0.
    let tr = offset_curve(&kidney(), 0.25, Point2::new(0.0, -1.75), Turn::Left, TraceParams::fine(0.05));
    assert_eq!(tr.stop, TraceStop::Singular);
}