crystal_ball 0.3.0

use crate::math::{Bounds3, Hit, Point3, Ray, Vec3, XYZEnum};
use crate::shapes::Shape;

/// A [bounding volume hierarchy](https://en.wikipedia.org/wiki/Bounding_volume_hierarchy) of [`Shape`]s.
///
/// In short, this allows for much faster render times
/// by storing the individual shapes in spatial nodes ([bounding volumes](Bounds3)).
/// Ray intersection is then performed to determine which nodes are hit.
/// This quickly rules out large numbers of shapes,
/// so there don't need to be any ray intersections against them.
///
/// The implementation is based on the [PBR book](https://pbr-book.org/3ed-2018/Primitives_and_Intersection_Acceleration/Bounding_Volume_Hierarchies).
#[derive(Clone, Debug)]
pub struct BVH<S: Shape> {
    /// The BVH's shapes.
    pub shapes: Vec<S>,
    /// The BVH's nodes.
    pub nodes: Vec<LinearBVHNode>,
}

impl<S: Shape> BVH<S> {
    /// Create a new [`BVH`].
    pub fn init(max_shapes_per_node: u8, mut shapes: Vec<S>) -> Self {
        let shapes_len = shapes.len();

        if shapes_len == 0 {
            return BVH {
                shapes,
                nodes: Vec::new(),
            };
        }

        let mut shape_infos = shapes
            .iter()
            .enumerate()
            .map(|(i, s)| BVHShapeInfo::new(i, s.bounds()))
            .collect();

        let mut total_nodes = 0;
        let mut ordered_shapes_indices = Vec::with_capacity(shapes_len);

        let root = Self::build(
            &mut shape_infos,
            0,
            shapes_len,
            &mut total_nodes,
            &mut ordered_shapes_indices,
            max_shapes_per_node,
        );

        let mut nodes = vec![LinearBVHNode::default(); total_nodes];
        let mut offset = 0;

        BVH::<S>::flatten_bvh_tree(&root, &mut nodes, &mut offset);

        // Sort shapes by ordered_shapes_indices
        for i in 0..shapes.len() {
            if ordered_shapes_indices[i] != i {
                let mut current_index = i;

                loop {
                    let target_idx = ordered_shapes_indices[current_index];
                    ordered_shapes_indices[current_index] = current_index;

                    if ordered_shapes_indices[target_idx] == target_idx {
                        break;
                    }

                    shapes.swap(current_index, target_idx);
                    current_index = target_idx;
                }
            }
        }

        BVH { shapes, nodes }
    }

    fn build(
        shape_infos: &mut Vec<BVHShapeInfo>,
        start: usize,
        end: usize,
        total_nodes: &mut usize,
        ordered_shapes_indices: &mut Vec<usize>,
        max_shapes_per_node: u8,
    ) -> BVHBuildNode {
        *total_nodes += 1;

        let mut bounds = Bounds3::default();

        for shape_info in shape_infos.iter().take(end).skip(start) {
            bounds = bounds.include_bounds(shape_info.bounds)
        }

        let shape_count = end - start;

        return if shape_count == 1 {
            let first_shape_offset = ordered_shapes_indices.len();

            for shape_info in shape_infos.iter().take(end).skip(start) {
                ordered_shapes_indices.push(shape_info.shape_number);
            }

            BVHBuildNode::new_leaf(first_shape_offset, shape_count, bounds)
        } else {
            let mut center_bounds = Bounds3::default();

            for shape_info in shape_infos.iter().take(end).skip(start) {
                center_bounds = center_bounds.include_point(shape_info.center)
            }

            let split_axis = center_bounds.maximum_extent();

            let mut middle = (start + end) / 2;

            if center_bounds.max[split_axis as usize] == center_bounds.min[split_axis as usize] {
                let first_shape_offset = ordered_shapes_indices.len();

                for shape_info in shape_infos.iter().take(end).skip(start) {
                    ordered_shapes_indices.push(shape_info.shape_number);
                }

                BVHBuildNode::new_leaf(first_shape_offset, shape_count, bounds)
            } else {
                if shape_count <= 2 {
                    middle = (start + end) / 2;

                    if start != end - 1
                        && shape_infos[end - 1].center[split_axis as usize]
                            < shape_infos[start].center[split_axis as usize]
                    {
                        shape_infos.swap(start, end - 1);
                    }
                } else {
                    let buckets_count = 12;
                    let mut bucket_infos = [BucketInfo::default(); 12];

                    for shape_info in shape_infos.iter().take(end).skip(start) {
                        let mut bucket = (buckets_count as f64
                            * center_bounds.offset(shape_info.center)[split_axis as usize])
                            as usize;

                        if bucket == buckets_count {
                            bucket = buckets_count - 1;
                        }

                        bucket_infos[bucket].count += 1;
                        bucket_infos[bucket].bounds = bucket_infos[bucket]
                            .bounds
                            .include_bounds(shape_info.bounds);
                    }

                    let mut costs = [0.0; 11];

                    for (i, cost) in costs.iter_mut().take(buckets_count - 1).enumerate() {
                        let (mut bounds1, mut bounds2) = (Bounds3::default(), Bounds3::default());
                        let (mut count1, mut count2) = (0, 0);

                        for bucket_info in bucket_infos.iter().take(i + 1) {
                            bounds1 = bounds1.include_bounds(bucket_info.bounds);
                            count1 += bucket_info.count;
                        }
                        for bucket_info in bucket_infos.iter().take(buckets_count).skip(i + 1) {
                            bounds2 = bounds2.include_bounds(bucket_info.bounds);
                            count2 += bucket_info.count;
                        }

                        *cost = 1.0
                            + (count1 as f64 * bounds1.surface_area()
                                + count2 as f64 * bounds2.surface_area())
                                / bounds.surface_area();
                    }

                    let (min_cost_index, min_cost) = (0..buckets_count - 1).fold(
                        (0, costs[0]),
                        |(min_cost_index, min_cost), i| {
                            if costs[i] < min_cost {
                                (i, costs[i])
                            } else {
                                (min_cost_index, min_cost)
                            }
                        },
                    );

                    let leaf_cost = shape_count as f64;

                    if shape_count > max_shapes_per_node as usize || min_cost < leaf_cost {
                        let (mut left, mut right): (Vec<BVHShapeInfo>, Vec<BVHShapeInfo>) =
                            shape_infos[start..end].iter().partition(|s| {
                                let mut buckets = (buckets_count as f64
                                    * center_bounds.offset(s.center)[split_axis as usize])
                                    as usize;

                                if buckets == buckets_count {
                                    buckets = buckets_count - 1;
                                }

                                buckets <= min_cost_index
                            });

                        middle = start + left.len();

                        if (left.len() + right.len()) == shape_infos.len() {
                            shape_infos.clear();
                            shape_infos.append(&mut left);
                            shape_infos.append(&mut right);
                        } else {
                            shape_infos.splice(start..middle, left);
                            shape_infos.splice(middle..end, right);
                        }
                    } else {
                        let first_shape_offset = ordered_shapes_indices.len();

                        for shape_info in shape_infos.iter().take(end).skip(start) {
                            ordered_shapes_indices.push(shape_info.shape_number);
                        }

                        return BVHBuildNode::new_leaf(first_shape_offset, shape_count, bounds);
                    }
                }

                BVHBuildNode::new_interior(
                    split_axis,
                    Box::new(Self::build(
                        shape_infos,
                        start,
                        middle,
                        total_nodes,
                        ordered_shapes_indices,
                        max_shapes_per_node,
                    )),
                    Box::new(Self::build(
                        shape_infos,
                        middle,
                        end,
                        total_nodes,
                        ordered_shapes_indices,
                        max_shapes_per_node,
                    )),
                )
            }
        };
    }

    fn flatten_bvh_tree(
        root: &BVHBuildNode,
        nodes: &mut [LinearBVHNode],
        offset: &mut usize,
    ) -> usize {
        let my_offset = *offset;
        *offset += 1;

        if root.shapes_count > 0 {
            let linear_node = LinearBVHNode::new(
                root.bounds,
                root.first_shape_offset as i32,
                root.shapes_count as u16,
                XYZEnum::X,
            );
            nodes[my_offset] = linear_node;
        } else {
            if let Some(ref child1) = root.child1 {
                BVH::<S>::flatten_bvh_tree(child1, nodes, offset);
            }
            if let Some(ref child2) = root.child2 {
                let linear_node = LinearBVHNode::new(
                    root.bounds,
                    BVH::<S>::flatten_bvh_tree(child2, nodes, offset) as i32,
                    0,
                    root.split_axis,
                );
                nodes[my_offset] = linear_node;
            }
        }

        my_offset
    }

    // TODO: Store the material in [`Hit`] so there's no need to return the shape.
    /// Performs a ray intersection against `self`,
    /// returning the closest hit if there is one as well as the corresponding shape.
    pub fn intersects(&self, ray: Ray) -> Option<(Hit, &S)> {
        if self.nodes.is_empty() {
            return None;
        }

        let mut closest_hit = None;
        let mut closest_hit_distance = f64::INFINITY;

        let inverse_direction = Vec3::new(
            1.0 / ray.direction.x,
            1.0 / ray.direction.y,
            1.0 / ray.direction.z,
        );
        let direction_is_negative = [
            inverse_direction.x < 0.0,
            inverse_direction.y < 0.0,
            inverse_direction.z < 0.0,
        ];

        let mut to_visit_offset = 0;
        let mut current_node_index = 0;
        let mut nodes_to_visit = [0; 64];

        loop {
            let node = self.nodes[current_node_index];

            if node.bounds.intersects(
                ray,
                closest_hit_distance,
                inverse_direction,
                direction_is_negative,
            ) {
                if node.shapes_count > 0 {
                    for i in 0..node.shapes_count {
                        let shape = &self.shapes[node.offset as usize + i as usize];

                        if let Some(hit) = shape.intersects(ray) {
                            if hit.distance < closest_hit_distance {
                                closest_hit = Some((hit, shape));
                                closest_hit_distance = hit.distance;
                            }
                        }
                    }

                    if to_visit_offset == 0 {
                        break;
                    }

                    to_visit_offset -= 1;
                    current_node_index = nodes_to_visit[to_visit_offset];
                } else if direction_is_negative[node.axis as usize] {
                    nodes_to_visit[to_visit_offset] = current_node_index + 1;

                    to_visit_offset += 1;
                    current_node_index = node.offset as usize;
                } else {
                    nodes_to_visit[to_visit_offset] = node.offset as usize;

                    to_visit_offset += 1;
                    current_node_index += 1;
                }
            } else {
                if to_visit_offset == 0 {
                    break;
                }

                to_visit_offset -= 1;
                current_node_index = nodes_to_visit[to_visit_offset];
            }
        }

        closest_hit
    }
}

#[derive(Copy, Clone, Debug)]
struct BVHShapeInfo {
    shape_number: usize,
    bounds: Bounds3,
    center: Point3,
}

impl BVHShapeInfo {
    fn new(shape_number: usize, bounds: Bounds3) -> Self {
        BVHShapeInfo {
            shape_number,
            bounds,
            center: 0.5 * bounds.min + 0.5 * bounds.max,
        }
    }
}

#[derive(Clone, Debug)]
struct BVHBuildNode {
    bounds: Bounds3,
    child1: Option<Box<BVHBuildNode>>,
    child2: Option<Box<BVHBuildNode>>,
    split_axis: XYZEnum,
    first_shape_offset: usize,
    shapes_count: usize,
}

impl BVHBuildNode {
    fn new_leaf(first_shape_offset: usize, shapes_count: usize, bounds: Bounds3) -> Self {
        BVHBuildNode {
            bounds,
            child1: None,
            child2: None,
            split_axis: XYZEnum::X,
            first_shape_offset,
            shapes_count,
        }
    }

    fn new_interior(
        split_axis: XYZEnum,
        child1: Box<BVHBuildNode>,
        child2: Box<BVHBuildNode>,
    ) -> Self {
        BVHBuildNode {
            bounds: child1.bounds.include_bounds(child2.bounds),
            child1: Some(child1),
            child2: Some(child2),
            split_axis,
            first_shape_offset: 0,
            shapes_count: 0,
        }
    }
}

#[derive(Copy, Clone, Default, Debug)]
struct BucketInfo {
    count: usize,
    bounds: Bounds3,
}

/// A node of the [`BVH`].
#[derive(Copy, Clone, Default, Debug)]
pub struct LinearBVHNode {
    pub bounds: Bounds3,
    pub offset: i32,
    pub shapes_count: u16,
    pub axis: XYZEnum,
    _padding: u8,
}

impl LinearBVHNode {
    /// Create a new [`LinearBvhNode`].
    pub fn new(bounds: Bounds3, offset: i32, shapes_count: u16, axis: XYZEnum) -> Self {
        LinearBVHNode {
            bounds,
            offset,
            shapes_count,
            axis,
            _padding: 0,
        }
    }
}