//! The synthetic depth-camera: a software Kinect.
//!
//! Every frame, [`SyntheticSource`] raycasts the scene — one ray per
//! depth pixel — and emits a raw [`DepthFrame`]: a millimetre depth map
//! plus RGB from the center (Kinect) viewpoint, *and* RGB from the two
//! side webcams, exactly what the live rig produces. `splatgen` then
//! turns those frames into surfels and occlusion-proxy splats, so the
//! dev/CI path exercises the real depth-processing pipeline.
//!
//! Realism knobs: the source injects deterministic sensor noise
//! ([`NoiseParams`]) — isolated fliers and per-pixel jitter — so the
//! flier removal and bilateral filter have real work to do. The noise is
//! a pure function of `(frame, pixel)`, never the wall clock: playback
//! stays deterministic.
//!
//! The subject sways a little, also deterministically, so temporal
//! behaviour is testable. The head's flanks carry distinct tints (see
//! [`crate::scene`]), visible mainly to the side cameras — ground truth
//! for the ellipsoid-proxy recolouring.

use anyhow::Result;
use glam::{UVec2, Vec3};

use rig::{Camera, DepthFrame, DepthSource, Rig, SideRgb};

use crate::scene::{Scene, Surface};

/// Where the synthetic subject sits: 2 m in front of the rig's center
/// camera at the origin.
pub const SUBJECT: Vec3 = Vec3::new(0.0, 0.0, 2.0);

/// Subject sway amplitude, metres, and period in frames.
const SWAY_AMPLITUDE_M: f32 = 0.04;
const SWAY_PERIOD_FRAMES: u64 = 90;

/// The depth range a Kinect v2 reports; synthetic noise stays in it.
const KINECT_MIN_MM: u16 = 500;
const KINECT_MAX_MM: u16 = 4000;

/// Sensor noise injected into the synthetic depth maps.
///
/// Deterministic: a pure function of the frame index and pixel, so
/// playback is reproducible.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct NoiseParams {
    /// Fraction of valid pixels replaced with a random depth (fliers),
    /// in `[0, 1]`.
    pub flier_rate: f32,
    /// Per-pixel depth jitter, ± millimetres.
    pub jitter_mm: u32,
}

impl Default for NoiseParams {
    fn default() -> Self {
        // Mild but realistic: scattered fliers, a few mm of jitter.
        Self {
            flier_rate: 0.005,
            jitter_mm: 4,
        }
    }
}

impl NoiseParams {
    /// Clean frames: no fliers, no jitter.
    pub const OFF: NoiseParams = NoiseParams {
        flier_rate: 0.0,
        jitter_mm: 0,
    };
}

/// Deterministic uniform noise in `[0, 1)` for a pixel: a hash of
/// `(frame, x, y, salt)` — no RNG state, so playback is reproducible.
fn noise01(frame: u64, x: u32, y: u32, salt: u32) -> f32 {
    let mut h = (frame as u32)
        ^ x.wrapping_mul(0x9E37_79B1)
        ^ y.wrapping_mul(0x85EB_CA77)
        ^ salt.wrapping_mul(0xC2B2_AE3D);
    h ^= h >> 16;
    h = h.wrapping_mul(0x7FEB_352D);
    h ^= h >> 15;
    h = h.wrapping_mul(0x846C_A68B);
    h ^= h >> 16;
    h as f32 / u32::MAX as f32
}

/// Center-camera focal ratios (per resolution pixel): a ~70° horizontal
/// and 60° vertical field of view, like the Kinect v2 depth sensor.
const CENTER_FX: f32 = 0.714_074; // 1 / (2 * tan(35°)), FOV_x = 70°
const CENTER_FY: f32 = 0.866_025_4; // 1 / (2 * tan(30°)), FOV_y = 60°
/// Side-camera focal ratio: a 60° horizontal FOV, square pixels.
const SIDE_FX: f32 = 0.866_025_4; // 1 / (2 * tan(30°))

/// The development rig: a Kinect-like center camera flanked by two
/// webcams 0.8 m apart, converged on [`SUBJECT`].
///
/// The center camera approximates the Kinect v2 depth sensor (70° × 60°
/// field of view); the side cameras are typical 60° webcams. Pass smaller
/// resolutions to keep the same fields of view with fewer pixels — what
/// the test suites use to stay fast.
pub fn dev_rig_with_resolutions(center: UVec2, side: UVec2) -> Rig {
    // Kinect v2 depth: ~70° x 60° field of view.
    let center = Camera {
        id: "kinect".into(),
        pose: rig::CameraPose::looking_at(Vec3::ZERO, SUBJECT),
        intrinsics: rig::CameraIntrinsics {
            fx: CENTER_FX * center.x as f32,
            fy: CENTER_FY * center.y as f32,
            cx: center.x as f32 / 2.0,
            cy: center.y as f32 / 2.0,
        },
        resolution: center,
    };
    // Typical webcams with a 60 degree horizontal FOV.
    let make_side = |x: f32| Camera {
        id: if x < 0.0 { "left" } else { "right" }.to_owned(),
        pose: rig::CameraPose::looking_at(Vec3::new(x, 0.0, 0.0), SUBJECT),
        intrinsics: rig::CameraIntrinsics {
            fx: SIDE_FX * side.x as f32,
            fy: SIDE_FX * side.y as f32,
            cx: side.x as f32 / 2.0,
            cy: side.y as f32 / 2.0,
        },
        resolution: side,
    };
    Rig {
        left: make_side(-0.4),
        center,
        right: make_side(0.4),
    }
}

/// The development rig at full size: Kinect-like depth at 256x212
/// (~54k depth pixels, inside the plan's 50–150k splat budget) and
/// 640x480 webcams.
pub fn dev_rig() -> Rig {
    dev_rig_with_resolutions(UVec2::new(256, 212), UVec2::new(640, 480))
}

/// A synthetic source over the default development scene (see
/// [`dev_rig`]), at `fps` frames per second.
pub fn dev_source(fps: f32) -> SyntheticSource {
    dev_source_with(&dev_rig(), SUBJECT, fps)
}

/// A synthetic source over the default development scene with a custom
/// rig and subject position — what the configurable app uses. The rig's
/// center camera is the depth viewpoint; its side cameras provide the
/// side RGB.
pub fn dev_source_with(rig: &Rig, subject: Vec3, fps: f32) -> SyntheticSource {
    let (head, shoulders) = Scene::head_and_shoulders(subject);
    SyntheticSource::new(
        Scene::new().with(Scene::backdrop(subject)),
        vec![head, shoulders],
        rig.clone(),
        fps,
    )
}

/// A synthetic depth-camera frame source.
pub struct SyntheticSource {
    /// Static surroundings (backdrop).
    scene: Scene,
    /// The animated subject surfaces, translated by the sway offset each
    /// frame.
    subject: Vec<Surface>,
    /// The whole rig: the center camera is the depth viewpoint, the
    /// side cameras provide side RGB.
    rig: Rig,
    fps: f32,
    noise: NoiseParams,
    frame_index: u64,
}

impl SyntheticSource {
    pub fn new(scene: Scene, subject: Vec<Surface>, rig: Rig, fps: f32) -> Self {
        Self {
            scene,
            subject,
            rig,
            fps: fps.max(1.0),
            noise: NoiseParams::default(),
            frame_index: 0,
        }
    }

    /// Override the sensor noise; [`NoiseParams::OFF`] for clean frames.
    pub fn with_noise(mut self, noise: NoiseParams) -> Self {
        self.noise = noise;
        self
    }

    /// The number of frames produced so far.
    pub fn frame_index(&self) -> u64 {
        self.frame_index
    }

    /// The depth viewpoint: the rig's center camera.
    pub fn center(&self) -> &Camera {
        &self.rig.center
    }

    /// The deterministic subject sway offset for a given frame.
    fn sway_offset(&self, frame: u64) -> Vec3 {
        let phase = 2.0 * std::f32::consts::PI * (frame % SWAY_PERIOD_FRAMES) as f32
            / SWAY_PERIOD_FRAMES as f32;
        Vec3::new(SWAY_AMPLITUDE_M * phase.sin(), 0.0, 0.0)
    }

    /// The full scene for a given frame: static surroundings plus the
    /// swayed subject.
    fn frame_scene(&self, frame: u64) -> Scene {
        let offset = self.sway_offset(frame);
        let mut scene = self.scene.clone();
        for surface in &self.subject {
            scene.surfaces.push(surface.translated(offset));
        }
        scene
    }

    /// The colour image one camera sees: RGB, 3 bytes per pixel.
    fn cast_rgb(&self, scene: &Scene, camera: &Camera) -> Vec<u8> {
        let (width, height) = (camera.resolution.x, camera.resolution.y);
        let mut rgb = vec![0_u8; (width * height * 3) as usize];
        for y in 0..height {
            for x in 0..width {
                let (origin, dir) = camera.pixel_ray(x as f32 + 0.5, y as f32 + 0.5);
                let Some(hit) = scene.cast_ray(origin, dir) else {
                    continue;
                };
                let idx = (y * width + x) as usize;
                rgb[idx * 3] = (hit.color[0] * 255.0) as u8;
                rgb[idx * 3 + 1] = (hit.color[1] * 255.0) as u8;
                rgb[idx * 3 + 2] = (hit.color[2] * 255.0) as u8;
            }
        }
        rgb
    }

    /// Raycast one frame: the clean depth map and RGB the software
    /// Kinect sees.
    fn cast_frame(&self, frame: u64) -> (Vec<u16>, Vec<u8>) {
        let scene = self.frame_scene(frame);
        let camera = &self.rig.center;
        let (width, height) = (camera.resolution.x, camera.resolution.y);
        let camera_from_world = camera.pose.camera_from_world();
        let mut depth = vec![0_u16; (width * height) as usize];
        let mut rgb = vec![0_u8; (width * height * 3) as usize];

        for y in 0..height {
            for x in 0..width {
                // Sample through pixel centres.
                let (origin, dir) = camera.pixel_ray(x as f32 + 0.5, y as f32 + 0.5);
                let Some(hit) = scene.cast_ray(origin, dir) else {
                    continue; // no return: depth stays 0
                };
                // Planar depth along the viewing axis, in millimetres —
                // what a Kinect depth map reports.
                let z = camera_from_world.transform_point3(hit.position).z;
                let idx = (y * width + x) as usize;
                depth[idx] = (z * 1000.0).round().clamp(1.0, u16::MAX as f32) as u16;
                rgb[idx * 3] = (hit.color[0] * 255.0) as u8;
                rgb[idx * 3 + 1] = (hit.color[1] * 255.0) as u8;
                rgb[idx * 3 + 2] = (hit.color[2] * 255.0) as u8;
            }
        }
        (depth, rgb)
    }

    /// Apply the deterministic sensor noise to a clean frame.
    fn add_noise(&self, depth: &mut [u16], frame: u64) {
        if self.noise == NoiseParams::OFF {
            return;
        }
        let width = self.rig.center.resolution.x;
        for (i, d) in depth.iter_mut().enumerate() {
            if *d == 0 {
                continue; // invalid pixels stay invalid
            }
            let x = (i as u32) % width;
            let y = (i as u32) / width;
            if self.noise.jitter_mm > 0 {
                let jitter = noise01(frame, x, y, 1) * 2.0 - 1.0;
                let delta = (jitter * self.noise.jitter_mm as f32).round() as i32;
                *d = (*d as i32 + delta).clamp(1, u16::MAX as i32) as u16;
            }
            if self.noise.flier_rate > 0.0 && noise01(frame, x, y, 2) < self.noise.flier_rate {
                let span = (KINECT_MAX_MM - KINECT_MIN_MM) as f32;
                *d = KINECT_MIN_MM + (noise01(frame, x, y, 3) * span) as u16;
            }
        }
    }

    /// Build the raw depth frame for `frame`, including the side RGB.
    fn build_frame(&self, frame: u64) -> DepthFrame {
        let (mut depth, rgb) = self.cast_frame(frame);
        self.add_noise(&mut depth, frame);
        let scene = self.frame_scene(frame);
        let side_rgb = SideRgb {
            left: self.cast_rgb(&scene, &self.rig.left),
            left_resolution: self.rig.left.resolution,
            right: self.cast_rgb(&scene, &self.rig.right),
            right_resolution: self.rig.right.resolution,
        };
        DepthFrame {
            timestamp_secs: frame as f64 / self.fps as f64,
            depth_mm: depth,
            rgb,
            side_rgb: Some(side_rgb),
        }
    }
}

impl DepthSource for SyntheticSource {
    fn camera(&self) -> &Camera {
        &self.rig.center
    }

    fn rig(&self) -> Rig {
        self.rig.clone()
    }

    fn next_depth_frame(&mut self) -> Result<DepthFrame> {
        let frame = self.frame_index;
        self.frame_index += 1;
        Ok(self.build_frame(frame))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::scene;
    use glam::uvec2;
    use rig::DepthSource;

    /// A small-resolution source so tests stay fast: 32x24 depth pixels.
    fn tiny_source() -> SyntheticSource {
        let rig = dev_rig_with_resolutions(uvec2(32, 24), uvec2(64, 48));
        dev_source_with(&rig, SUBJECT, 30.0)
    }

    /// The straight-ahead ray passes below the head and hits the front
    /// of the shoulders: ~1.85 m out.
    const SHOULDER_FRONT_MM: f32 = 1852.0;

    #[test]
    fn frames_are_depth_maps_like_a_kinect() {
        let mut src = tiny_source();
        let frame = src.next_depth_frame().expect("frame");
        assert_eq!(frame.depth_mm.len(), 32 * 24);
        assert_eq!(frame.rgb.len(), 32 * 24 * 3);

        // Every ray hits something (the backdrop is infinite): all
        // depths valid.
        assert!(frame.depth_mm.iter().all(|&d| d > 0));

        // The pixel looking straight at the subject sees the front of
        // the shoulders (± noise), in the shoulders' colour.
        let mid = 12 * 32 + 16;
        assert!(
            (frame.depth_mm[mid] as f32 - SHOULDER_FRONT_MM).abs() < 40.0,
            "center depth {} mm",
            frame.depth_mm[mid]
        );
        let shoulders = &frame.rgb[mid * 3..mid * 3 + 3];
        assert!(
            (shoulders[0] as i32 - 51).abs() < 6
                && (shoulders[1] as i32 - 89).abs() < 6
                && (shoulders[2] as i32 - 140).abs() < 6,
            "shoulders rgb {shoulders:?}"
        );

        // A top-corner pixel sees the backdrop behind the subject.
        let corner = frame.depth_mm[32 + 2] as f32;
        assert!((corner - 2800.0).abs() < 200.0, "corner depth {corner}");
    }

    #[test]
    fn side_rgb_covers_the_webcams() {
        // Higher-resolution webcams so the thin tint flanks are sampled.
        let rig = dev_rig_with_resolutions(uvec2(32, 24), uvec2(160, 120));
        let mut src = dev_source_with(&rig, SUBJECT, 30.0);
        let frame = src.next_depth_frame().expect("frame");
        let sides = frame.side_rgb.as_ref().expect("side rgb");
        assert_eq!(sides.left.len(), 160 * 120 * 3);
        assert_eq!(sides.right.len(), 160 * 120 * 3);
        assert_eq!(sides.left_resolution, uvec2(160, 120));

        // The head's flanks carry the side tints: the left camera's
        // image must contain the left tint, the right camera's the
        // right tint.
        let has_tint = |rgb: &[u8], tint: [f32; 3]| {
            rgb.as_chunks::<3>().0.iter().any(|px| {
                (px[0] as i32 - (tint[0] * 255.0) as i32).abs() < 8
                    && (px[2] as i32 - (tint[2] * 255.0) as i32).abs() < 8
            })
        };
        assert!(
            has_tint(&sides.left, scene::HEAD_LEFT_TINT),
            "left tint in left image"
        );
        assert!(has_tint(&sides.right, scene::HEAD_RIGHT_TINT));
    }

    #[test]
    fn timestamps_advance_with_index() {
        let mut src = tiny_source();
        assert_eq!(src.next_depth_frame().unwrap().timestamp_secs, 0.0);
        assert!((src.next_depth_frame().unwrap().timestamp_secs - 1.0 / 30.0).abs() < 1e-9);
        assert!((src.next_depth_frame().unwrap().timestamp_secs - 2.0 / 30.0).abs() < 1e-9);
    }

    #[test]
    fn playback_is_deterministic() {
        let mut a = tiny_source();
        let mut b = tiny_source();
        for _ in 0..5 {
            assert_eq!(a.next_depth_frame().unwrap(), b.next_depth_frame().unwrap());
        }
    }

    #[test]
    fn subject_sways_and_repeats() {
        // The sway is periodic in the frame index; the sensor noise is
        // not, so exact repetition is asserted on clean frames.
        let mut src = tiny_source().with_noise(NoiseParams::OFF);
        let frame0 = src.next_depth_frame().unwrap();
        // After a quarter period the sway is near its maximum.
        let mut quarter = None;
        while src.frame_index() <= SWAY_PERIOD_FRAMES / 4 {
            quarter = Some(src.next_depth_frame().unwrap());
        }
        let quarter = quarter.expect("quarter frame produced");
        assert_ne!(frame0.depth_mm, quarter.depth_mm, "no visible sway");

        // The frame at index SWAY_PERIOD_FRAMES repeats frame 0 exactly.
        let mut period = None;
        while src.frame_index() <= SWAY_PERIOD_FRAMES {
            period = Some(src.next_depth_frame().unwrap());
        }
        let period = period.expect("period frame produced");
        assert_eq!(frame0.depth_mm, period.depth_mm, "sway must be periodic");
        assert_eq!(frame0.rgb, period.rgb);
        assert_eq!(frame0.side_rgb, period.side_rgb);

        // With noise on, the sway still moves the same pixels but the
        // exact map differs — noise must not be frame-periodic.
        let mut noisy = tiny_source();
        let n0 = noisy.next_depth_frame().unwrap();
        let mut quarter = None;
        while noisy.frame_index() <= SWAY_PERIOD_FRAMES / 4 {
            quarter = Some(noisy.next_depth_frame().unwrap());
        }
        assert_ne!(n0.depth_mm, quarter.unwrap().depth_mm);
    }

    #[test]
    fn sway_offset_is_deterministic_and_bounded() {
        let src = tiny_source();
        assert_eq!(src.sway_offset(0), Vec3::ZERO);
        assert_eq!(src.sway_offset(SWAY_PERIOD_FRAMES), Vec3::ZERO);
        let max = src.sway_offset(SWAY_PERIOD_FRAMES / 4);
        assert!(max.x > 0.99 * SWAY_AMPLITUDE_M, "sway max was {max:?}");
        for frame in 0..SWAY_PERIOD_FRAMES * 2 {
            let o = src.sway_offset(frame);
            assert!(o.x.abs() <= SWAY_AMPLITUDE_M + 1e-5);
        }
    }

    #[test]
    fn noise_is_deterministic_and_can_be_disabled() {
        // Noise must not break determinism, and OFF must give clean maps.
        let mut noisy = tiny_source();
        let mut clean = tiny_source().with_noise(NoiseParams::OFF);

        let noisy_frame = noisy.next_depth_frame().unwrap();
        let clean_frame = clean.next_depth_frame().unwrap();
        assert_ne!(noisy_frame.depth_mm, clean_frame.depth_mm);

        // Clean frames repeat exactly across sources.
        let mut clean2 = tiny_source().with_noise(NoiseParams::OFF);
        assert_eq!(clean_frame, clean2.next_depth_frame().unwrap());

        // And noisy ones too.
        let mut noisy2 = tiny_source();
        assert_eq!(noisy_frame, noisy2.next_depth_frame().unwrap());

        // Jitter is bounded: every noisy depth within ±jitter of the
        // clean one (fliers excepted).
        let mut fliers = 0;
        for (n, c) in noisy_frame.depth_mm.iter().zip(&clean_frame.depth_mm) {
            if (*n as i32 - *c as i32).abs() > NoiseParams::default().jitter_mm as i32 + 1 {
                fliers += 1;
            }
        }
        assert!(
            fliers as f32 <= 3.0 * NoiseParams::default().flier_rate * 32.0 * 24.0,
            "too many changed pixels: {fliers}"
        );
    }

    #[test]
    fn dev_rig_is_symmetric() {
        let rig = dev_rig();
        assert_eq!(rig.left.resolution, uvec2(640, 480));
        assert_eq!(rig.right.resolution, uvec2(640, 480));
        assert_eq!(rig.center.resolution, uvec2(256, 212));
        // The center camera sits exactly between the side cameras.
        let mid = rig.virtual_pose(0.5);
        assert!(mid.position.distance(rig.center.pose.position) < 1e-6);
        // And it looks at the subject.
        let forward = mid.rotation * Vec3::Z;
        let to_subject = (SUBJECT - mid.position).normalize();
        assert!(forward.dot(to_subject) > 0.999);
    }
}