//! Depth preprocessing and splat construction (PLAN.md, pipeline
//! steps 2–3).
//!
//! [`SplatGen`] turns one raw [`DepthFrame`] — a millimetre depth map
//! plus RGB, exactly what a Kinect (or the synthetic source, or a
//! recording) emits — into the frame's [`SplatBatch`]:
//!
//! 1. **cull** — drop invalid pixels (depth 0) and everything beyond
//!    `max_depth_mm` (background culling),
//! 2. **flier removal** — drop pixels that disagree hard with their
//!    3×3 neighbourhood median (isolated wrong returns),
//! 3. **bilateral filter** — edge-preserving smoothing over a small
//!    pixel radius, so depth edges survive but sensor jitter does not,
//! 4. **normals** — from screen-space derivatives of the filtered
//!    point cloud, oriented towards the camera,
//! 5. **surfels** — one footprint-sized anisotropic splat per
//!    surviving pixel via `rig::Splat::surfel` (the plan's core
//!    construction), coloured from the RGB pixel.
//!
//! [`SplatSource`] adapts any [`DepthSource`] into a
//! [`FrameSource`]: capture, recordings, and the synthetic dataset all
//! plug into the renderer the same way.

use anyhow::Result;
use glam::Vec3;

use rig::{Camera, DepthFrame, DepthSource, Frame, FrameSource, Splat, SplatBatch};

/// Tuning knobs for the depth pipeline; the defaults follow PLAN.md
/// (bilateral radius 3–5 px, opacity ~0.95, few-mm thickness).
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct SplatGenParams {
    /// Bilateral filter radius in pixels.
    pub filter_radius: u32,
    /// Bilateral depth-domain sigma, millimetres.
    pub filter_sigma_depth_mm: f32,
    /// Bilateral spatial-domain sigma, pixels.
    pub filter_sigma_space_px: f32,
    /// A pixel disagreeing with its 3×3 neighbourhood median by more
    /// than this is a flier and is dropped, millimetres.
    pub flier_threshold_mm: u32,
    /// Pixels beyond this depth are culled as background, millimetres.
    pub max_depth_mm: u16,
    /// Surfet opacity (PLAN.md: ~0.9–0.99).
    pub opacity: f32,
    /// Surfel thickness along the normal, metres.
    pub thickness_m: f32,
}

impl Default for SplatGenParams {
    fn default() -> Self {
        Self {
            filter_radius: 3,
            filter_sigma_depth_mm: 25.0,
            filter_sigma_space_px: 2.0,
            flier_threshold_mm: 250,
            max_depth_mm: 4000,
            opacity: 0.95,
            thickness_m: rig::SURFEL_THICKNESS_M,
        }
    }
}

/// The depth pipeline: depth frames in, splat batches out.
pub struct SplatGen {
    params: SplatGenParams,
}

impl Default for SplatGen {
    fn default() -> Self {
        Self::new(SplatGenParams::default())
    }
}

impl SplatGen {
    pub fn new(params: SplatGenParams) -> Self {
        Self { params }
    }

    /// The tuning parameters in effect.
    pub fn params(&self) -> &SplatGenParams {
        &self.params
    }

    /// Preprocess one depth frame and assemble its splat batch.
    pub fn generate(&self, frame: &DepthFrame, camera: &Camera) -> SplatBatch {
        let (width, height) = (camera.resolution.x, camera.resolution.y);
        let pixel_count = (width * height) as usize;
        assert_eq!(frame.depth_mm.len(), pixel_count, "depth map size");
        assert_eq!(frame.rgb.len(), pixel_count * 3, "rgb size");

        // f32 working copy; `None` = invalid.
        let mut depth: Vec<Option<f32>> = frame
            .depth_mm
            .iter()
            .map(|&d| (d > 0 && d <= self.params.max_depth_mm).then_some(d as f32))
            .collect();

        remove_fliers(&mut depth, width, self.params.flier_threshold_mm);
        let filtered = bilateral_filter(&depth, width, self.params);

        let mut batch = SplatBatch::with_capacity(pixel_count);
        for y in 0..height {
            for x in 0..width {
                let Some(z_mm) = filtered[(y * width + x) as usize] else {
                    continue;
                };
                let (u, v) = (x as f32 + 0.5, y as f32 + 0.5);
                let z = z_mm / 1000.0;
                let position = camera.backproject(u, v, z);
                let normal = estimate_normal(&filtered, width, camera, x, y);
                let idx = (y * width + x) as usize;
                let color = [
                    frame.rgb[idx * 3] as f32 / 255.0,
                    frame.rgb[idx * 3 + 1] as f32 / 255.0,
                    frame.rgb[idx * 3 + 2] as f32 / 255.0,
                ];
                batch.push(Splat::surfel(
                    position,
                    normal,
                    z,
                    &camera.intrinsics,
                    self.params.thickness_m,
                    color,
                    self.params.opacity,
                ));
            }
        }
        batch
    }
}

/// A [`FrameSource`] adapter: pull raw depth frames from `S` and run
/// them through [`SplatGen`]. This is the pipeline seam — live capture,
/// recordings, and the synthetic dataset all become renderable sources.
pub struct SplatSource<S: DepthSource> {
    inner: S,
    generator: SplatGen,
}

impl<S: DepthSource> SplatSource<S> {
    pub fn new(inner: S, generator: SplatGen) -> Self {
        Self { inner, generator }
    }

    /// The wrapped depth source.
    pub fn inner(&self) -> &S {
        &self.inner
    }
}

impl<S: DepthSource> FrameSource for SplatSource<S> {
    fn next_frame(&mut self) -> Result<Frame> {
        let depth_frame = self.inner.next_depth_frame()?;
        let splats = self.generator.generate(&depth_frame, self.inner.camera());
        Ok(Frame {
            timestamp_secs: depth_frame.timestamp_secs,
            splats,
        })
    }
}

/// Drop pixels that disagree with their 3×3 neighbourhood median by
/// more than `threshold_mm` (isolated fliers). Neighbourhood positions
/// are clamped at the borders.
fn remove_fliers(depth: &mut [Option<f32>], width: u32, threshold_mm: u32) {
    let width = width as i32;
    let height = (depth.len() / width as usize) as i32;
    let threshold = threshold_mm as f32;
    let mut neighbourhood = [0.0_f32; 9];

    let mut flier = vec![false; depth.len()];
    for y in 0..height {
        for x in 0..width {
            let idx = (y * width + x) as usize;
            let Some(own) = depth[idx] else { continue };
            let mut n = 0;
            for dy in -1..=1 {
                for dx in -1..=1 {
                    let (nx, ny) = (x + dx, y + dy);
                    if !(0..width).contains(&nx) || !(0..height).contains(&ny) {
                        continue;
                    }
                    if let Some(d) = depth[(ny * width + nx) as usize] {
                        neighbourhood[n] = d;
                        n += 1;
                    }
                }
            }
            if n < 3 {
                continue; // not enough context to judge
            }
            neighbourhood[..n].sort_by(|a, b| a.partial_cmp(b).unwrap());
            let median = if n % 2 == 1 {
                neighbourhood[n / 2]
            } else {
                (neighbourhood[n / 2 - 1] + neighbourhood[n / 2]) * 0.5
            };
            if (own - median).abs() > threshold {
                flier[idx] = true;
            }
        }
    }
    for (depth, flier) in depth.iter_mut().zip(flier) {
        if flier {
            *depth = None;
        }
    }
}

/// Edge-preserving bilateral filter over valid pixels: spatial × depth
/// Gaussians, radius from the params.
fn bilateral_filter(depth: &[Option<f32>], width: u32, params: SplatGenParams) -> Vec<Option<f32>> {
    let width_i = width as i32;
    let height = (depth.len() / width as usize) as i32;
    let radius = params.filter_radius as i32;
    let inv_two_sigma_d2 = 1.0 / (2.0 * params.filter_sigma_depth_mm.powi(2));
    let inv_two_sigma_s2 = 1.0 / (2.0 * params.filter_sigma_space_px.powi(2));

    let mut out = depth.to_vec();
    for y in 0..height {
        for x in 0..width_i {
            let idx = (y * width_i + x) as usize;
            let Some(own) = depth[idx] else { continue };
            let mut total = 0.0;
            let mut weight_sum = 0.0;
            for dy in -radius..=radius {
                for dx in -radius..=radius {
                    let (nx, ny) = (x + dx, y + dy);
                    if !(0..width_i).contains(&nx) || !(0..height).contains(&ny) {
                        continue;
                    }
                    let Some(neighbour) = depth[(ny * width_i + nx) as usize] else {
                        continue;
                    };
                    let spatial = (dx * dx + dy * dy) as f32;
                    let dd = neighbour - own;
                    let weight = (-spatial * inv_two_sigma_s2 - dd * dd * inv_two_sigma_d2).exp();
                    total += weight * neighbour;
                    weight_sum += weight;
                }
            }
            if weight_sum > 1e-6 {
                out[idx] = Some(total / weight_sum);
            }
        }
    }
    out
}

/// Surface normal at a pixel from the filtered depth: the cross product
/// of screen-space derivatives of the backprojected point cloud,
/// oriented towards the camera. Falls back to facing the camera when no
/// valid neighbours exist.
fn estimate_normal(depth: &[Option<f32>], width: u32, camera: &Camera, x: u32, y: u32) -> Vec3 {
    let width_i = width as i32;
    let height = (depth.len() / width as usize) as i32;
    let (xi, yi) = (x as i32, y as i32);

    let point = |px: i32, py: i32| -> Option<Vec3> {
        if !(0..width_i).contains(&px) || !(0..height).contains(&py) {
            return None;
        }
        let d = depth[(py * width_i + px) as usize]? / 1000.0;
        Some(camera.backproject(px as f32 + 0.5, py as f32 + 0.5, d))
    };

    let center = point(xi, yi).expect("caller checked the pixel is valid");
    // Central differences where both neighbours are valid, otherwise
    // one-sided, otherwise no estimate.
    let du = point(xi + 1, yi)
        .zip(point(xi - 1, yi))
        .map(|(a, b)| a - b)
        .or_else(|| point(xi + 1, yi).map(|a| (a - center) * 2.0))
        .or_else(|| point(xi - 1, yi).map(|b| (center - b) * 2.0));
    let dv = point(xi, yi + 1)
        .zip(point(xi, yi - 1))
        .map(|(a, b)| a - b)
        .or_else(|| point(xi, yi + 1).map(|a| (a - center) * 2.0))
        .or_else(|| point(xi, yi - 1).map(|b| (center - b) * 2.0));

    let toward_camera = (camera.pose.position - center).normalize_or_zero();
    let Some(du) = du else { return toward_camera };
    let Some(dv) = dv else { return toward_camera };
    let mut normal = du.cross(dv).normalize_or_zero();
    if normal.length_squared() < 1e-8 {
        // Degenerate (no valid neighbours): face the camera.
        normal = toward_camera;
    } else if normal.dot(toward_camera) < 0.0 {
        // The cross product's handedness is arbitrary; orient towards
        // the camera.
        normal = -normal;
    }
    normal
}

#[cfg(test)]
mod tests {
    use super::*;
    use glam::{uvec2, vec3};
    use rig::{CameraIntrinsics, CameraPose};

    fn camera(width: u32, height: u32) -> Camera {
        Camera {
            id: "test".into(),
            pose: CameraPose::looking_at(Vec3::ZERO, vec3(0.0, 0.0, 2.0)),
            intrinsics: CameraIntrinsics {
                fx: 100.0,
                fy: 100.0,
                cx: width as f32 / 2.0,
                cy: height as f32 / 2.0,
            },
            resolution: uvec2(width, height),
        }
    }

    fn frame(width: u32, height: u32, depth_of: impl Fn(u32, u32) -> u16) -> DepthFrame {
        let mut depth = Vec::with_capacity((width * height) as usize);
        let mut rgb = Vec::with_capacity((width * height * 3) as usize);
        for y in 0..height {
            for x in 0..width {
                depth.push(depth_of(x, y));
                rgb.extend_from_slice(&[200, 100, 50]);
            }
        }
        DepthFrame {
            timestamp_secs: 0.0,
            depth_mm: depth,
            rgb,
        }
    }

    #[test]
    fn flat_wall_makes_camera_facing_surfels() {
        let (w, h) = (32u32, 24u32);
        let camera = camera(w, h);
        let frame = frame(w, h, |_, _| 2000);
        let batch = SplatGen::default().generate(&frame, &camera);

        assert_eq!(batch.len(), (w * h) as usize);
        for splat in batch.iter() {
            // On the wall plane: the projection lands back on the pixel.
            let projected = camera.project(splat.mean).expect("in front of the camera");
            // The wall's normal is constant (0, 0, -1) in world space…
            let wall_normal = Vec3::new(0.0, 0.0, -1.0);
            assert!(
                (splat.normal() - wall_normal).length() < 0.05,
                "wall normal {:?}",
                splat.normal()
            );
            // …and faces the camera everywhere on the wall.
            let to_camera = (camera.pose.position - splat.mean).normalize();
            assert!(splat.normal().dot(to_camera) > 0.9);
            // Footprint from the plan: depth / focal.
            let z = 2.0;
            assert!((splat.scale.x - z / 100.0).abs() < 1e-4);
            assert!((splat.scale.y - z / 100.0).abs() < 1e-4);
            assert_eq!(splat.color, [200.0 / 255.0, 100.0 / 255.0, 50.0 / 255.0]);
            assert!((splat.opacity - 0.95).abs() < 1e-6);
            let _ = projected;
        }
    }

    #[test]
    fn culling_drops_invalid_and_far_pixels() {
        let (w, h) = (16u32, 8u32);
        let camera = camera(w, h);
        let frame = frame(w, h, |x, y| match (x, y) {
            (0, 0) => 0,    // no return
            (1, 0) => 5000, // beyond max depth
            _ => 2000,
        });
        let batch = SplatGen::default().generate(&frame, &camera);
        assert_eq!(batch.len(), (w * h - 2) as usize);
    }

    #[test]
    fn fliers_are_removed() {
        let (w, h) = (16u32, 16u32);
        let camera = camera(w, h);
        // A few isolated wrong returns among a flat wall.
        let frame = frame(w, h, |x, y| {
            if matches!((x, y), (4, 4) | (10, 12) | (15, 0)) {
                600 // wildly wrong
            } else {
                2000
            }
        });
        let batch = SplatGen::default().generate(&frame, &camera);
        assert_eq!(
            batch.len(),
            (w * h - 3) as usize,
            "isolated fliers must be culled"
        );
        // And none of the surviving splats sit at the flier depth.
        for splat in batch.iter() {
            let z = camera
                .pose
                .camera_from_world()
                .transform_point3(splat.mean)
                .z;
            assert!((z - 2.0).abs() < 0.05, "stray depth {z}");
        }
    }

    #[test]
    fn bilateral_smooths_jitter_but_not_edges() {
        let (w, h) = (32u32, 8u32);
        let camera = camera(w, h);
        // Left half at 2000 mm, right half at 3000 mm, everything
        // jittered ±3 mm.
        let clean = frame(w, h, |x, _| if x < 16 { 2000 } else { 3000 });
        let noisy = frame(w, h, |x, y| {
            let base = if x < 16 { 2000 } else { 3000 };
            let jitter = ((x * 7 + y * 13) % 7) as i16 - 3;
            ((base as i16) + jitter) as u16
        });

        let generator = SplatGen::default();
        let clean_batch = generator.generate(&clean, &camera);
        let noisy_batch = generator.generate(&noisy, &camera);
        assert_eq!(clean_batch.len(), noisy_batch.len());

        // The noisy splats sit within a fraction of a pixel of the
        // clean ones: the jitter (≈3 mm at 2 m ≈ 0.15 px) is smoothed
        // away, and the step edge is preserved.
        for (a, b) in clean_batch.iter().zip(noisy_batch.iter()) {
            assert!(
                (a.mean - b.mean).length() < 0.002,
                "positions differ: {:?} vs {:?}",
                a.mean,
                b.mean
            );
        }
    }

    #[test]
    fn normals_follow_a_slanted_plane() {
        let (w, h) = (24u32, 24u32);
        let camera = camera(w, h);
        // Depth increases by 40 mm per pixel: the surface tilts away to
        // the right. dz/du = 0.04 m per pixel; world dx per pixel is
        // z/fx, so the plane's normal is ∝ (-slope, 0, 1).
        let frame = frame(w, h, |x, _| 2000 + 40 * x as u16);

        let batch = SplatGen::default().generate(&frame, &camera);
        // Check a pixel near the centre, away from border effects.
        let splat = batch.iter().nth(12 * 24 + 12).expect("center pixel");

        let z_world = 2.0 + 0.04 * 12.5; // metres at that pixel
        // The plane's camera-facing normal in world space: image-right
        // maps to world -X, so rising depth with u means dz/dx_world is
        // negative there.
        let dz_du = 0.04_f32; // metres per pixel
        let dz_dx_world = dz_du / (-(z_world / camera.intrinsics.fx));
        let expected = Vec3::new(dz_dx_world, 0.0, -1.0).normalize();
        let dot = splat.normal().dot(expected);
        assert!(
            dot > 0.995,
            "normal {:?} vs expected {expected:?} (dot {dot})",
            splat.normal()
        );
    }

    #[test]
    fn params_are_honoured() {
        let (w, h) = (16u32, 16u32);
        let camera = camera(w, h);
        let frame = frame(w, h, |_, _| 2500);
        let params = SplatGenParams {
            max_depth_mm: 2000, // culls everything
            ..SplatGenParams::default()
        };
        let batch = SplatGen::new(params).generate(&frame, &camera);
        assert!(batch.is_empty());

        // Opacity flows through to the surfels.
        let params = SplatGenParams {
            opacity: 0.9,
            ..SplatGenParams::default()
        };
        let batch = SplatGen::new(params).generate(&frame, &camera);
        assert!((batch.iter().next().unwrap().opacity - 0.9).abs() < 1e-6);
    }

    #[test]
    fn generates_from_the_synthetic_source() {
        // Integration: the synthetic Kinect feeds the pipeline.
        use rig::DepthSource;
        let rig = replay::synth::dev_rig_with_resolutions(uvec2(48, 40), uvec2(64, 48));
        let mut source =
            replay::synth::dev_source_with(rig.center.clone(), replay::synth::SUBJECT, 30.0);
        let frame = source.next_depth_frame().expect("depth frame");
        let batch = SplatGen::default().generate(&frame, &rig.center);

        // Roughly one splat per depth pixel (fliers removed).
        assert!(batch.len() > 48 * 40 * 95 / 100);
        assert!(batch.len() <= 48 * 40);
        // All in front of the camera, all at valid depths.
        let camera_from_world = rig.center.pose.camera_from_world();
        for splat in batch.iter() {
            let p = camera_from_world.transform_point3(splat.mean);
            assert!(p.z > 0.5 && p.z < 3.1, "depth {p:?}");
        }
    }
}