thread.rs

use std::sync::Arc;
use std::thread;
use std::time::Duration;

use arc_swap::ArcSwap;
use audionimbus::{
    AudioSettings, Context, Convolution, ConvolutionParameters, ConvolutionSettings,
    CoordinateSystem, DefaultRayTracer, DeviationModel, PathEffectParams, Pathing,
    PathingSimulationParameters, PathingSimulationSettings, Point, ProbeBatch,
    ReflectionEffectParams, Reflections, ReflectionsSharedInputs, Scene, SimulationInputs,
    SimulationParameters, SimulationSettings, SimulationSharedInputs, Simulator, Source,
    StaticMesh, StaticMeshSettings, SteamAudioError, Vector3, num_ambisonics_channels,
};
use crossbeam_channel::{Receiver, TryRecvError};
use rtrb::{Consumer, Producer};

use crate::core::mapping::EAR_HEIGHT;
use crate::core::scene2d::{Vec2, Wall};
use crate::core::shared::WorldState;
use crate::sim::builder::build_geometry;
use crate::sim::pathing::bake_probe_batch;

pub const REFLECTION_ORDER: u32 = 1;
pub const REFLECTION_CHANNELS: u32 = num_ambisonics_channels(REFLECTION_ORDER);
pub const REFLECTION_DURATION: f32 = 1.5;

pub const PATHING_ORDER: u32 = REFLECTION_ORDER;
pub const PATHING_CHANNELS: u32 = num_ambisonics_channels(PATHING_ORDER);

const NUM_VISIBILITY_SAMPLES: u32 = 16;
const PATHING_VISIBILITY_RADIUS: f32 = 0.5;
const PATHING_VISIBILITY_THRESHOLD: f32 = 0.1;
const PATHING_VISIBILITY_RANGE: f32 = 20.0;

const SIM_INTERVAL: Duration = Duration::from_millis(66);
const NUM_RAYS: u32 = 2048;
const NUM_BOUNCES: u32 = 8;
const MAX_RAYS: u32 = 4096;
const NUM_DIFFUSE_SAMPLES: u32 = 1024;
const MAX_SOURCES: u32 = 64;
const IRRADIANCE_MIN_DISTANCE: f32 = 1.0;

pub type ReflectionSnapshot = Vec<Option<ReflectionEffectParams<Convolution>>>;
pub type ReflectionProducer = Producer<ReflectionSnapshot>;
pub type ReflectionConsumer = Consumer<ReflectionSnapshot>;

pub type PathingSnapshot = Vec<Option<PathEffectParams>>;
pub type PathingProducer = Producer<PathingSnapshot>;
pub type PathingConsumer = Consumer<PathingSnapshot>;

pub struct SceneUpdate {
    pub walls: Vec<Wall>,
    pub emitter_count: usize,
    pub enclosure: bool,
}

type ReflectionSimulator = Simulator<DefaultRayTracer, (), Reflections, Pathing, Convolution>;
type ReflectionSource = Source<(), Reflections, Pathing, Convolution>;

pub fn spawn(
    context: Context,
    audio_settings: AudioSettings,
    world: Arc<ArcSwap<WorldState>>,
    updates: Receiver<SceneUpdate>,
    mut reflections: ReflectionProducer,
    mut paths: PathingProducer,
) {
    thread::spawn(move || {
        let mut state = match SimState::new(context, audio_settings) {
            Ok(state) => state,
            Err(error) => {
                eprintln!("reflections simulator failed to start: {error}");
                return;
            }
        };
        if let Err(error) = state.set_geometry(&[], false) {
            eprintln!("reflections scene init failed: {error}");
            return;
        }

        loop {
            match drain_updates(&updates) {
                Ok(Some(update)) => {
                    if let Err(error) = state.set_source_count(update.emitter_count) {
                        eprintln!("reflections source update failed: {error}");
                    }
                    if let Err(error) = state.set_geometry(&update.walls, update.enclosure) {
                        eprintln!("reflections geometry update failed: {error}");
                    }
                }
                Ok(None) => {}
                Err(()) => return,
            }

            let world_state = world.load_full();
            let (reflection_snapshot, pathing_snapshot) = state.step(&world_state);
            let _ = reflections.push(reflection_snapshot);
            let _ = paths.push(pathing_snapshot);

            thread::sleep(SIM_INTERVAL);
        }
    });
}

fn drain_updates(updates: &Receiver<SceneUpdate>) -> Result<Option<SceneUpdate>, ()> {
    let mut latest = None;
    loop {
        match updates.try_recv() {
            Ok(update) => latest = Some(update),
            Err(TryRecvError::Empty) => return Ok(latest),
            Err(TryRecvError::Disconnected) => return Err(()),
        }
    }
}

struct SimState {
    context: Context,
    simulator: ReflectionSimulator,
    scene: Scene<DefaultRayTracer>,
    sources: Vec<ReflectionSource>,
    probes: Option<ProbeBatch>,
}

impl SimState {
    fn new(context: Context, audio_settings: AudioSettings) -> Result<Self, SteamAudioError> {
        let settings = SimulationSettings::new(&audio_settings)
            .with_reflections(ConvolutionSettings {
                max_num_rays: MAX_RAYS,
                num_diffuse_samples: NUM_DIFFUSE_SAMPLES,
                max_duration: REFLECTION_DURATION,
                max_num_sources: MAX_SOURCES,
                num_threads: 1,
                max_order: REFLECTION_ORDER,
            })
            .with_pathing(PathingSimulationSettings {
                num_visibility_samples: NUM_VISIBILITY_SAMPLES,
            });
        let simulator = Simulator::try_new(&context, &settings)?;
        let scene: Scene<DefaultRayTracer> = Scene::try_new(&context)?;
        Ok(Self {
            context,
            simulator,
            scene,
            sources: Vec::new(),
            probes: None,
        })
    }

    fn set_geometry(&mut self, walls: &[Wall], enclosure: bool) -> Result<(), SteamAudioError> {
        let geometry = build_geometry(walls, enclosure);
        let mut scene: Scene<DefaultRayTracer> = Scene::try_new(&self.context)?;
        if !geometry.triangles.is_empty() {
            let mesh = StaticMesh::try_new(
                &scene,
                &StaticMeshSettings {
                    vertices: &geometry.vertices,
                    triangles: &geometry.triangles,
                    material_indices: &geometry.material_indices,
                    materials: &geometry.materials,
                },
            )?;
            scene.add_static_mesh(mesh);
        }
        scene.commit();
        self.simulator.set_scene(&scene);

        if let Some(previous) = self.probes.take() {
            self.simulator.remove_probe_batch(&previous);
        }
        let probes = bake_probe_batch(&self.context, walls);
        if let Some(batch) = &probes {
            self.simulator.add_probe_batch(batch);
        }
        self.probes = probes;

        self.simulator.commit();
        self.scene = scene;
        Ok(())
    }

    fn set_source_count(&mut self, count: usize) -> Result<(), SteamAudioError> {
        while self.sources.len() < count {
            let source = Source::try_new(&self.simulator)?;
            self.simulator.add_source(&source);
            self.sources.push(source);
        }
        while self.sources.len() > count {
            if let Some(source) = self.sources.pop() {
                self.simulator.remove_source(&source);
            }
        }
        self.simulator.commit();
        Ok(())
    }

    fn step(&self, world: &WorldState) -> (ReflectionSnapshot, PathingSnapshot) {
        if self.sources.is_empty() {
            return (Vec::new(), Vec::new());
        }
        let reflections = self.step_reflections(world);
        let paths = self.step_pathing(world);
        (reflections, paths)
    }

    fn step_reflections(&self, world: &WorldState) -> ReflectionSnapshot {
        let shared = SimulationSharedInputs::new(coordinates(world.listener)).with_reflections(
            ReflectionsSharedInputs {
                num_rays: NUM_RAYS,
                num_bounces: NUM_BOUNCES,
                duration: REFLECTION_DURATION,
                order: REFLECTION_ORDER,
                irradiance_min_distance: IRRADIANCE_MIN_DISTANCE,
            },
        );
        if self
            .simulator
            .set_shared_reflections_inputs(&shared)
            .is_err()
        {
            return std::iter::repeat_with(|| None)
                .take(self.sources.len())
                .collect();
        }

        for (index, source) in self.sources.iter().enumerate() {
            let position = world.emitters.get(index).copied().unwrap_or_default();
            let inputs = SimulationInputs {
                source: coordinates(position),
                parameters: SimulationParameters::new().with_reflections(ConvolutionParameters {
                    baked_data_identifier: None,
                }),
            };
            let _ = source.set_reflections_inputs(&inputs);
        }

        if self.simulator.run_reflections().is_err() {
            return std::iter::repeat_with(|| None)
                .take(self.sources.len())
                .collect();
        }

        self.sources
            .iter()
            .map(|source| source.get_reflections_outputs().ok())
            .collect()
    }

    fn step_pathing(&self, world: &WorldState) -> PathingSnapshot {
        let Some(probes) = &self.probes else {
            return std::iter::repeat_with(|| None)
                .take(self.sources.len())
                .collect();
        };

        let shared = SimulationSharedInputs::new(coordinates(world.listener)).with_pathing();
        if self.simulator.set_shared_pathing_inputs(&shared).is_err() {
            return std::iter::repeat_with(|| None)
                .take(self.sources.len())
                .collect();
        }

        for (index, source) in self.sources.iter().enumerate() {
            let position = world.emitters.get(index).copied().unwrap_or_default();
            let inputs = SimulationInputs {
                source: coordinates(position),
                parameters: SimulationParameters::new().with_pathing(PathingSimulationParameters {
                    pathing_probes: probes.clone(),
                    visibility_radius: PATHING_VISIBILITY_RADIUS,
                    visibility_threshold: PATHING_VISIBILITY_THRESHOLD,
                    visibility_range: PATHING_VISIBILITY_RANGE,
                    pathing_order: PATHING_ORDER,
                    enable_validation: false,
                    find_alternate_paths: false,
                    deviation: DeviationModel::Default,
                }),
            };
            let _ = source.set_pathing_inputs(&inputs);
        }

        if self.simulator.run_pathing().is_err() {
            return std::iter::repeat_with(|| None)
                .take(self.sources.len())
                .collect();
        }

        self.sources
            .iter()
            .map(|source| source.get_pathing_outputs().ok())
            .collect()
    }
}

fn coordinates(position: Vec2) -> CoordinateSystem {
    CoordinateSystem {
        right: Vector3::new(1.0, 0.0, 0.0),
        up: Vector3::new(0.0, 1.0, 0.0),
        ahead: Vector3::new(0.0, 0.0, -1.0),
        origin: Point::new(position.x, EAR_HEIGHT, position.y),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use audionimbus::{
        AmbisonicsDecodeEffect, AmbisonicsDecodeEffectParams, AmbisonicsDecodeEffectSettings,
        AudioBuffer, AudioBufferSettings, Hrtf, HrtfSettings, PathEffect, PathEffectSettings,
        ReflectionEffect, ReflectionEffectSettings, Rendering, SpeakerLayout,
    };

    use crate::core::scene2d::MaterialKind;

    fn box_walls() -> Vec<Wall> {
        let corners = [
            Vec2::new(-3.0, -3.0),
            Vec2::new(3.0, -3.0),
            Vec2::new(3.0, 3.0),
            Vec2::new(-3.0, 3.0),
        ];
        (0..4)
            .map(|i| Wall {
                a: corners[i],
                b: corners[(i + 1) % 4],
                material: MaterialKind::Concrete,
                thickness_mm: 100.0,
            })
            .collect()
    }

    #[test]
    fn enclosed_room_simulates_and_renders_reverb() {
        let audio_settings = AudioSettings {
            sampling_rate: 48_000,
            frame_size: 1024,
        };
        let context = Context::default();
        let mut state =
            SimState::new(context.clone(), audio_settings).expect("simulator should initialise");
        state.set_geometry(&box_walls(), true).expect("geometry");
        state.set_source_count(1).expect("source");

        let world = WorldState {
            listener: Vec2::new(0.0, 0.0),
            emitters: vec![Vec2::new(1.5, 1.0)],
        };
        let (reflections, paths) = state.step(&world);
        assert_eq!(reflections.len(), 1);
        assert_eq!(paths.len(), 1);
        let params = reflections[0]
            .as_ref()
            .expect("reflections should produce params");

        let mut effect = ReflectionEffect::<Convolution>::try_new(
            &context,
            &audio_settings,
            &ReflectionEffectSettings {
                impulse_response_size: (REFLECTION_DURATION * 48_000.0).ceil() as u32,
                num_channels: REFLECTION_CHANNELS,
            },
        )
        .expect("reflection effect");

        let input = vec![0.2_f32; 1024];
        let mut output = vec![0.0_f32; REFLECTION_CHANNELS as usize * 1024];
        let input_buffer = AudioBuffer::try_with_data(input.as_slice()).unwrap();
        let output_buffer = AudioBuffer::try_with_data_and_settings(
            output.as_mut_slice(),
            AudioBufferSettings::with_num_channels(REFLECTION_CHANNELS),
        )
        .unwrap();
        effect
            .apply(params, &input_buffer, &output_buffer)
            .expect("reflection apply should succeed");

        let hrtf =
            Hrtf::try_new(&context, &audio_settings, &HrtfSettings::default()).expect("hrtf");
        let mut decode = AmbisonicsDecodeEffect::try_new(
            &context,
            &audio_settings,
            &AmbisonicsDecodeEffectSettings {
                speaker_layout: SpeakerLayout::Stereo,
                hrtf: hrtf.clone(),
                max_order: REFLECTION_ORDER,
                rendering: Rendering::Binaural,
            },
        )
        .expect("decode effect");

        let mut stereo = vec![0.0_f32; 2 * 1024];
        let ambisonics_in = AudioBuffer::try_with_data_and_settings(
            output.as_slice(),
            AudioBufferSettings::with_num_channels(REFLECTION_CHANNELS),
        )
        .unwrap();
        let stereo_out = AudioBuffer::try_with_data_and_settings(
            stereo.as_mut_slice(),
            AudioBufferSettings::with_num_channels(2),
        )
        .unwrap();
        decode
            .apply(
                &AmbisonicsDecodeEffectParams {
                    order: REFLECTION_ORDER,
                    hrtf: hrtf.clone(),
                    orientation: CoordinateSystem::default(),
                },
                &ambisonics_in,
                &stereo_out,
            )
            .expect("ambisonics decode should succeed");

        let path_params = paths[0]
            .as_ref()
            .expect("pathing should find a line-of-sight path");
        let mut path_effect = PathEffect::try_new(
            &context,
            &audio_settings,
            &PathEffectSettings {
                max_order: PATHING_ORDER,
                spatialization: None,
            },
        )
        .expect("path effect");
        let mut path_ambisonics = vec![0.0_f32; PATHING_CHANNELS as usize * 1024];
        let path_in = AudioBuffer::try_with_data(input.as_slice()).unwrap();
        let path_out = AudioBuffer::try_with_data_and_settings(
            path_ambisonics.as_mut_slice(),
            AudioBufferSettings::with_num_channels(PATHING_CHANNELS),
        )
        .unwrap();
        path_effect
            .apply(path_params, &path_in, &path_out)
            .expect("path apply should succeed");
    }
}