psychopath/src/scene/assembly.rs

use std::collections::HashMap;

use kioku::Arena;

use crate::{
    accel::BVH4,
    accel::{LightAccel, LightTree},
    bbox::{transform_bbox_slice_from, BBox},
    boundable::Boundable,
    color::SpectralSample,
    lerp::lerp_slice,
    light::SurfaceLight,
    math::{Normal, Point, Xform, XformFull},
    shading::SurfaceShader,
    surface::{Surface, SurfaceIntersection},
};

#[derive(Copy, Clone, Debug)]
pub struct Assembly<'a> {
    // Instance list
    pub instances: &'a [Instance],
    pub light_instances: &'a [Instance],
    pub xforms: &'a [Xform],

    // Surface shader list
    pub surface_shaders: &'a [&'a dyn SurfaceShader],

    // Object list
    pub objects: &'a [Object<'a>],

    // Assembly list
    pub assemblies: &'a [Assembly<'a>],

    // Object accel
    pub object_accel: BVH4<'a>,

    // Light accel
    pub light_accel: LightTree<'a>,
}

// TODO: actually fix this clippy warning, rather than `allow`ing it.
#[allow(clippy::type_complexity)]
impl<'a> Assembly<'a> {
    // Returns (light_color, (sample_point, normal, point_err), pdf, selection_pdf)
    pub fn sample_lights(
        &self,
        n: f32,
        uvw: (f32, f32, f32),
        wavelength: f32,
        time: f32,
        space: &XformFull,
        intr: &SurfaceIntersection,
    ) -> Option<(SpectralSample, (Point, Normal, f32), f32, f32)> {
        if let SurfaceIntersection::Hit {
            intersection_data: idata,
            closure,
        } = *intr
        {
            if let Some((light_i, sel_pdf, whittled_n)) = self.light_accel.select(
                idata.incoming.xform_inv(space),
                idata.pos.xform_inv(space),
                idata.nor.xform_inv_fast(space),
                idata.nor_g.xform_inv_fast(space),
                &closure,
                time,
                n,
            ) {
                let inst = self.light_instances[light_i];

                // Handle transforms.
                let local_space = if let Some((a, b)) = inst.transform_indices {
                    if let Some(new_space) = lerp_slice(&self.xforms[a..b], time)
                        .compose(&space.fwd)
                        .to_full()
                    {
                        new_space
                    } else {
                        // Invalid transform.  Give up.
                        return None;
                    }
                } else {
                    *space
                };

                match inst.instance_type {
                    InstanceType::Object => {
                        match self.objects[inst.data_index] {
                            Object::SurfaceLight(light) => {
                                // Sample the light
                                let (color, sample_geo, pdf) = light.sample_from_point(
                                    &local_space,
                                    idata.pos,
                                    uvw.0,
                                    uvw.1,
                                    wavelength,
                                    time,
                                );
                                return Some((color, sample_geo, pdf, sel_pdf));
                            }

                            _ => unimplemented!(),
                        }
                    }

                    InstanceType::Assembly => {
                        // Sample sub-assembly lights
                        let sample = self.assemblies[inst.data_index].sample_lights(
                            whittled_n,
                            uvw,
                            wavelength,
                            time,
                            &local_space,
                            intr,
                        );

                        // Return sample
                        return sample.map(|(ss, v, pdf, spdf)| (ss, v, pdf, spdf * sel_pdf));
                    }
                }
            } else {
                None
            }
        } else {
            None
        }
    }
}

impl<'a> Boundable for Assembly<'a> {
    fn bounds(&self) -> &[BBox] {
        self.object_accel.bounds()
    }
}

#[derive(Debug)]
pub struct AssemblyBuilder<'a> {
    arena: &'a Arena,

    // Instance list
    instances: Vec<Instance>,
    xforms: Vec<Xform>,

    // Shader list
    surface_shaders: Vec<&'a dyn SurfaceShader>,
    surface_shader_map: HashMap<String, usize>, // map Name -> Index

    // Object list
    objects: Vec<Object<'a>>,
    object_map: HashMap<String, usize>, // map Name -> Index

    // Assembly list
    assemblies: Vec<Assembly<'a>>,
    assembly_map: HashMap<String, usize>, // map Name -> Index
}

impl<'a> AssemblyBuilder<'a> {
    pub fn new(arena: &'a Arena) -> AssemblyBuilder<'a> {
        AssemblyBuilder {
            arena: arena,
            instances: Vec::new(),
            xforms: Vec::new(),
            surface_shaders: Vec::new(),
            surface_shader_map: HashMap::new(),
            objects: Vec::new(),
            object_map: HashMap::new(),
            assemblies: Vec::new(),
            assembly_map: HashMap::new(),
        }
    }

    pub fn add_surface_shader(&mut self, name: &str, shader: &'a dyn SurfaceShader) {
        // Make sure the name hasn't already been used.
        if self.surface_shader_map.contains_key(name) {
            panic!("Attempted to add surface shader to assembly with a name that already exists.");
        }

        // Add shader
        self.surface_shader_map
            .insert(name.to_string(), self.surface_shaders.len());
        self.surface_shaders.push(shader);
    }

    pub fn add_object(&mut self, name: &str, obj: Object<'a>) {
        // Make sure the name hasn't already been used.
        if self.name_exists(name) {
            panic!("Attempted to add object to assembly with a name that already exists.");
        }

        // Add object
        self.object_map.insert(name.to_string(), self.objects.len());
        self.objects.push(obj);
    }

    pub fn add_assembly(&mut self, name: &str, asmb: Assembly<'a>) {
        // Make sure the name hasn't already been used.
        if self.name_exists(name) {
            panic!(
                "Attempted to add assembly to another assembly with a name that already \
                 exists."
            );
        }

        // Add assembly
        self.assembly_map
            .insert(name.to_string(), self.assemblies.len());
        self.assemblies.push(asmb);
    }

    pub fn add_instance(
        &mut self,
        name: &str,
        surface_shader_name: Option<&str>,
        xforms: Option<&[Xform]>,
    ) {
        // Make sure name exists
        if !self.name_exists(name) {
            panic!("Attempted to add instance with a name that doesn't exist.");
        }

        // Map zero-length transforms to None
        let xforms = if let Some(xf) = xforms {
            if !xf.is_empty() {
                Some(xf)
            } else {
                None
            }
        } else {
            None
        };

        // Create instance
        let instance = if self.object_map.contains_key(name) {
            Instance {
                instance_type: InstanceType::Object,
                data_index: self.object_map[name],
                surface_shader_index: surface_shader_name.map(|name| {
                    *self
                        .surface_shader_map
                        .get(name)
                        .unwrap_or_else(|| panic!("Unknown surface shader '{}'.", name))
                }),
                id: self.instances.len(),
                transform_indices: xforms
                    .map(|xf| (self.xforms.len(), self.xforms.len() + xf.len())),
            }
        } else {
            Instance {
                instance_type: InstanceType::Assembly,
                data_index: self.assembly_map[name],
                surface_shader_index: surface_shader_name.map(|name| {
                    *self
                        .surface_shader_map
                        .get(name)
                        .unwrap_or_else(|| panic!("Unknown surface shader '{}'.", name))
                }),
                id: self.instances.len(),
                transform_indices: xforms
                    .map(|xf| (self.xforms.len(), self.xforms.len() + xf.len())),
            }
        };

        self.instances.push(instance);

        // Store transforms
        if let Some(xf) = xforms {
            self.xforms.extend(xf);
        }
    }

    pub fn name_exists(&self, name: &str) -> bool {
        self.object_map.contains_key(name) || self.assembly_map.contains_key(name)
    }

    pub fn build(mut self) -> Assembly<'a> {
        // Calculate instance bounds, used for building object accel and light accel.
        let (bis, bbs) = self.instance_bounds();

        // Build object accel
        let object_accel = BVH4::from_objects(self.arena, &mut self.instances[..], 1, |inst| {
            &bbs[bis[inst.id]..bis[inst.id + 1]]
        });

        // Get list of instances that are for light sources or assemblies that contain light
        // sources.
        let mut light_instances: Vec<_> = self
            .instances
            .iter()
            .filter(|inst| match inst.instance_type {
                InstanceType::Object => {
                    if let Object::SurfaceLight(_) = self.objects[inst.data_index] {
                        true
                    } else {
                        false
                    }
                }

                InstanceType::Assembly => {
                    self.assemblies[inst.data_index]
                        .light_accel
                        .approximate_energy()
                        > 0.0
                }
            })
            .cloned()
            .collect();

        // Build light accel
        let light_accel = LightTree::from_objects(self.arena, &mut light_instances[..], |inst| {
            let bounds = &bbs[bis[inst.id]..bis[inst.id + 1]];
            let energy = match inst.instance_type {
                InstanceType::Object => {
                    if let Object::SurfaceLight(light) = self.objects[inst.data_index] {
                        light.approximate_energy()
                    } else {
                        0.0
                    }
                }

                InstanceType::Assembly => self.assemblies[inst.data_index]
                    .light_accel
                    .approximate_energy(),
            };
            (bounds, energy)
        });

        Assembly {
            instances: self.arena.copy_slice(&self.instances),
            light_instances: self.arena.copy_slice(&light_instances),
            xforms: self.arena.copy_slice(&self.xforms),
            surface_shaders: self.arena.copy_slice(&self.surface_shaders),
            objects: self.arena.copy_slice(&self.objects),
            assemblies: self.arena.copy_slice(&self.assemblies),
            object_accel: object_accel,
            light_accel: light_accel,
        }
    }

    /// Returns a pair of vectors with the bounds of all instances.
    /// This is used for building the assembly's BVH4.
    fn instance_bounds(&self) -> (Vec<usize>, Vec<BBox>) {
        let mut indices = vec![0];
        let mut bounds = Vec::new();

        for inst in &self.instances {
            let mut bbs = Vec::new();
            let mut bbs2 = Vec::new();

            // Get bounding boxes
            match inst.instance_type {
                InstanceType::Object => {
                    // Push bounds onto bbs
                    let obj = &self.objects[inst.data_index];
                    match *obj {
                        Object::Surface(s) => bbs.extend(s.bounds()),
                        Object::SurfaceLight(l) => bbs.extend(l.bounds()),
                    }
                }

                InstanceType::Assembly => {
                    // Push bounds onto bbs
                    let asmb = &self.assemblies[inst.data_index];
                    bbs.extend(asmb.bounds());
                }
            }

            // Transform the bounding boxes, if necessary
            if let Some((xstart, xend)) = inst.transform_indices {
                let xf = &self.xforms[xstart..xend];
                transform_bbox_slice_from(&bbs, xf, &mut bbs2);
            } else {
                bbs2.clear();
                bbs2.extend(bbs);
            }

            // Push transformed bounds onto vec
            bounds.extend(bbs2);
            indices.push(bounds.len());
        }

        (indices, bounds)
    }
}

#[derive(Copy, Clone, Debug)]
pub enum Object<'a> {
    Surface(&'a dyn Surface),
    SurfaceLight(&'a dyn SurfaceLight),
}

#[derive(Debug, Copy, Clone)]
pub struct Instance {
    pub instance_type: InstanceType,
    pub data_index: usize,
    pub surface_shader_index: Option<usize>,
    pub id: usize,
    pub transform_indices: Option<(usize, usize)>,
}

#[derive(Debug, Copy, Clone)]
pub enum InstanceType {
    Object,
    Assembly,
}