Probabilistic Sensor Models

use nalgebra::{Matrix3, Point2};
use pr_core::geom::se2::SE2;
use sim::World; // Chapter 4's polyline world *is* the map, until Chapter 13

/// One sweep of a range finder: ranges plus the bearings they were taken at.
///
/// Bearings are stored, not assumed. A sub-sampled scan is still a valid `Scan`,
/// which is what makes tempering-by-subsampling a two-line operation instead of
/// an index-arithmetic bug farm.
#[derive(Clone, Debug)]
pub struct Scan {
    pub ranges: Vec<f64>,
    /// Bearings relative to the robot's heading, radians.
    pub bearings: Vec<f64>,
    pub max_range: f64,
}

/// Anything Chapters 11–12 can localize with.
///
/// Log-likelihood, not likelihood: the product for any hypothesis that is even
/// slightly wrong underflows `f64` within a hundred beams, and tempering is a
/// single multiply in log space.
pub trait SensorModel {
    fn log_likelihood(&self, scan: &Scan, x: &SE2, map: &World) -> f64;

    /// Cheap enough to call per particle per step? Chapter 12 branches on it.
    fn is_cheap(&self) -> bool {
        false
    }
}

/// Θ — the six intrinsic parameters of the beam model.
///
/// The four weights are *not* a `[f64; 4]`: naming them stops the classic
/// transposition bug where `z_short` and `z_max` get swapped and the model
/// starts believing every dropout is a person.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct BeamIntrinsics {
    pub z_hit: f64,
    pub z_short: f64,
    pub z_max: f64,
    pub z_rand: f64,
    pub sigma_hit: f64,
    pub lambda_short: f64,
    pub max_range: f64,
}

impl BeamIntrinsics {
    /// The weights must live on the simplex; everything downstream assumes it.
    pub fn is_normalized(&self) -> bool {
        (self.z_hit + self.z_short + self.z_max + self.z_rand - 1.0).abs() < 1e-9
    }
}

/// A range-bearing-signature detection, Thrun's f = (r, φ, s).
#[derive(Clone, Copy, Debug)]
pub struct Feature {
    pub r: f64,
    pub phi: f64,
    pub s: f64,
}

/// The map's view of a landmark. `Q` is the feature noise covariance Q_t.
#[derive(Clone, Copy, Debug)]
pub struct Landmark {
    pub xy: Point2<f64>,
    pub signature: f64,
}

pub type FeatureNoise = Matrix3<f64>;

use crate::{BeamIntrinsics, Scan, SensorModel};
use pr_core::geom::se2::SE2;
use sim::{ray_cast, World};
use statrs::distribution::{ContinuousCDF, Normal};

pub struct BeamModel {
    pub intr: BeamIntrinsics,
    /// Use every `stride`-th beam. See `tempering.rs` for why this is a lie
    /// that is sometimes worth telling.
    pub stride: usize,
}

/// Within this of `max_range`, a reading *is* a max-range reading. Floating
/// point equality against a sensor's advertised limit is not a plan.
const MAX_EPS: f64 = 1e-9;

impl BeamModel {
    /// p(z | x, m) for one beam — Table 6.1, lines 4–7.
    pub fn beam_likelihood(&self, z: f64, z_star: f64) -> f64 {
        let i = &self.intr;

        // p_hit: a Gaussian truncated to [0, z_max]. η is the reciprocal of the
        // mass that survives truncation; without it, hypotheses whose predicted
        // range sits near a boundary are silently favoured.
        let p_hit = if (0.0..=i.max_range).contains(&z) {
            let n = Normal::new(z_star, i.sigma_hit).expect("sigma_hit > 0");
            let mass = n.cdf(i.max_range) - n.cdf(0.0);
            if mass > 1e-12 {
                let d = (z - z_star) / i.sigma_hit;
                (-0.5 * d * d).exp() / (i.sigma_hit * (2.0 * std::f64::consts::PI).sqrt()) / mass
            } else {
                0.0
            }
        } else {
            0.0
        };

        // p_short: the nearest hit of a Poisson clutter field of rate λ along
        // the ray, truncated at the wall.
        let p_short = if z >= 0.0 && z <= z_star && z_star > 0.0 {
            let eta = 1.0 / (1.0 - (-i.lambda_short * z_star).exp());
            eta * i.lambda_short * (-i.lambda_short * z).exp()
        } else {
            0.0
        };

        // p_max is a point mass: it integrates to 1 on its own.
        let p_max = if z >= i.max_range - MAX_EPS { 1.0 } else { 0.0 };
        let p_rand = if z >= 0.0 && z < i.max_range { 1.0 / i.max_range } else { 0.0 };

        i.z_hit * p_hit + i.z_short * p_short + i.z_max * p_max + i.z_rand * p_rand
    }

    /// The four *unweighted* component densities, for the EM E-step.
    pub fn components(&self, z: f64, z_star: f64) -> [f64; 4] {
        let one_hot = |c: usize| {
            let mut i = self.intr;
            i.z_hit = (c == 0) as u8 as f64;
            i.z_short = (c == 1) as u8 as f64;
            i.z_max = (c == 2) as u8 as f64;
            i.z_rand = (c == 3) as u8 as f64;
            BeamModel { intr: i, stride: 1 }.beam_likelihood(z, z_star)
        };
        [one_hot(0), one_hot(1), one_hot(2), one_hot(3)]
    }
}

impl SensorModel for BeamModel {
    /// `beam_range_finder_model` — Table 6.1, in log space.
    fn log_likelihood(&self, scan: &Scan, x: &SE2, map: &World) -> f64 {
        let mut q = 0.0;
        for k in (0..scan.ranges.len()).step_by(self.stride.max(1)) {
            let z_star = ray_cast(map, x.x(), x.y(), x.theta() + scan.bearings[k], scan.max_range);
            q += self.beam_likelihood(scan.ranges[k], z_star).max(1e-300).ln();
        }
        q
    }
}

/// `learn_intrinsic_parameters` — Table 6.2.
///
/// Takes ranges paired with the ray-cast range each *should* have had. Caching
/// `z_star` outside the loop is what turns EM from "one ray cast per beam per
/// iteration" into "one ray cast per beam, ever".
pub fn learn_intrinsics(
    z: &[f64],
    z_star: &[f64],
    init: BeamIntrinsics,
    iters: usize,
) -> (BeamIntrinsics, Vec<f64>) {
    let mut intr = init;
    let mut trace = Vec::with_capacity(iters);

    for _ in 0..iters {
        let model = BeamModel { intr, stride: 1 };
        let (mut s_hit, mut s_short, mut s_max, mut s_rand) = (0.0, 0.0, 0.0, 0.0);
        let (mut sq_hit, mut z_short_sum, mut ll) = (0.0, 0.0, 0.0);

        for (&zi, &si) in z.iter().zip(z_star) {
            let c = model.components(zi, si);
            let w = [intr.z_hit * c[0], intr.z_short * c[1], intr.z_max * c[2], intr.z_rand * c[3]];
            let total: f64 = w.iter().sum();
            if total <= 0.0 {
                continue; // no cause explains this reading; refuse to guess
            }
            let e = [w[0] / total, w[1] / total, w[2] / total, w[3] / total];

            s_hit += e[0];
            s_short += e[1];
            s_max += e[2];
            s_rand += e[3];
            sq_hit += e[0] * (zi - si).powi(2);
            z_short_sum += e[1] * zi;
            ll += total.max(1e-300).ln();
        }

        trace.push(ll / z.len() as f64);
        let n = z.len() as f64;
        intr = BeamIntrinsics {
            z_hit: s_hit / n,
            z_short: s_short / n,
            z_max: s_max / n,
            z_rand: s_rand / n,
            // A component with no responsibility keeps its shape parameter
            // rather than recomputing it from an empty sum.
            sigma_hit: if s_hit > 1e-9 { (sq_hit / s_hit).sqrt().max(1e-3) } else { intr.sigma_hit },
            lambda_short: if z_short_sum > 1e-9 {
                (s_short / z_short_sum).clamp(1e-3, 50.0)
            } else {
                intr.lambda_short
            },
            ..intr
        };
    }
    (intr, trace)
}

use crate::{BeamIntrinsics, Scan, SensorModel};
use pr_core::geom::se2::SE2;
use sim::World;

/// Distance to the nearest occupied cell, in metres, on a regular grid.
pub struct DistanceField {
    cells: Vec<f32>,
    nx: usize,
    ny: usize,
    res: f64,
    origin: (f64, f64),
}

/// Stands in for +∞ with a real number, so the envelope arithmetic below never
/// evaluates ∞ − ∞.
const INF: f64 = 1e20;

/// Felzenszwalb–Huttenlocher, Figure 1: the lower envelope of the parabolas
/// z ↦ (z − q)² + f(q). All parabolas share a curvature, so any two meet
/// exactly once and the envelope can be built by one left-to-right sweep.
fn dt_1d(f: &[f64], out: &mut [f64]) {
    let n = f.len();
    let mut v = vec![0usize; n]; // parabola indices currently on the envelope
    let mut z = vec![0.0f64; n + 1]; // breakpoints between them
    let mut k = 0;
    v[0] = 0;
    z[0] = -INF;
    z[1] = INF;

    for q in 1..n {
        let mut s = intersect(f, q, v[k]);
        while k > 0 && s <= z[k] {
            k -= 1; // the top parabola is entirely above the new one: pop it
            s = intersect(f, q, v[k]);
        }
        k += 1;
        v[k] = q;
        z[k] = s;
        z[k + 1] = INF;
    }

    let mut k = 0;
    for q in 0..n {
        while z[k + 1] < q as f64 {
            k += 1;
        }
        let d = q as f64 - v[k] as f64;
        out[q] = d * d + f[v[k]];
    }
}

fn intersect(f: &[f64], q: usize, p: usize) -> f64 {
    ((f[q] + (q * q) as f64) - (f[p] + (p * p) as f64)) / (2 * q as f64 - 2 * p as f64)
}

impl DistanceField {
    /// Rasterize the map, transform down the columns, transform along the rows,
    /// take one square root per cell. O(N), and exact on the lattice.
    ///
    /// `rasterize` walks each wall at a third of a cell and returns the seed
    /// grid: 0.0 at a cell a wall passes through, INF elsewhere. Indices are
    /// clamped, not dropped — the Apartment's shell lies exactly on the
    /// bounding box, and a wall that falls off the grid leaves the whole border
    /// looking like open space.
    pub fn from_map(map: &World, res: f64) -> Self {
        let (nx, ny, mut f) = rasterize(map, res);
        let mut buf = vec![0.0; nx.max(ny)];

        for i in 0..nx {
            let col: Vec<f64> = (0..ny).map(|j| f[j * nx + i]).collect();
            dt_1d(&col, &mut buf[..ny]);
            for j in 0..ny {
                f[j * nx + i] = buf[j];
            }
        }
        for j in 0..ny {
            let row: Vec<f64> = (0..nx).map(|i| f[j * nx + i]).collect();
            dt_1d(&row, &mut buf[..nx]);
            f[j * nx..j * nx + nx].copy_from_slice(&buf[..nx]);
        }

        let cells = f.iter().map(|d| (d.sqrt() * res) as f32).collect();
        Self { cells, nx, ny, res, origin: (map.bounds.min_x, map.bounds.min_y) }
    }

    /// Nearest-cell lookup, clamped at the border. Queries outside the map are
    /// charged the extra Euclidean distance so an off-map endpoint reads as
    /// *unlikely* rather than merely uninformative.
    #[inline]
    pub fn dist(&self, x: f64, y: f64) -> f64 {
        let i = (((x - self.origin.0) / self.res - 0.5).round() as isize)
            .clamp(0, self.nx as isize - 1) as usize;
        let j = (((y - self.origin.1) / self.res - 0.5).round() as isize)
            .clamp(0, self.ny as isize - 1) as usize;
        self.cells[j * self.nx + i] as f64
    }
}

pub struct LikelihoodField {
    pub field: DistanceField,
    pub intr: BeamIntrinsics,
    /// Sensor pose in the robot frame — Thrun's (x_k,sens, y_k,sens).
    pub sensor_offset: (f64, f64),
}

impl SensorModel for LikelihoodField {
    /// `likelihood_field_range_finder_model` — Table 6.3. No ray casting, no
    /// correspondence, and max-range readings are skipped outright: "I saw
    /// nothing" is not evidence about where the nearest wall is.
    fn log_likelihood(&self, scan: &Scan, x: &SE2, _map: &World) -> f64 {
        let (c, s) = (x.theta().cos(), x.theta().sin());
        let sx = x.x() + c * self.sensor_offset.0 - s * self.sensor_offset.1;
        let sy = x.y() + s * self.sensor_offset.0 + c * self.sensor_offset.1;

        let mut q = 0.0;
        for (k, &z) in scan.ranges.iter().enumerate() {
            if z >= scan.max_range - 1e-9 {
                continue;
            }
            let a = x.theta() + scan.bearings[k];
            let d = self.field.dist(sx + z * a.cos(), sy + z * a.sin());
            let g = (-0.5 * (d / self.intr.sigma_hit).powi(2)).exp()
                / (self.intr.sigma_hit * (2.0 * std::f64::consts::PI).sqrt());
            q += (self.intr.z_hit * g + self.intr.z_rand / self.intr.max_range)
                .max(1e-300)
                .ln();
        }
        q
    }

    fn is_cheap(&self) -> bool {
        true
    }
}

use crate::{Feature, FeatureNoise, Landmark};
use pr_core::geom::se2::{wrap_angle, SE2};
use rand::distr::{Distribution, Uniform};
use rand::Rng;
use rand_distr::Normal;

pub struct LandmarkModel {
    /// Q_t = diag(σ_r², σ_φ², σ_s²).
    pub q: FeatureNoise,
}

impl LandmarkModel {
    /// `landmark_model_known_correspondence` — Table 6.4, in log space.
    pub fn log_likelihood(&self, f: &Feature, j: &Landmark, x: &SE2) -> f64 {
        let (dx, dy) = (j.xy.x - x.x(), j.xy.y - x.y());
        let r_hat = dx.hypot(dy);
        // The wrap is not optional: an unwrapped 2π error in the bearing
        // residual rejects the correct landmark and accepts a wrong one.
        let phi_hat = wrap_angle(dy.atan2(dx) - x.theta());

        let e = [f.r - r_hat, wrap_angle(f.phi - phi_hat), f.s - j.signature];
        (0..3)
            .map(|i| {
                let var = self.q[(i, i)];
                -0.5 * (e[i] * e[i] / var + (2.0 * std::f64::consts::PI * var).ln())
            })
            .sum()
    }

    /// `sample_landmark_model_known_correspondence` — Table 6.5.
    ///
    /// Two constraints, three degrees of freedom: γ̂ is the one the reading
    /// cannot see, so it is drawn uniformly. The result is a *proposal*, not a
    /// belief — the pose prior it assumes (uniform over the plane) is wrong in
    /// every building ever built.
    pub fn sample_pose<R: Rng + ?Sized>(&self, f: &Feature, j: &Landmark, rng: &mut R) -> SE2 {
        let gamma = Uniform::new(0.0, std::f64::consts::TAU).unwrap().sample(rng);
        let r_hat = (f.r + Normal::new(0.0, self.q[(0, 0)].sqrt()).unwrap().sample(rng)).max(0.0);
        let phi_hat = f.phi + Normal::new(0.0, self.q[(1, 1)].sqrt()).unwrap().sample(rng);

        SE2::new(
            j.xy.x + r_hat * gamma.cos(),
            j.xy.y + r_hat * gamma.sin(),
            wrap_angle(gamma - std::f64::consts::PI - phi_hat),
        )
    }
}

use crate::{Scan, SensorModel};
use pr_core::geom::se2::SE2;
use sim::World;

/// A sensor model that has been told to be less sure of itself.
///
/// `kappa` divides the log-likelihood — the p^{1/κ} of Thrun §6.3.4 and the
/// scaled loss of generalized Bayesian updating. `stride` throws beams away
/// instead. They are not interchangeable: tempering keeps the peak where the
/// full scan puts it, subsampling moves it, so prefer κ and use `stride` only
/// when the cost of the beams themselves is the problem.
pub struct Tempered<S: SensorModel> {
    pub inner: S,
    pub kappa: f64,
    pub stride: usize,
}

impl<S: SensorModel> SensorModel for Tempered<S> {
    fn log_likelihood(&self, scan: &Scan, x: &SE2, map: &World) -> f64 {
        let stride = self.stride.max(1);
        let thinned = if stride == 1 {
            scan.clone()
        } else {
            Scan {
                ranges: scan.ranges.iter().step_by(stride).copied().collect(),
                bearings: scan.bearings.iter().step_by(stride).copied().collect(),
                max_range: scan.max_range,
            }
        };
        self.inner.log_likelihood(&thinned, x, map) / self.kappa.max(1e-6)
    }

    fn is_cheap(&self) -> bool {
        self.inner.is_cheap()
    }
}