//! Camera poses, pinhole intrinsics, and projection math.

use glam::{Affine3A, Mat3, Quat, UVec2, Vec2, Vec3};

/// A pinhole camera's intrinsics, in pixels at the camera's native
/// resolution.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct CameraIntrinsics {
    /// Focal length in pixels, horizontal.
    pub fx: f32,
    /// Focal length in pixels, vertical.
    pub fy: f32,
    /// Principal point in pixels (image centre for a well-calibrated rig).
    pub cx: f32,
    pub cy: f32,
}

impl CameraIntrinsics {
    /// Project a camera-space point to pixel coordinates `(u, v)`.
    ///
    /// Returns `None` for points on or behind the camera plane.
    pub fn project_camera_point(&self, p: Vec3) -> Option<Vec2> {
        if p.z <= f32::EPSILON {
            return None;
        }
        Some(Vec2::new(
            self.fx * p.x / p.z + self.cx,
            self.fy * p.y / p.z + self.cy,
        ))
    }

    /// The camera-space direction (unnormalized) of a pixel's view ray.
    pub fn camera_ray(&self, u: f32, v: f32) -> Vec3 {
        Vec3::new((u - self.cx) / self.fx, (v - self.cy) / self.fy, 1.0)
    }
}

/// The position and orientation of a camera in world space.
///
/// The rotation is camera-to-world; the camera looks along its local +Z
/// (COLMAP convention). With [`CameraPose::identity`] the camera frame
/// equals the world frame, so it looks down world +Z with image-right at
/// world +X — and the image upside-down in a Y-up world (an identity pose
/// is just a pose; use [`CameraPose::looking_at`] for a right-side-up
/// view).
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct CameraPose {
    pub position: Vec3,
    pub rotation: Quat,
}

impl CameraPose {
    /// Camera at the origin, looking down world +Z.
    pub fn identity() -> Self {
        Self {
            position: Vec3::ZERO,
            rotation: Quat::IDENTITY,
        }
    }

    /// A pose at `position` that looks at `target` with the world up
    /// (+Y) at the top of the image.
    ///
    /// Falls back to +Z as "up" when the view direction is near-vertical.
    pub fn looking_at(position: Vec3, target: Vec3) -> Self {
        let forward = (target - position).normalize_or_zero();
        // Camera basis: x = image-right, y = image-down, z = forward,
        // right-handed, so x = y × z and y = z × x.
        let up = if forward.abs().dot(Vec3::Y) > 0.999 {
            Vec3::Z
        } else {
            Vec3::Y
        };
        let right = (-up).cross(forward).normalize_or_zero();
        let down = forward.cross(right).normalize_or_zero();
        let rotation = Quat::from_mat3(&Mat3::from_cols(right, down, forward));
        Self { position, rotation }
    }

    /// A pose at `position` with aviation-style rotation in degrees:
    /// yaw turns the camera (positive towards world +X), pitch tilts it
    /// (positive up), roll rotates about the viewing axis. Rotations are
    /// intrinsic, applied yaw → pitch → roll.
    ///
    /// `(0, 0, 0)` looks straight down world +Z with the world up at
    /// the top of the image — the same orientation
    /// [`CameraPose::looking_at`] produces for a target straight ahead,
    /// and pure yaw matches `looking_at` exactly.
    pub fn from_euler_deg(position: Vec3, pitch_deg: f32, yaw_deg: f32, roll_deg: f32) -> Self {
        let yaw = Quat::from_rotation_y(yaw_deg.to_radians());
        let pitch = Quat::from_rotation_x(-pitch_deg.to_radians());
        let roll = Quat::from_rotation_z(roll_deg.to_radians());
        // The "looking at +Z, right-side-up" frame is the COLMAP identity
        // rolled 180° about the viewing axis (see the crate docs).
        let upright = Quat::from_rotation_z(std::f32::consts::PI);
        Self {
            position,
            rotation: yaw * pitch * roll * upright,
        }
    }

    /// The camera-to-world transform.
    pub fn world_from_camera(&self) -> Affine3A {
        Affine3A::from_rotation_translation(self.rotation, self.position)
    }

    /// The world-to-camera transform.
    pub fn camera_from_world(&self) -> Affine3A {
        self.world_from_camera().inverse()
    }

    /// Interpolate from `self` at `t = 0` to `other` at `t = 1`.
    ///
    /// Positions lerp, orientations slerp. Used to place the virtual
    /// camera between the physical rig cameras.
    pub fn blend(&self, other: &Self, t: f32) -> Self {
        Self {
            position: self.position.lerp(other.position, t),
            rotation: self.rotation.slerp(other.rotation, t),
        }
    }
}

/// A pinhole camera: pose, intrinsics, and image resolution.
#[derive(Debug, Clone, PartialEq)]
pub struct Camera {
    /// Camera identifier: the device name/serial it maps to on the
    /// physical rig (e.g. the webcam path for capture, the Kinect for
    /// the depth source).
    pub id: String,
    pub pose: CameraPose,
    pub intrinsics: CameraIntrinsics,
    /// Image size in pixels: `(width, height)`.
    pub resolution: UVec2,
}

impl Camera {
    /// Project a world-space point to pixel coordinates, or `None` if the
    /// point is behind the camera.
    pub fn project(&self, world: Vec3) -> Option<Vec2> {
        let p = self.pose.camera_from_world().transform_point3(world);
        self.intrinsics.project_camera_point(p)
    }

    /// The world-space ray through the centre of pixel `(u, v)`:
    /// `(origin, normalized_direction)`.
    pub fn pixel_ray(&self, u: f32, v: f32) -> (Vec3, Vec3) {
        let dir_cam = self.intrinsics.camera_ray(u, v).normalize_or_zero();
        let dir_world = self.pose.rotation * dir_cam;
        (self.pose.position, dir_world)
    }

    /// Horizontal field of view in radians, matching Brush's convention
    /// (`focal = (width / 2) / tan(fov / 2)`).
    pub fn fov_x(&self) -> f32 {
        2.0 * ((self.resolution.x as f32 / 2.0) / self.intrinsics.fx).atan()
    }

    /// Vertical field of view in radians.
    pub fn fov_y(&self) -> f32 {
        2.0 * ((self.resolution.y as f32 / 2.0) / self.intrinsics.fy).atan()
    }
}

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

    fn cam_looking_at_subject() -> Camera {
        Camera {
            id: "test".into(),
            pose: CameraPose::looking_at(Vec3::ZERO, vec3(0.0, 0.0, 2.0)),
            intrinsics: CameraIntrinsics {
                fx: 500.0,
                fy: 500.0,
                cx: 320.0,
                cy: 240.0,
            },
            resolution: uvec2(640, 480),
        }
    }

    #[test]
    fn identity_pose_looks_down_plus_z() {
        let cam = Camera {
            id: "test".into(),
            pose: CameraPose::identity(),
            intrinsics: CameraIntrinsics {
                fx: 100.0,
                fy: 100.0,
                cx: 50.0,
                cy: 50.0,
            },
            resolution: uvec2(100, 100),
        };
        // A point straight ahead at +Z maps to the principal point.
        assert_eq!(cam.project(vec3(0.0, 0.0, 5.0)), Some(vec2(50.0, 50.0)));
        // Camera frame = world frame: image-right is +X, image-down is +Y
        // (in a Y-up world this pose renders upside-down; it's the COLMAP
        // identity, see CameraPose docs).
        assert!(cam.project(vec3(1.0, 0.0, 5.0)).unwrap().x > 50.0);
        assert!(cam.project(vec3(0.0, 1.0, 5.0)).unwrap().y > 50.0);
        // Behind the camera: no projection.
        assert_eq!(cam.project(vec3(0.0, 0.0, -5.0)), None);
    }

    #[test]
    fn looking_at_is_right_side_up() {
        let cam = cam_looking_at_subject();
        // The subject's head (above the look-at point) must land above the
        // principal point: world up is image up.
        let head = cam.project(vec3(0.0, 0.5, 2.0)).expect("in front");
        assert!(head.y < cam.intrinsics.cy, "head projected to {head:?}");
        // Something to the camera's right (world -X when looking down +Z,
        // see CameraPose docs) lands on the right half of the image.
        let right = cam.project(vec3(-0.5, 0.0, 2.0)).expect("in front");
        assert!(right.x > cam.intrinsics.cx);
        // ...and world +X lands on the left half.
        let left = cam.project(vec3(0.5, 0.0, 2.0)).expect("in front");
        assert!(left.x < cam.intrinsics.cx);
    }

    #[test]
    fn pixel_ray_round_trips() {
        let cam = cam_looking_at_subject();
        // Ray through the principal point hits the look-at target.
        let (o, d) = cam.pixel_ray(320.0, 240.0);
        let hit = o + d * 2.0;
        assert!(hit.distance_squared(vec3(0.0, 0.0, 2.0)) < 1e-8);

        // Marching a ray and projecting back gives the same pixel.
        let (o, d) = cam.pixel_ray(123.0, 45.0);
        let p = o + d * 3.7;
        let projected = cam.project(p).expect("in front");
        assert!((projected - vec2(123.0, 45.0)).length() < 1e-3);
    }

    #[test]
    fn euler_matches_looking_at_for_pure_yaw() {
        for yaw_deg in [0.0_f32, 11.31, -30.0, 75.0] {
            let position = vec3(-0.4, 0.1, 0.0);
            let yaw = yaw_deg.to_radians();
            let target = position + vec3(yaw.sin(), 0.0, yaw.cos()) * 2.0;
            let euler = CameraPose::from_euler_deg(position, 0.0, yaw_deg, 0.0);
            let look = CameraPose::looking_at(position, target);
            assert!((euler.rotation * Vec3::Z - look.rotation * Vec3::Z).length() < 1e-5);
            assert!((euler.rotation * Vec3::Y - look.rotation * Vec3::Y).length() < 1e-5);
            assert!((euler.rotation * Vec3::X - look.rotation * Vec3::X).length() < 1e-5);
        }
    }

    #[test]
    fn euler_pitch_and_roll_axes() {
        let zero = CameraPose::from_euler_deg(Vec3::ZERO, 0.0, 0.0, 0.0);
        // Zero euler: looking at +Z, right-side-up (image-down = world -Y).
        assert!((zero.rotation * Vec3::Z - Vec3::Z).length() < 1e-6);
        assert!((zero.rotation * Vec3::Y + Vec3::Y).length() < 1e-6);

        // Positive pitch tilts the camera up (forward gains world +Y).
        let pitch = CameraPose::from_euler_deg(Vec3::ZERO, 10.0, 0.0, 0.0);
        assert!((pitch.rotation * Vec3::Z).y > 0.1);

        // Positive yaw turns towards world +X.
        let yaw = CameraPose::from_euler_deg(Vec3::ZERO, 0.0, 30.0, 0.0);
        assert!((yaw.rotation * Vec3::Z).x > 0.4);

        // Roll rotates about the viewing axis: the forward vector holds.
        let roll = CameraPose::from_euler_deg(Vec3::ZERO, 0.0, 0.0, 90.0);
        assert!((roll.rotation * Vec3::Z - Vec3::Z).length() < 1e-6);
    }

    #[test]
    fn pose_blend_hits_endpoints() {
        let a = CameraPose::looking_at(vec3(-1.0, 0.0, 0.0), vec3(0.0, 0.0, 2.0));
        let b = CameraPose::looking_at(vec3(1.0, 0.0, 0.0), vec3(0.0, 0.0, 2.0));
        assert_eq!(a.blend(&b, 0.0), a);
        assert_eq!(a.blend(&b, 1.0), b);
        let mid = a.blend(&b, 0.5);
        // Halfway along the baseline, still looking at the subject.
        assert!(mid.position.length() < 1e-6);
        let forward = mid.rotation * Vec3::Z;
        assert!(forward.z > 0.99);
    }

    #[test]
    fn fov_matches_focal() {
        let cam = cam_looking_at_subject();
        // Brush's formula: focal = (w/2) / tan(fov/2).
        let focal = (640.0 / 2.0) / (cam.fov_x() / 2.0).tan();
        assert!((focal - 500.0).abs() < 1e-3);
        let focal_y = (480.0 / 2.0) / (cam.fov_y() / 2.0).tan();
        assert!((focal_y - 500.0).abs() < 1e-3);
    }
}