blackout/src/pitch_shift.rs

use std::collections::HashMap;

use bevy::prelude::*;

use crate::{
    bevy_synthizer::{Listener, Sound},
    commands::RunIfExistsExt,
    sound::SoundIcon,
};

#[derive(Component, Clone, Copy, Debug, Default, Reflect)]
#[reflect(Component)]
struct Behind;

fn tag_behind(
    mut commands: Commands,
    config: Res<PitchShiftPlugin>,
    sounds: Query<(Entity, &GlobalTransform, Option<&Sound>, Option<&SoundIcon>)>,
    listener: Query<&GlobalTransform, With<Listener>>,
    behind: Query<Entity, With<Behind>>,
) {
    if config.downshift_behind {
        if let Ok(listener_transform) = listener.get_single() {
            let listener_forward = listener_transform.forward();
            for (entity, transform, sound, icon) in sounds.iter() {
                if icon.is_none() && sound.is_none() {
                    continue;
                }
                let v = transform.translation() - listener_transform.translation();
                let dot = v.dot(listener_forward);
                let is_behind = dot <= 0.;
                if is_behind {
                    commands.run_if_exists(entity, |mut entity| {
                        entity.insert(Behind);
                    });
                } else if !is_behind {
                    commands.run_if_exists(entity, |mut entity| {
                        entity.remove::<Behind>();
                    });
                }
            }
        }
    } else {
        for entity in behind.iter() {
            commands.run_if_exists(entity, |mut entity| {
                entity.remove::<Behind>();
            });
        }
    }
}

#[derive(Resource, Clone, Debug, Default, Deref, DerefMut)]
struct LastIconPitch(HashMap<Entity, f64>);

#[derive(Resource, Clone, Debug, Default, Deref, DerefMut)]
struct LastSoundPitch(HashMap<Entity, f64>);

fn behind_added(
    config: Res<PitchShiftPlugin>,
    mut last_icon_pitch: ResMut<LastIconPitch>,
    mut last_sound_pitch: ResMut<LastSoundPitch>,
    mut query: Query<(Entity, Option<&mut SoundIcon>, Option<&mut Sound>), Added<Behind>>,
) {
    for (entity, icon, sound) in query.iter_mut() {
        if let Some(mut icon) = icon {
            icon.pitch *= config.downshift;
            last_icon_pitch.insert(entity, icon.pitch);
        } else if let Some(mut sound) = sound {
            sound.pitch *= config.downshift;
            last_sound_pitch.insert(entity, sound.pitch);
        }
    }
}

fn behind_removed(
    config: Res<PitchShiftPlugin>,
    mut last_icon_pitch: ResMut<LastIconPitch>,
    mut last_sound_pitch: ResMut<LastSoundPitch>,
    mut removed: RemovedComponents<Behind>,
    mut icons: Query<&mut SoundIcon>,
    mut sounds: Query<&mut Sound>,
) {
    let downshift = 1. / config.downshift;
    for entity in &mut removed {
        if let Ok(mut icon) = icons.get_mut(entity) {
            icon.pitch *= downshift;
            last_icon_pitch.remove(&entity);
        } else if let Ok(mut sound) = sounds.get_mut(entity) {
            sound.pitch *= downshift;
            last_sound_pitch.remove(&entity);
        }
    }
}

fn sound_icon_changed(
    config: Res<PitchShiftPlugin>,
    mut last_icon_pitch: ResMut<LastIconPitch>,
    mut icons: Query<(Entity, &mut SoundIcon), (With<Behind>, Changed<SoundIcon>)>,
) {
    for (entity, mut icon) in icons.iter_mut() {
        let should_change = if let Some(pitch) = last_icon_pitch.get(&entity) {
            *pitch != icon.pitch
        } else {
            false
        };
        if should_change {
            icon.pitch *= config.downshift;
            last_icon_pitch.insert(entity, icon.pitch);
        }
    }
}

fn sound_changed(
    config: Res<PitchShiftPlugin>,
    mut last_sound_pitch: ResMut<LastSoundPitch>,
    mut sounds: Query<(Entity, &mut Sound), (With<Behind>, Without<SoundIcon>, Changed<Sound>)>,
) {
    for (entity, mut sound) in sounds.iter_mut() {
        let should_change = if let Some(pitch) = last_sound_pitch.get(&entity) {
            *pitch != sound.pitch
        } else {
            false
        };
        if should_change {
            sound.pitch *= config.downshift;
            last_sound_pitch.insert(entity, sound.pitch);
        }
    }
}

fn sync_config(
    mut commands: Commands,
    config: Res<PitchShiftPlugin>,
    behind: Query<Entity, With<Behind>>,
) {
    if config.is_changed() {
        for entity in behind.iter() {
            commands.entity(entity).remove::<Behind>();
        }
    }
}

#[derive(Resource, Clone, Copy, Debug)]
pub struct PitchShiftPlugin {
    pub downshift_behind: bool,
    pub downshift: f64,
}

impl Default for PitchShiftPlugin {
    fn default() -> Self {
        Self {
            downshift_behind: false,
            downshift: 0.95,
        }
    }
}

impl Plugin for PitchShiftPlugin {
    fn build(&self, app: &mut App) {
        app.insert_resource(*self)
            .register_type::<Behind>()
            .init_resource::<LastIconPitch>()
            .init_resource::<LastSoundPitch>()
            .add_system(tag_behind.in_base_set(CoreSet::PreUpdate))
            .add_systems((behind_added, sound_icon_changed, sound_changed, sync_config))
            .add_system(behind_removed.in_base_set(CoreSet::PostUpdate));
    }
}