//! The synthetic frame source: a software depth camera.
//!
//! Every frame, [`SyntheticSource`] raycasts the scene from the rig's
//! center (Kinect) viewpoint — one ray per depth pixel — and turns each
//! hit into a footprint-sized surfel via [`rig::Splat::surfel`], the same
//! construction the real `splatgen` uses on Kinect depth. The subject
//! sways a little so temporal behaviour is testable, deterministically
//! (a function of the frame index, never the wall clock).

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

use rig::{
    Camera, CameraIntrinsics, CameraPose, Frame, FrameSource, Rig, SURFEL_THICKNESS_M, Splat,
    SplatBatch,
};

use crate::scene::{self, 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);

/// Opacity used for synthetic surfels (PLAN.md: ~0.9–0.99).
pub const SURFEL_OPACITY: f32 = 0.95;

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

/// 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: CameraPose::looking_at(Vec3::ZERO, SUBJECT),
        intrinsics: 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: CameraPose::looking_at(Vec3::new(x, 0.0, 0.0), SUBJECT),
        intrinsics: 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),
    }
}

/// 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 at full size: Kinect-like depth at 256x212
/// (~54k active splats, inside the plan's 50–150k 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_center_camera(dev_rig().center, fps)
}

/// Like [`dev_source`], but with a custom center (depth) camera and
/// subject position — what the configurable app uses.
pub fn dev_source_with(center: Camera, subject: Vec3, fps: f32) -> SyntheticSource {
    let (head, shoulders) = scene::head_and_shoulders(subject);
    SyntheticSource::new(
        Scene::new().with(scene::backdrop(subject)),
        vec![head, shoulders],
        center,
        fps,
    )
}

/// Like [`dev_source`], but with a custom center (depth) camera — e.g.
/// `dev_rig_with_resolutions(..).center` at a small resolution for tests.
pub fn dev_source_with_center_camera(center: Camera, fps: f32) -> SyntheticSource {
    dev_source_with(center, SUBJECT, 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 depth viewpoint: the rig's center camera.
    center: Camera,
    fps: f32,
    frame_index: u64,
}

impl SyntheticSource {
    pub fn new(scene: Scene, subject: Vec<Surface>, center: Camera, fps: f32) -> Self {
        Self {
            scene,
            subject,
            center,
            fps: fps.max(1.0),
            frame_index: 0,
        }
    }

    /// 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.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
    }

    /// Build the splat batch the software depth camera sees in `frame`.
    fn build_frame(&self, frame: u64) -> SplatBatch {
        let scene = self.frame_scene(frame);
        let (width, height) = (self.center.resolution.x, self.center.resolution.y);
        let world_from_camera = self.center.pose.world_from_camera();
        let mut batch = SplatBatch::with_capacity((width * height) as usize);

        for y in 0..height {
            for x in 0..width {
                // Sample through pixel centres.
                let (origin, dir) = self.center.pixel_ray(x as f32 + 0.5, y as f32 + 0.5);
                let Some(hit) = scene.cast_ray(origin, dir) else {
                    continue;
                };
                // Depth along the viewing axis (like a Kinect depth map,
                // which is what PLAN.md's footprint formula expects).
                let depth = world_from_camera.inverse().transform_point3(hit.position).z;
                batch.push(Splat::surfel(
                    hit.position,
                    hit.normal,
                    depth,
                    &self.center.intrinsics,
                    SURFEL_THICKNESS_M,
                    hit.color,
                    SURFEL_OPACITY,
                ));
            }
        }
        batch
    }
}

impl FrameSource for SyntheticSource {
    fn next_frame(&mut self) -> Result<Frame> {
        let frame = self.frame_index;
        self.frame_index += 1;
        Ok(Frame {
            timestamp_secs: frame as f64 / self.fps as f64,
            splats: self.build_frame(frame),
        })
    }
}

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

    /// A small-resolution source so tests stay fast: 32x24 depth pixels.
    fn tiny_source() -> SyntheticSource {
        let (head, shoulders) = scene::head_and_shoulders(SUBJECT);
        let mut center = dev_rig().center;
        // Keep the center camera's FOV, shrink the resolution.
        let scale = 32.0 / center.resolution.x as f32;
        center.intrinsics.fx *= scale;
        center.intrinsics.fy *= scale;
        center.intrinsics.cx = 16.0;
        center.intrinsics.cy = 12.0;
        center.resolution = UVec2::new(32, 24);
        SyntheticSource::new(
            Scene::new().with(scene::backdrop(SUBJECT)),
            vec![head, shoulders],
            center,
            30.0,
        )
    }

    #[test]
    fn frames_are_nonempty_and_in_front() {
        let mut src = tiny_source();
        let frame = src.next_frame().unwrap();
        assert!(!frame.splats.is_empty());
        // One splat per depth pixel that hit — never more.
        assert!(frame.splats.len() <= 32 * 24);
        // Most rays hit the backdrop.
        assert!(frame.splats.len() > 32 * 24 / 2);

        // Every splat lies in front of the center camera and matches the
        // footprint formula.
        let cam = src.center().clone();
        let camera_from_world = cam.pose.camera_from_world();
        for splat in frame.splats.iter() {
            let p = camera_from_world.transform_point3(splat.mean);
            assert!(p.z > 0.5, "splat behind camera: {p:?}");
            assert!((splat.scale.x - p.z / cam.intrinsics.fx).abs() < 1e-5);
            assert!((splat.scale.y - p.z / cam.intrinsics.fy).abs() < 1e-5);
            assert!((splat.opacity - SURFEL_OPACITY).abs() < 1e-6);
        }
    }

    #[test]
    fn subject_and_backdrop_both_present() {
        let mut src = tiny_source();
        let frame = src.next_frame().unwrap();
        let mut skin = 0;
        let mut checker = 0;
        for splat in frame.splats.iter() {
            if (splat.color[0] - 0.85).abs() < 0.01 {
                skin += 1;
            }
            // Checker greys: 0.72 or 0.28.
            if (splat.color[0] - splat.color[1]).abs() < 1e-4
                && ((splat.color[0] - 0.72).abs() < 0.01 || (splat.color[0] - 0.28).abs() < 0.01)
            {
                checker += 1;
            }
        }
        assert!(skin > 0, "no subject splats");
        assert!(checker > 0, "no backdrop splats");
    }

    #[test]
    fn timestamps_advance_with_index() {
        let mut src = tiny_source();
        assert_eq!(src.next_frame().unwrap().timestamp_secs, 0.0);
        assert!((src.next_frame().unwrap().timestamp_secs - 1.0 / 30.0).abs() < 1e-9);
        assert!((src.next_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_frame().unwrap(), b.next_frame().unwrap());
        }
    }

    #[test]
    fn subject_sways_and_repeats() {
        let mut src = tiny_source();
        let frame0 = src.next_frame().unwrap();

        // The frame produced at a quarter period (index 22) is at (nearly)
        // maximum sway.
        let mut quarter = None;
        while src.frame_index() <= SWAY_PERIOD_FRAMES / 4 {
            quarter = Some(src.next_frame().unwrap());
        }
        let quarter = quarter.expect("quarter frame produced");
        assert_ne!(
            frame0.splats, quarter.splats,
            "no visible sway at a quarter period"
        );

        // 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_frame().unwrap());
        }
        let period = period.expect("period frame produced");
        assert_eq!(
            frame0.splats, period.splats,
            "sway must be periodic in the frame index"
        );
    }

    #[test]
    fn sway_offset_is_deterministic_and_bounded() {
        let src = tiny_source();
        let offset = src.sway_offset(0);
        assert_eq!(offset, Vec3::ZERO);
        assert_eq!(src.sway_offset(SWAY_PERIOD_FRAMES), Vec3::ZERO);
        // A quarter period (90/4 = 22 frames) is close to the maximum.
        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 dev_rig_is_symmetric() {
        let rig = dev_rig();
        assert_eq!(rig.left.resolution, UVec2::new(640, 480));
        assert_eq!(rig.right.resolution, UVec2::new(640, 480));
        assert_eq!(rig.center.resolution, UVec2::new(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);
    }
}