Probabilistic Motion Models

use nalgebra::{Matrix3, Vector3};
use pr_core::geom::se2::SE2; // Chapter 3: exp / log / boxplus / boxminus
use rand::Rng;

/// A velocity command held constant over an interval.
#[derive(Clone, Copy, Debug)]
pub struct VelocityCmd {
    pub v: f64,
    pub omega: f64,
    pub dt: f64,
}

/// An odometry reading, already decomposed into rotate–drive–rotate.
#[derive(Clone, Copy, Debug)]
pub struct OdomDelta {
    pub rot1: f64,
    pub trans: f64,
    pub rot2: f64,
}

/// The six α's of the velocity model. Newtyped because the *order* is load
/// bearing and a bare `[f64; 6]` at a call site tells the reader nothing.
#[derive(Clone, Copy, Debug)]
pub struct VelocityAlphas(pub [f64; 6]);

/// The four α's of the odometry model.
#[derive(Clone, Copy, Debug)]
pub struct OdomAlphas(pub [f64; 4]);

/// Thrun's `sample(b²)` and `prob(a, b²)` take a **variance**. Making that a
/// type stops the single most common bug in a first implementation: handing
/// them a standard deviation, and being wrong by a square root ever after.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Variance(pub f64);

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum NoiseKind {
    /// Exact, via the Ziggurat sampler in `rand_distr`.
    Normal,
    /// Bounded support at ±√6·b, so no sample is ever absurd (Table 5.4).
    Triangular,
}

/// Every motion model in this book: a density and a sampler that must agree.
///
/// The associated `Control` is why this is a trait and not two free functions.
/// A velocity model consumes a command that exists *before* the motion; an
/// odometry model consumes a measurement that exists only after it. Nothing
/// downstream should be able to confuse them, and here nothing can.
pub trait MotionModel {
    type Control: Copy;

    /// p(x_t | u_t, x_{t-1}) — the scoring function of Tables 5.1 / 5.5.
    fn prob(&self, x_next: &SE2, u: &Self::Control, x: &SE2) -> f64;

    /// One draw from that distribution (Tables 5.3 / 5.6).
    fn sample<R: Rng + ?Sized>(&self, u: &Self::Control, x: &SE2, rng: &mut R) -> SE2;

    /// The tangent-space covariance this model induces, where it has one.
    /// `None` for models whose noise is not additive in se(2) — the odometry
    /// model, for instance.
    fn tangent_cov(&self, _u: &Self::Control) -> Option<Matrix3<f64>> {
        None
    }
}

use crate::{MotionModel, NoiseKind, VelocityAlphas, VelocityCmd, Variance};
use nalgebra::{Matrix3, Vector3};
use pr_core::geom::se2::{wrap_angle, SE2};
use rand::Rng;

pub struct VelocityModel {
    pub alphas: VelocityAlphas,
    pub noise: NoiseKind,
}

/// Below this |ω| the arc form subtracts two near-equal sines and multiplies the
/// remainder by a radius of order v/ω, so cancellation error is amplified by the
/// radius. The two error curves cross near √(2ε)/Δt ≈ 2e-8 rad/s; switching two
/// decades earlier costs under a micrometre per step and keeps the branch well
/// clear of the crossing, where neither form is comfortable.
const STRAIGHT: f64 = 1e-6;

impl VelocityModel {
    fn variances(&self, u: &VelocityCmd) -> [Variance; 3] {
        let [a1, a2, a3, a4, a5, a6] = self.alphas.0;
        let (v2, w2) = (u.v * u.v, u.omega * u.omega);
        [
            Variance(a1 * v2 + a2 * w2),
            Variance(a3 * v2 + a4 * w2),
            Variance(a5 * v2 + a6 * w2),
        ]
    }
}

impl MotionModel for VelocityModel {
    type Control = VelocityCmd;

    fn sample<R: Rng + ?Sized>(&self, u: &VelocityCmd, x: &SE2, rng: &mut R) -> SE2 {
        let [sv, sw, sg] = self.variances(u);
        let v_hat = u.v + self.noise.sample(sv, rng);
        let w_hat = u.omega + self.noise.sample(sw, rng);
        let g_hat = self.noise.sample(sg, rng);

        let theta = x.theta();
        let next_theta = theta + w_hat * u.dt;
        let (dx, dy) = if w_hat.abs() < STRAIGHT {
            (v_hat * theta.cos() * u.dt, v_hat * theta.sin() * u.dt)
        } else {
            let r = v_hat / w_hat;
            (
                -r * theta.sin() + r * next_theta.sin(),
                r * theta.cos() - r * next_theta.cos(),
            )
        };
        // γ̂ is a rotation applied *after* the arc — it moves no position.
        SE2::new(x.x() + dx, x.y() + dy, wrap_angle(next_theta + g_hat * u.dt))
    }

    fn prob(&self, x_next: &SE2, u: &VelocityCmd, x: &SE2) -> f64 {
        let inv = invert_arc(x, x_next, u.dt);
        let [sv, sw, sg] = self.variances(u);
        self.noise.prob(u.v - inv.v_hat, sv)
            * self.noise.prob(u.omega - inv.omega_hat, sw)
            * self.noise.prob(inv.gamma_hat, sg)
    }

    /// R_u = Δt² · diag(α₁v² + α₂ω², 0, α₃v² + α₄ω²). The zero is the point:
    /// a differential drive has no lateral degree of freedom to be noisy in.
    fn tangent_cov(&self, u: &VelocityCmd) -> Option<Matrix3<f64>> {
        let [sv, sw, _] = self.variances(u);
        let dt2 = u.dt * u.dt;
        Some(Matrix3::from_diagonal(&Vector3::new(dt2 * sv.0, 0.0, dt2 * sw.0)))
    }
}

pub struct ArcInversion {
    pub v_hat: f64,
    pub omega_hat: f64,
    pub gamma_hat: f64,
}

/// Table 5.1, lines 1–8: find the unique arc leaving `x` tangentially that
/// reaches `x_next`'s position, and report the controls it implies.
fn invert_arc(x: &SE2, x_next: &SE2, dt: f64) -> ArcInversion {
    let (c, s) = (x.theta().cos(), x.theta().sin());
    let (dx, dy) = (x.x() - x_next.x(), x.y() - x_next.y());
    let numer = dx * c + dy * s;
    let denom = dy * c - dx * s;

    // Displacement parallel to the heading: the centre is at infinity.
    if denom.abs() < 1e-9 {
        return ArcInversion {
            v_hat: -numer / dt,
            omega_hat: 0.0,
            gamma_hat: wrap_angle(x_next.theta() - x.theta()) / dt,
        };
    }

    let mu = 0.5 * numer / denom;
    let cx = 0.5 * (x.x() + x_next.x()) + mu * dy;
    let cy = 0.5 * (x.y() + x_next.y()) - mu * dx;
    let r_star = (x.x() - cx).hypot(x.y() - cy);
    let d_theta = wrap_angle(
        (x_next.y() - cy).atan2(x_next.x() - cx) - (x.y() - cy).atan2(x.x() - cx),
    );
    ArcInversion {
        v_hat: d_theta / dt * r_star,
        omega_hat: d_theta / dt,
        gamma_hat: wrap_angle(x_next.theta() - x.theta()) / dt - d_theta / dt,
    }
}

use crate::{MotionModel, NoiseKind, OdomAlphas, OdomDelta, Variance};
use pr_core::geom::se2::{wrap_angle, SE2};
use rand::Rng;

pub struct OdometryModel {
    pub alphas: OdomAlphas,
    pub noise: NoiseKind,
}

impl OdomDelta {
    /// Table 5.5, lines 2–4: decompose a relative odometry reading.
    ///
    /// A pure in-place spin has no direction of travel, so `atan2` would return
    /// whatever the floating-point noise in two zeros happens to encode. We give
    /// the whole rotation to `rot2` instead — otherwise a stationary turn, which
    /// is a very common command, produces a nonsense first hinge and a noise
    /// variance built from it.
    pub fn from_poses(prev: &SE2, curr: &SE2) -> Self {
        let (dx, dy) = (curr.x() - prev.x(), curr.y() - prev.y());
        let trans = dx.hypot(dy);
        if trans < 1e-9 {
            return Self { rot1: 0.0, trans: 0.0, rot2: wrap_angle(curr.theta() - prev.theta()) };
        }
        let rot1 = wrap_angle(dy.atan2(dx) - prev.theta());
        Self { rot1, trans, rot2: wrap_angle(curr.theta() - prev.theta() - rot1) }
    }

    /// Apply an exact (noise-free) decomposition to a pose.
    pub fn apply(&self, x: &SE2) -> SE2 {
        let heading = x.theta() + self.rot1;
        SE2::new(
            x.x() + self.trans * heading.cos(),
            x.y() + self.trans * heading.sin(),
            wrap_angle(heading + self.rot2),
        )
    }
}

impl MotionModel for OdometryModel {
    /// The control *is* the decomposed reading. Callers build it with
    /// `OdomDelta::from_poses(&odom_prev, &odom_now)` — never from map-frame
    /// poses, which is the mistake that makes a localizer fight its own drift.
    type Control = OdomDelta;

    fn sample<R: Rng + ?Sized>(&self, u: &OdomDelta, x: &SE2, rng: &mut R) -> SE2 {
        let [a1, a2, a3, a4] = self.alphas.0;
        let (r1sq, r2sq, tsq) = (u.rot1 * u.rot1, u.rot2 * u.rot2, u.trans * u.trans);

        // Note the minus signs: Table 5.6 *subtracts* the noise. For a symmetric
        // distribution it makes no difference in law, and keeping it makes this
        // diffable against the book.
        let perturbed = OdomDelta {
            rot1: u.rot1 - self.noise.sample(Variance(a1 * r1sq + a2 * tsq), rng),
            trans: u.trans - self.noise.sample(Variance(a3 * tsq + a4 * (r1sq + r2sq)), rng),
            rot2: u.rot2 - self.noise.sample(Variance(a1 * r2sq + a2 * tsq), rng),
        };
        perturbed.apply(x)
    }

    fn prob(&self, x_next: &SE2, u: &OdomDelta, x: &SE2) -> f64 {
        let [a1, a2, a3, a4] = self.alphas.0;
        let hat = OdomDelta::from_poses(x, x_next);
        let (r1sq, r2sq, tsq) = (hat.rot1 * hat.rot1, hat.rot2 * hat.rot2, hat.trans * hat.trans);

        self.noise.prob(wrap_angle(u.rot1 - hat.rot1), Variance(a1 * r1sq + a2 * tsq))
            * self.noise.prob(u.trans - hat.trans, Variance(a3 * tsq + a4 * (r1sq + r2sq)))
            * self.noise.prob(wrap_angle(u.rot2 - hat.rot2), Variance(a1 * r2sq + a2 * tsq))
    }
    // No `tangent_cov`: three noises entering at three points of a composition
    // are not one additive perturbation in se(2). Saying so with `None` is more
    // useful than returning an approximation a caller would trust.
}

use crate::VelocityCmd;
use nalgebra::{Cholesky, Matrix3, Vector3};
use pr_core::geom::se2::SE2;
use rand::Rng;
use rand_distr::{Distribution, StandardNormal};

/// Gaussian noise in se(2), pushed onto the group by ⊞.
///
/// `r` is an R and not a Q: motion noise, per the book-wide convention of
/// Chapter 6. It is expressed in the **body** frame, matching the right
/// perturbation x ⊞ τ = x · exp(τ) used everywhere in this book.
pub struct TangentNoiseModel {
    pub r: Matrix3<f64>,
    chol: Cholesky<f64, nalgebra::U3>,
}

impl TangentNoiseModel {
    pub fn new(r: Matrix3<f64>) -> Self {
        // A rank-2 R (no lateral slip) is legitimate, so nudge the diagonal
        // rather than refusing to factor it.
        let padded = r + Matrix3::identity() * 1e-12;
        let chol = Cholesky::new(padded).expect("R must be positive semi-definite");
        Self { r, chol }
    }

    /// τ = (vΔt, 0, ωΔt): the twist whose exponential *is* the arc update.
    pub fn twist(u: &VelocityCmd) -> Vector3<f64> {
        Vector3::new(u.v * u.dt, 0.0, u.omega * u.dt)
    }

    pub fn sample<R: Rng + ?Sized>(&self, u: &VelocityCmd, x: &SE2, rng: &mut R) -> SE2 {
        let z = Vector3::new(
            StandardNormal.sample(rng),
            StandardNormal.sample(rng),
            StandardNormal.sample(rng),
        );
        x.boxplus(Self::twist(u) + self.chol.l() * z)
    }

    /// log p(x_t | u, x_{t-1}) — a plain 3-D Gaussian in ξ = x_t ⊟ x_{t-1}.
    ///
    /// No arc inversion, no unpaid change of variables: because the chart is
    /// anchored at x_{t-1} rather than at the predicted mean, ξ = τ + w holds
    /// identically and this really is the sampler's density.
    pub fn log_prob(&self, x_next: &SE2, u: &VelocityCmd, x: &SE2) -> f64 {
        let d = x_next.boxminus(x) - Self::twist(u);
        let y = self.chol.l().solve_lower_triangular(&d).expect("L is nonsingular");
        let log_det: f64 = self.chol.l().diagonal().iter().map(|l| 2.0 * l.ln()).sum();
        -0.5 * (y.dot(&y) + log_det + 3.0 * (2.0 * std::f64::consts::PI).ln())
    }

    /// Σ_k = A Σ_{k-1} Aᵀ + R with A = Ad(exp(τ)⁻¹): the Kalman prediction
    /// equation, with the adjoint standing in for the motion Jacobian.
    pub fn compound(&self, sigma: &Matrix3<f64>, tau: &Vector3<f64>) -> Matrix3<f64> {
        let a = SE2::exp(tau).inverse().adjoint();
        a * sigma * a.transpose() + self.r
    }
}

use approx::assert_relative_eq;
use motion::{MotionModel, NoiseKind, VelocityAlphas, VelocityCmd, VelocityModel};
use nalgebra::Vector3;
use pr_core::geom::se2::SE2;
use rand::{rngs::SmallRng, SeedableRng};

const U: VelocityCmd = VelocityCmd { v: 1.0, omega: 0.5, dt: 1.0 };
const ALPHAS: VelocityAlphas = VelocityAlphas([0.02, 0.02, 0.02, 0.02, 0.0, 0.0]);

#[test]
fn arc_update_is_the_exponential_map() {
    let model = VelocityModel { alphas: VelocityAlphas([0.0; 6]), noise: NoiseKind::Normal };
    let mut rng = SmallRng::seed_from_u64(0); // noiseless: the draws are unused
    let by_arc = model.sample(&U, &SE2::identity(), &mut rng);
    let by_exp = SE2::exp(&Vector3::new(U.v * U.dt, 0.0, U.omega * U.dt));

    assert_relative_eq!(by_arc.x(), 0.958851_077_208, epsilon = 1e-11);
    assert_relative_eq!(by_arc.y(), 0.244834_876_219, epsilon = 1e-11);
    assert_relative_eq!(by_arc.theta(), 0.5, epsilon = 1e-12);
    assert_relative_eq!(by_arc.x(), by_exp.x(), epsilon = 1e-12);
    assert_relative_eq!(by_arc.y(), by_exp.y(), epsilon = 1e-12);
}

/// The chapter's headline claim, as an assertion: the sample cloud's
/// exponential coordinates have mean τ and covariance R_u.
#[test]
fn the_banana_is_gaussian_in_exponential_coordinates() {
    let model = VelocityModel { alphas: ALPHAS, noise: NoiseKind::Normal };
    let x0 = SE2::identity();
    let mut rng = SmallRng::seed_from_u64(7);

    let xis: Vec<Vector3<f64>> = (0..200_000)
        .map(|_| model.sample(&U, &x0, &mut rng).boxminus(&x0))
        .collect();

    let n = xis.len() as f64;
    let mean = xis.iter().sum::<Vector3<f64>>() / n;
    let cov = xis
        .iter()
        .map(|xi| (xi - mean) * (xi - mean).transpose())
        .sum::<nalgebra::Matrix3<f64>>()
        / (n - 1.0);

    let r_u = model.tangent_cov(&U).unwrap();
    assert_relative_eq!(mean, Vector3::new(1.0, 0.0, 0.5), epsilon = 3e-3);
    assert_relative_eq!(cov[(0, 0)], r_u[(0, 0)], max_relative = 0.02);
    assert_relative_eq!(cov[(2, 2)], r_u[(2, 2)], max_relative = 0.02);
    assert!(cov[(1, 1)] < 1e-12, "a differential drive has no lateral noise");
}

/// And the counterpart: in (x, y) the same cloud is visibly non-Gaussian.
#[test]
fn the_same_cloud_is_skewed_in_cartesian_coordinates() {
    let model = VelocityModel { alphas: ALPHAS, noise: NoiseKind::Normal };
    let mut rng = SmallRng::seed_from_u64(7);
    let poses: Vec<SE2> = (0..200_000)
        .map(|_| model.sample(&U, &SE2::identity(), &mut rng))
        .collect();

    // Skewness of the residual perpendicular to the mean heading: identically
    // zero for any Gaussian, and decidedly not zero here.
    let skew = cross_track_skew(&poses);
    assert!(skew < -0.25, "expected a pronounced left-bending banana, got {skew}");
}

/// Third standardized moment of the residual perpendicular to the mean heading.
fn cross_track_skew(poses: &[SE2]) -> f64 {
    let n = poses.len() as f64;
    let (mx, my) = poses.iter().fold((0.0, 0.0), |(a, b), p| (a + p.x() / n, b + p.y() / n));
    // Circular mean of the headings — never average θ arithmetically.
    let (sc, ss) = poses
        .iter()
        .fold((0.0, 0.0), |(c, s), p| (c + p.theta().cos(), s + p.theta().sin()));
    let (c, s) = {
        let h = ss.atan2(sc);
        (h.cos(), h.sin())
    };

    let d: Vec<f64> = poses.iter().map(|p| -(p.x() - mx) * s + (p.y() - my) * c).collect();
    let m2 = d.iter().map(|x| x * x).sum::<f64>() / n;
    let m3 = d.iter().map(|x| x * x * x).sum::<f64>() / n;
    m3 / m2.powf(1.5)
}