//! Synthetic scenes with exact, analytic geometry.
//!
//! Surfaces are simple enough to intersect exactly (ellipsoids, planes)
//! yet shaped like the real subject — a head, shoulders, and a backdrop —
//! so surfel construction, the virtual camera, the occlusion proxy, and
//! the UI can be developed and tested against known ground truth.

use glam::Vec3;

use rig::Ellipsoid;

/// A ray-surface intersection.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Hit {
    /// Distance along the ray.
    pub t: f32,
    /// World-space hit position.
    pub position: Vec3,
    /// Unit surface normal at the hit.
    pub normal: Vec3,
    /// Surface colour at the hit, RGB in `[0, 1]`.
    pub color: [f32; 3],
}

/// A scene surface.
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Surface {
    /// An axis-aligned ellipsoid (sphere with per-axis radii) in a
    /// single solid colour.
    Ellipsoid {
        center: Vec3,
        radii: Vec3,
        color: [f32; 3],
    },
    /// An ellipsoid whose left and right flanks carry a distinct tint:
    /// beyond `tint_from` fractions of the x radius from the centre,
    /// the colour switches to the side tint. Those flanks are what the
    /// side webcams see but the center camera barely does — the ground
    /// truth for the ellipsoid-proxy recolouring.
    TintedEllipsoid {
        center: Vec3,
        radii: Vec3,
        base: [f32; 3],
        left_tint: [f32; 3],
        right_tint: [f32; 3],
        /// Flank tinting starts beyond this fraction of the x radius,
        /// in `[0, 1]`.
        tint_from: f32,
    },
    /// A backdrop plane with a checkerboard pattern, `cell` metres per
    /// cell (0 disables the pattern: flat colour).
    CheckerPlane {
        point: Vec3,
        normal: Vec3,
        cell: f32,
        color_a: [f32; 3],
        color_b: [f32; 3],
    },
}

impl Surface {
    /// Intersect `origin + t * dir` with the surface; the nearest hit
    /// with `t > 0`.
    ///
    /// `dir` need not be normalized (as long as the caller compares
    /// distances consistently, which [`Scene::cast_ray`] does).
    pub fn intersect(&self, origin: Vec3, dir: Vec3) -> Option<Hit> {
        match *self {
            Surface::Ellipsoid {
                center,
                radii,
                color,
            } => {
                let ellipsoid = Ellipsoid::new(center, radii);
                let (t, normal) = ellipsoid.intersect_ray(origin, dir)?;
                Some(Hit {
                    t,
                    position: origin + dir * t,
                    normal,
                    color,
                })
            }
            Surface::TintedEllipsoid {
                center,
                radii,
                base,
                left_tint,
                right_tint,
                tint_from,
            } => {
                let ellipsoid = Ellipsoid::new(center, radii);
                let (t, normal) = ellipsoid.intersect_ray(origin, dir)?;
                let position = origin + dir * t;
                // Tint by the hit's offset along the local x axis.
                let dx = (position.x - center.x) / radii.x;
                let color = if dx < -tint_from {
                    left_tint
                } else if dx > tint_from {
                    right_tint
                } else {
                    base
                };
                Some(Hit {
                    t,
                    position,
                    normal,
                    color,
                })
            }
            Surface::CheckerPlane {
                point,
                normal,
                cell,
                color_a,
                color_b,
            } => {
                let normal = normal.normalize_or_zero();
                let denom = dir.dot(normal);
                if denom.abs() <= f32::EPSILON {
                    return None; // parallel
                }
                let t = (point - origin).dot(normal) / denom;
                if t <= f32::EPSILON {
                    return None;
                }
                let position = origin + dir * t;
                let color = if cell > 0.0 {
                    // Checkerboard over the plane's in-plane coordinates.
                    let (u, v) = plane_coords(position, normal);
                    let parity = ((u / cell).floor() + (v / cell).floor()) as i32 % 2;
                    if parity.rem_euclid(2) == 0 {
                        color_a
                    } else {
                        color_b
                    }
                } else {
                    color_a
                };
                Some(Hit {
                    t,
                    position,
                    normal,
                    color,
                })
            }
        }
    }

    /// The surface translated by `offset`.
    pub fn translated(&self, offset: Vec3) -> Self {
        match *self {
            Surface::Ellipsoid {
                center,
                radii,
                color,
            } => Surface::Ellipsoid {
                center: center + offset,
                radii,
                color,
            },
            Surface::TintedEllipsoid {
                center,
                radii,
                base,
                left_tint,
                right_tint,
                tint_from,
            } => Surface::TintedEllipsoid {
                center: center + offset,
                radii,
                base,
                left_tint,
                right_tint,
                tint_from,
            },
            Surface::CheckerPlane {
                point,
                normal,
                cell,
                color_a,
                color_b,
            } => Surface::CheckerPlane {
                point: point + offset,
                normal,
                cell,
                color_a,
                color_b,
            },
        }
    }
}

/// In-plane coordinates of a point on a plane with the given normal,
/// for the checkerboard pattern.
fn plane_coords(p: Vec3, normal: Vec3) -> (f32, f32) {
    // Build any orthonormal tangent basis.
    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();
    (p.dot(t1), p.dot(t2))
}

/// A synthetic scene: a list of surfaces, cast against per pixel.
#[derive(Debug, Clone, Default, PartialEq)]
pub struct Scene {
    pub surfaces: Vec<Surface>,
}

/// The flanks of the synthetic head (the region tinted for the side
/// cameras) start beyond this fraction of the head's x radius.
pub const HEAD_TINT_FROM: f32 = 0.6;

/// The synthetic subject's head colour and its side tints.
pub const HEAD_BASE: [f32; 3] = [0.85, 0.62, 0.52];
pub const HEAD_LEFT_TINT: [f32; 3] = [0.85, 0.05, 0.05];
pub const HEAD_RIGHT_TINT: [f32; 3] = [0.05, 0.05, 0.85];

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

    /// Add a surface.
    pub fn with(mut self, surface: Surface) -> Self {
        self.surfaces.push(surface);
        self
    }

    /// The default development scene's subject: a head-and-shoulders
    /// stand-in at `subject`, facing the rig's center camera at the
    /// origin. The head's flanks carry the side tints (see
    /// [`Surface::TintedEllipsoid`]).
    ///
    /// Returns `(head, shoulders)`; the backdrop is [`backdrop`].
    pub fn head_and_shoulders(subject: Vec3) -> (Surface, Surface) {
        (
            // Head: slightly egg-shaped ellipsoid, ~22 cm tall, with
            // tinted flanks for the proxy ground truth.
            Surface::TintedEllipsoid {
                center: subject + Vec3::new(0.0, 0.18, 0.0),
                radii: Vec3::new(0.09, 0.11, 0.10),
                base: HEAD_BASE,
                left_tint: HEAD_LEFT_TINT,
                right_tint: HEAD_RIGHT_TINT,
                tint_from: HEAD_TINT_FROM,
            },
            // Shoulders.
            Surface::Ellipsoid {
                center: subject + Vec3::new(0.0, -0.05, 0.0),
                radii: Vec3::new(0.24, 0.13, 0.16),
                color: [0.20, 0.35, 0.55],
            },
        )
    }

    /// The default development scene's backdrop: a checkerboard plane
    /// `0.8 m` behind the subject, facing the rig.
    pub fn backdrop(subject: Vec3) -> Surface {
        Surface::CheckerPlane {
            point: subject + Vec3::new(0.0, 0.0, 0.8),
            normal: Vec3::new(0.0, 0.0, -1.0),
            cell: 0.25,
            color_a: [0.72, 0.72, 0.72],
            color_b: [0.28, 0.28, 0.28],
        }
    }

    /// The default development scene: head, shoulders, backdrop, all
    /// centred at `subject`. See [`head_and_shoulders`] for the
    /// animated variant used by the sources.
    pub fn head_stand_in(subject: Vec3) -> Self {
        let (head, shoulders) = Self::head_and_shoulders(subject);
        Self::new()
            .with(head)
            .with(shoulders)
            .with(Self::backdrop(subject))
    }

    /// Cast a ray against the scene; the nearest surface hit.
    pub fn cast_ray(&self, origin: Vec3, dir: Vec3) -> Option<Hit> {
        self.surfaces
            .iter()
            .filter_map(|s| s.intersect(origin, dir))
            .min_by(|a, b| a.t.partial_cmp(&b.t).unwrap())
    }
}

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

    fn scene() -> Scene {
        Scene::head_stand_in(vec3(0.0, 0.0, 2.0))
    }

    #[test]
    fn center_ray_hits_the_head() {
        // Straight at the head's centre: (0, 0.18, 2.0).
        let head_center = vec3(0.0, 0.18, 2.0);
        let hit = scene()
            .cast_ray(Vec3::ZERO, head_center)
            .expect("head in the way");
        // The hit is on the ellipsoid's near side, on its surface.
        assert!(hit.position.distance(head_center) < 0.11);
        assert!(hit.position.z < head_center.z, "must hit the near side");
        // Normal faces the camera (front face of the ellipsoid).
        assert!(hit.normal.dot(hit.position - Vec3::ZERO) < 0.0);
    }

    #[test]
    fn head_flanks_are_tinted() {
        // Rays aimed straight at the head's left and right flanks (beyond
        // HEAD_TINT_FROM of the x radius) hit tinted colour; the front
        // is the base colour.
        let head_center = vec3(0.0, 0.18, 2.0);
        let left = scene()
            .cast_ray(vec3(-3.0, 0.18, 2.0), vec3(1.0, 0.0, 0.0))
            .expect("left flank hit");
        assert_eq!(left.color, HEAD_LEFT_TINT);
        assert!(left.position.x < head_center.x - 0.09 * HEAD_TINT_FROM);
        let right = scene()
            .cast_ray(vec3(3.0, 0.18, 2.0), vec3(-1.0, 0.0, 0.0))
            .expect("right flank hit");
        assert_eq!(right.color, HEAD_RIGHT_TINT);
        assert!(right.position.x > head_center.x + 0.09 * HEAD_TINT_FROM);
        let front = scene()
            .cast_ray(Vec3::ZERO, head_center)
            .expect("front hit");
        assert_eq!(front.color, HEAD_BASE);
    }

    #[test]
    fn backdrop_behind_the_subject() {
        let hit = scene()
            .cast_ray(Vec3::ZERO, vec3(0.0, 0.35, 1.0))
            .expect("backdrop in the way");
        assert!((hit.position.z - 2.8).abs() < 1e-5);
        assert!((hit.normal - vec3(0.0, 0.0, -1.0)).length() < 1e-6);
        // Checkerboard is one of the two greys.
        let c = hit.color;
        assert!(
            (c[0] - 0.72).abs() < 1e-5 || (c[0] - 0.28).abs() < 1e-5,
            "checker color was {c:?}"
        );
    }

    #[test]
    fn nearest_hit_wins() {
        // A ray above the head hits only the backdrop.
        let hit = scene()
            .cast_ray(Vec3::ZERO, vec3(0.0, 0.5, 1.0))
            .expect("backdrop");
        assert!((hit.position.z - 2.8).abs() < 1e-5);
        // A ray through the head hits the head, not the backdrop behind it.
        let hit = scene()
            .cast_ray(Vec3::ZERO, vec3(0.0, 0.09, 1.0))
            .expect("head");
        assert!(hit.position.z < 2.0, "nearest hit must be the head");
    }

    #[test]
    fn misses_report_none() {
        // Only rays pointing away from the scene miss: everything forward
        // eventually reaches the (infinite) backdrop plane.
        assert!(scene().cast_ray(Vec3::ZERO, vec3(0.0, 0.0, -1.0)).is_none());
        // Straight up, parallel to the backdrop: miss.
        assert!(scene().cast_ray(Vec3::ZERO, vec3(0.0, 1.0, 0.0)).is_none());
    }

    #[test]
    fn ellipsoid_hit_math() {
        let s = Surface::Ellipsoid {
            center: vec3(0.0, 0.0, 2.0),
            radii: vec3(1.0, 1.0, 1.0),
            color: [1.0, 0.0, 0.0],
        };
        let hit = s.intersect(Vec3::ZERO, vec3(0.0, 0.0, 1.0)).expect("hits");
        assert!((hit.position.z - 1.0).abs() < 1e-5);
        assert!((hit.normal - vec3(0.0, 0.0, -1.0)).length() < 1e-5);
        assert!((hit.t - 1.0).abs() < 1e-5);
        // Ray from inside: front face ahead of the origin.
        let hit = s
            .intersect(vec3(0.0, 0.0, 2.0), vec3(0.0, 0.0, 1.0))
            .expect("exits");
        assert!((hit.t - 1.0).abs() < 1e-5);
    }

    #[test]
    fn checker_pattern_alternates() {
        let s = Surface::CheckerPlane {
            point: vec3(0.0, 0.0, 3.0),
            normal: vec3(0.0, 0.0, -1.0),
            cell: 1.0,
            color_a: [1.0, 1.0, 1.0],
            color_b: [0.0, 0.0, 0.0],
        };
        let a = s.intersect(Vec3::ZERO, vec3(0.0, 0.0, 1.0)).unwrap();
        let b = s.intersect(Vec3::ZERO, vec3(1.0, 0.0, 3.0)).unwrap();
        let c = s.intersect(Vec3::ZERO, vec3(-1.0, 0.0, 3.0)).unwrap();
        assert_eq!(a.color, [1.0, 1.0, 1.0]);
        assert_eq!(b.color, [0.0, 0.0, 0.0]);
        assert_eq!(c.color, [0.0, 0.0, 0.0]);
    }
}