//! The per-frame Gaussian list: footprint-sized anisotropic surfel splats
//! (see PLAN.md, "Core design").
//!
//! A surfel is a flat disk aligned to the depth-map normal:
//!
//! * two tangential scales = `depth * pixel_size / focal_length`
//!   (pixel size 1 px, focal in pixels ⇒ `depth / fx`),
//! * a few-mm thickness along the normal,
//! * opacity ~0.9–0.99,
//! * rendered with front-to-back alpha blending (Brush's renderer).

use glam::{Mat3, Quat, Vec3};

use crate::camera::CameraIntrinsics;

/// Thickness along the normal for surfels taken from depth data, in
/// metres ("a few mm", PLAN.md).
pub const SURFEL_THICKNESS_M: f32 = 0.003;

/// One Gaussian splat, in the raw (non-log, non-logit) parameterization
/// the pipeline reasons about. Brush's log/logit/SH conversions happen at
/// the render boundary, not here.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Splat {
    /// Centre position in world space, metres.
    pub mean: Vec3,
    /// Orientation: maps local (X, Y, Z) to the splat's axes. For surfels
    /// local Z is the surface normal. Stored as a glam quaternion (x, y,
    /// z, w) — the layout Brush expects.
    pub rotation: Quat,
    /// Per-axis scale in metres (not log scale).
    pub scale: Vec3,
    /// Colour, RGB in `[0, 1]`.
    pub color: [f32; 3],
    /// Opacity in `[0, 1]` (not logit).
    pub opacity: f32,
}

impl Splat {
    /// A footprint-sized surfel from a depth-map sample.
    ///
    /// * `position` — the back-projected depth pixel's 3D position,
    /// * `normal` — the surface normal at that position (unit length),
    /// * `depth` — the pixel's depth (distance along the camera's viewing
    ///   axis), used for the footprint,
    /// * `intrinsics` — the depth camera's intrinsics, for the
    ///   `depth / focal` footprint from PLAN.md,
    /// * `thickness` — extent along the normal in metres.
    pub fn surfel(
        position: Vec3,
        normal: Vec3,
        depth: f32,
        intrinsics: &CameraIntrinsics,
        thickness: f32,
        color: [f32; 3],
        opacity: f32,
    ) -> Self {
        let normal = normal.normalize_or_zero();
        // Footprint: one pixel's world size at this depth.
        let footprint_x = depth / intrinsics.fx;
        let footprint_y = depth / intrinsics.fy;
        let rotation = surfel_rotation(normal);
        Self {
            mean: position,
            rotation,
            scale: Vec3::new(footprint_x, footprint_y, thickness),
            color,
            opacity: opacity.clamp(0.0, 1.0),
        }
    }

    /// The splat's normal axis in world space (its local +Z).
    pub fn normal(&self) -> Vec3 {
        self.rotation * Vec3::Z
    }
}

/// A flat disk oriented so local Z is `normal`: local (X, Y) span the
/// tangent plane, right-handed.
fn surfel_rotation(normal: Vec3) -> Quat {
    // Reference axis not parallel to the normal (world up, or world X
    // for near-vertical normals).
    let reference = if normal.abs().dot(Vec3::Y) > 0.999 {
        Vec3::X
    } else {
        Vec3::Y
    };
    let t1 = normal.cross(reference).normalize_or_zero();
    let t2 = normal.cross(t1).normalize_or_zero();
    Quat::from_mat3(&Mat3::from_cols(t1, t2, normal))
}

/// A frame's Gaussian list. Constructed fresh every frame — there is no
/// training (PLAN.md, "Core design") — so this is a plain growable
/// collection.
#[derive(Debug, Clone, Default, PartialEq)]
pub struct SplatBatch {
    splats: Vec<Splat>,
}

impl SplatBatch {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            splats: Vec::with_capacity(capacity),
        }
    }

    pub fn len(&self) -> usize {
        self.splats.len()
    }

    pub fn is_empty(&self) -> bool {
        self.splats.is_empty()
    }

    pub fn push(&mut self, splat: Splat) {
        self.splats.push(splat);
    }

    pub fn clear(&mut self) {
        self.splats.clear();
    }

    pub fn iter(&self) -> impl Iterator<Item = &Splat> {
        self.splats.iter()
    }

    /// The centroid of the splat means (world space). Useful for aiming
    /// debug views at the subject.
    pub fn centroid(&self) -> Option<Vec3> {
        if self.splats.is_empty() {
            return None;
        }
        let sum = self.splats.iter().fold(Vec3::ZERO, |acc, s| acc + s.mean);
        Some(sum / self.splats.len() as f32)
    }
}

impl<'a> IntoIterator for &'a SplatBatch {
    type Item = &'a Splat;
    type IntoIter = std::slice::Iter<'a, Splat>;

    fn into_iter(self) -> Self::IntoIter {
        self.splats.iter()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use glam::vec3;

    fn intrinsics() -> CameraIntrinsics {
        CameraIntrinsics {
            fx: 500.0,
            fy: 500.0,
            cx: 256.0,
            cy: 212.0,
        }
    }

    #[test]
    fn surfel_footprint_is_depth_over_focal() {
        // PLAN.md: two tangential scales = depth * pixel_size / focal.
        let s = Splat::surfel(
            vec3(0.0, 0.0, 2.0),
            vec3(0.0, 0.0, -1.0),
            2.0,
            &intrinsics(),
            SURFEL_THICKNESS_M,
            [1.0, 0.0, 0.0],
            0.95,
        );
        assert!((s.scale.x - 2.0 / 500.0).abs() < 1e-6);
        assert!((s.scale.y - 2.0 / 500.0).abs() < 1e-6);
        // Thickness along the normal, not the footprint.
        assert!((s.scale.z - SURFEL_THICKNESS_M).abs() < 1e-6);
        assert!((s.normal() - vec3(0.0, 0.0, -1.0)).length() < 1e-5);
        assert_eq!(s.opacity, 0.95);
    }

    #[test]
    fn surfel_rotation_maps_local_z_to_normal() {
        for normal in [
            vec3(0.0, 0.0, -1.0),
            vec3(0.0, 1.0, 0.0),
            vec3(1.0, 0.0, 0.0),
            vec3(0.707, 0.707, 0.0),
            vec3(0.577, 0.577, 0.577),
            vec3(0.0, -1.0, 0.0), // near the world-up fallback axis
        ] {
            let n = normal.normalize();
            let q = surfel_rotation(n);
            // Local Z maps to the normal; the quaternion is a proper
            // rotation (orthonormal basis).
            assert!((q * Vec3::Z - n).length() < 1e-5, "normal {n:?}");
            assert!((q.length() - 1.0).abs() < 1e-5);
        }
    }

    #[test]
    fn surfel_opacity_is_clamped() {
        let s = Splat::surfel(
            Vec3::ZERO,
            vec3(0.0, 0.0, -1.0),
            1.0,
            &intrinsics(),
            SURFEL_THICKNESS_M,
            [0.0; 3],
            1.4,
        );
        assert_eq!(s.opacity, 1.0);
    }

    #[test]
    fn batch_basics() {
        let mut batch = SplatBatch::with_capacity(4);
        assert!(batch.is_empty());
        assert_eq!(batch.centroid(), None);
        batch.push(Splat::surfel(
            vec3(1.0, 0.0, 0.0),
            vec3(0.0, 0.0, -1.0),
            2.0,
            &intrinsics(),
            SURFEL_THICKNESS_M,
            [1.0; 3],
            0.9,
        ));
        batch.push(Splat::surfel(
            vec3(0.0, 3.0, 0.0),
            vec3(0.0, 0.0, -1.0),
            2.0,
            &intrinsics(),
            SURFEL_THICKNESS_M,
            [1.0; 3],
            0.9,
        ));
        assert_eq!(batch.len(), 2);
        assert_eq!(batch.centroid(), Some(vec3(0.5, 1.5, 0.0)));
        batch.clear();
        assert!(batch.is_empty());
    }
}