vsrg 0.3.0

use super::*;
use crate::collections::sort_multi;
use crate::notes::{NoteGroup, NoteGroups};
use crate::rhythm::{ScrollSpeedTracks, TempoTracks};

impl<Note: ShortNoteData> Default for ShortNoteTypeStorage<Note> {
    fn default() -> Self {
        Self::new()
    }
}

impl<Note: ShortNoteData> ShortNoteTypeStorage<Note> {
    /// Create a new empty storage.
    pub fn new() -> Self {
        Self {
            note_id_counter: 0,
            note_id_lookup: HashMap::new(),
            time_sorted_groups: vec![],
            scroll_sorted_groups: vec![],
        }
    }

    /// Create a new storage populated with notes.
    pub fn with_notes<G: NoteGroup>(
        note_groups: &mut NoteGroups<G>,
        tempo_tracks: &mut TempoTracks,
        scroll_tracks: &mut ScrollSpeedTracks,
        notes: impl Iterator<Item = Note>,
    ) -> Self {
        let mut storage = Self::new();
        storage.add_notes(note_groups, tempo_tracks, scroll_tracks, notes);
        storage
    }

    pub(super) fn advance_to_next_id(&mut self) -> u64 {
        loop {
            self.note_id_counter = self.note_id_counter.wrapping_add(1);
            if self.note_id_counter != 0 && !self.note_id_lookup.contains_key(&self.note_id_counter)
            {
                return self.note_id_counter;
            }
        }
    }

    fn insert<G: NoteGroup>(
        &mut self,
        groups: &NoteGroups<G>,
        tempos: &TempoTracks,
        scrolls: &ScrollSpeedTracks,
        id: u64,
        note: Note,
        runtime: Note::RuntimeData,
    ) {
        let (group_id, scroll_id, tempo_id) = get_ids(&note, groups);
        let time = note.time();
        let clock = to_clock_time(tempos, time, tempo_id);
        let scroll_track = scrolls.get_track_or_main(scroll_id);
        let scroll = scroll_track.calculate_scroll_position(clock);
        let mut runtime = runtime;
        Note::apply_event_track(
            &mut runtime,
            tempos.get_track_or_main(tempo_id),
            scroll_track,
        );
        let g = self
            .time_sorted_groups
            .iter()
            .position(|g| &g.group_id == group_id)
            .unwrap_or_else(|| {
                self.time_sorted_groups
                    .push(ShortNoteGroup::new(group_id.clone()));
                self.scroll_sorted_groups
                    .push(ShortNoteGroup::new(group_id.clone()));
                self.time_sorted_groups.len() - 1
            });
        let value = note.value_data();
        for group in [
            &mut self.time_sorted_groups[g],
            &mut self.scroll_sorted_groups[g],
        ] {
            group.note_id.push(id);
            group.time.push(time);
            group.clock_time.push(clock);
            group.scroll_position.push(scroll);
            group.value_data.push(value.clone());
            group.runtime_data.push(runtime.clone());
        }
    }

    /// Add a note to the storage.
    ///
    /// # Returns
    /// The id of the added note.
    pub fn add_note<G: NoteGroup>(
        &mut self,
        groups: &mut NoteGroups<G>,
        tempos: &mut TempoTracks,
        scrolls: &mut ScrollSpeedTracks,
        note: Note,
    ) -> u64 {
        self.add_notes(groups, tempos, scrolls, std::iter::once(note))[0]
    }

    /// Add multiple notes to the storage.
    ///
    /// # Returns
    /// The list of ids of added notes.
    pub fn add_notes<G: NoteGroup>(
        &mut self,
        groups: &mut NoteGroups<G>,
        tempos: &mut TempoTracks,
        scrolls: &mut ScrollSpeedTracks,
        notes: impl Iterator<Item = Note>,
    ) -> Vec<u64> {
        let mut ids = vec![];
        for note in notes {
            let id = self.advance_to_next_id();
            self.insert(
                groups,
                tempos,
                scrolls,
                id,
                note,
                Note::RuntimeData::default(),
            );
            ids.push(id);
        }
        self.sort_and_link();
        self.health_check();
        ids
    }

    /// Remove a note from the storage, identified by its id.
    pub fn remove_note(&mut self, id: u64) {
        self.remove_notes(std::iter::once(id));
    }

    /// Remove multiple notes from the storage, identified by their ids.
    pub fn remove_notes(&mut self, ids: impl Iterator<Item = u64>) {
        let ids: std::collections::HashSet<_> = ids.collect();
        for group in self
            .time_sorted_groups
            .iter_mut()
            .chain(self.scroll_sorted_groups.iter_mut())
        {
            for i in (0..group.note_id.len()).rev() {
                if ids.contains(&group.note_id[i]) {
                    group.remove(i);
                }
            }
        }
        self.relink();
        self.health_check();
    }

    /// Replace a note while preserving its id and runtime state.
    pub fn replace_note<G: NoteGroup>(
        &mut self,
        groups: &mut NoteGroups<G>,
        tempos: &mut TempoTracks,
        scrolls: &mut ScrollSpeedTracks,
        id: u64,
        note: Note,
    ) {
        self.replace_notes(groups, tempos, scrolls, std::iter::once((id, note)));
    }

    /// Replace multiple notes while preserving their ids and runtime state.
    ///
    /// Replacement with unknown ids are skipped
    /// If multiple replacements repeat the same id, then the last replacement takes effect.
    pub fn replace_notes<G: NoteGroup>(
        &mut self,
        groups: &mut NoteGroups<G>,
        tempos: &mut TempoTracks,
        scrolls: &mut ScrollSpeedTracks,
        notes: impl Iterator<Item = (u64, Note)>,
    ) {
        for (id, note) in notes {
            let Some(loc) = self.note_id_lookup.get(&id).copied() else {
                continue;
            };
            let runtime =
                self.time_sorted_groups[loc.group_index].remove(loc.index_in_time_sorted_group);
            self.scroll_sorted_groups[loc.group_index].remove(loc.index_in_scroll_sorted_group);
            self.insert(groups, tempos, scrolls, id, note, runtime);
            self.relink();
        }
        self.sort_and_link();
        self.health_check();
    }

    /// Recalculate notes clock time reindex after edits to tempo & scroll tracks.
    pub fn rebuild_clock_time<G: NoteGroup>(
        &mut self,
        groups: &NoteGroups<G>,
        tempos: &TempoTracks,
        scrolls: &ScrollSpeedTracks,
    ) {
        for group in &mut self.time_sorted_groups {
            let ng = groups.get_group_or_main(&group.group_id);
            let tempo_track = tempos.get_track_or_main(ng.tempo_track_id());
            let scroll_track = scrolls.get_track_or_main(ng.scroll_track_id());
            for i in 0..group.note_id.len() {
                group.clock_time[i] = match group.time[i] {
                    Time::Clock(clock_time) => clock_time,
                    Time::Beat(beat_time) => tempo_track.beat_to_clock(beat_time),
                };
                group.scroll_position[i] =
                    scroll_track.calculate_scroll_position(group.clock_time[i]);
                apply_event_track_at(group, i, tempo_track, scroll_track);
            }
        }
        for group in &mut self.scroll_sorted_groups {
            let ng = groups.get_group_or_main(&group.group_id);
            let tempo_track = tempos.get_track_or_main(ng.tempo_track_id());
            let scroll_track = scrolls.get_track_or_main(ng.scroll_track_id());
            for i in 0..group.note_id.len() {
                group.clock_time[i] = match group.time[i] {
                    Time::Clock(clock_time) => clock_time,
                    Time::Beat(beat_time) => tempo_track.beat_to_clock(beat_time),
                };
                group.scroll_position[i] =
                    scroll_track.calculate_scroll_position(group.clock_time[i]);
                apply_event_track_at(group, i, tempo_track, scroll_track);
            }
        }
        self.sort_and_link();
        self.health_check();
    }

    /// Recalculate notes clock time reindex after edits to tempo & scroll tracks.
    pub fn rebuild_scroll_pos<G: NoteGroup>(
        &mut self,
        groups: &NoteGroups<G>,
        tempos: &TempoTracks,
        scrolls: &ScrollSpeedTracks,
    ) {
        for group in &mut self.time_sorted_groups {
            let ng = groups.get_group_or_main(&group.group_id);
            let tempo_track = tempos.get_track_or_main(ng.tempo_track_id());
            let scroll_track = scrolls.get_track_or_main(ng.scroll_track_id());
            for i in 0..group.note_id.len() {
                group.scroll_position[i] =
                    scroll_track.calculate_scroll_position(group.clock_time[i]);
                apply_event_track_at(group, i, tempo_track, scroll_track);
            }
        }
        for group in &mut self.scroll_sorted_groups {
            let ng = groups.get_group_or_main(&group.group_id);
            let tempo_track = tempos.get_track_or_main(ng.tempo_track_id());
            let scroll_track = scrolls.get_track_or_main(ng.scroll_track_id());
            for i in 0..group.note_id.len() {
                group.scroll_position[i] =
                    scroll_track.calculate_scroll_position(group.clock_time[i]);
                apply_event_track_at(group, i, tempo_track, scroll_track);
            }
        }
        self.sort_and_link();
        self.health_check();
    }

    pub(super) fn sort_and_link(&mut self) {
        for group in &mut self.time_sorted_groups {
            group.sort(false);
        }
        for group in &mut self.scroll_sorted_groups {
            group.sort(true);
        }
        self.relink();
    }

    pub(super) fn relink(&mut self) {
        self.time_sorted_groups.retain(|g| !g.note_id.is_empty());
        self.scroll_sorted_groups.retain(|g| !g.note_id.is_empty());
        self.note_id_lookup.clear();
        for (g, group) in self.time_sorted_groups.iter().enumerate() {
            for (i, &id) in group.note_id.iter().enumerate() {
                self.note_id_lookup.insert(
                    id,
                    NoteLocation {
                        group_index: g,
                        index_in_time_sorted_group: i,
                        index_in_scroll_sorted_group: 0,
                    },
                );
            }
        }
        for group in &self.scroll_sorted_groups {
            for (i, id) in group.note_id.iter().enumerate() {
                self.note_id_lookup
                    .get_mut(id)
                    .unwrap()
                    .index_in_scroll_sorted_group = i;
            }
        }
    }

    pub fn remove_group(&mut self, group_id: &GroupId, mut handle_removed: impl FnMut(Note)) {
        let ids = self
            .time_sorted_groups
            .iter()
            .filter(|g| &g.group_id == group_id)
            .flat_map(|g| g.note_id.iter().copied())
            .collect::<Vec<_>>();
        self.extract_notes(ids.into_iter(), |_, note| handle_removed(note));
    }

    pub fn extract_notes(
        &mut self,
        ids: impl Iterator<Item = u64> + Clone,
        mut handle_removed: impl FnMut(u64, Note),
    ) {
        let mut seen = std::collections::HashSet::new();
        for (id, note) in self
            .query_notes(ids.clone())
            .filter(|(id, _)| seen.insert(*id))
        {
            handle_removed(id, note);
        }
        self.remove_notes(ids);
    }

    pub(crate) fn health_check(&self) {
        debug_assert_eq!(
            self.time_sorted_groups.len(),
            self.scroll_sorted_groups.len()
        );
        debug_assert_eq!(
            self.note_id_lookup.len(),
            self.time_sorted_groups
                .iter()
                .map(|g| g.note_id.len())
                .sum::<usize>()
        );
        for (g, (time, scroll)) in self
            .time_sorted_groups
            .iter()
            .zip(&self.scroll_sorted_groups)
            .enumerate()
        {
            debug_assert_eq!(time.group_id, scroll.group_id);
            debug_assert!(time.clock_time.is_sorted());
            debug_assert!(scroll.scroll_position.is_sorted());
            debug_assert_eq!(time.note_id.len(), scroll.note_id.len());
            for group in [time, scroll] {
                let len = group.note_id.len();
                debug_assert!(len > 0);
                debug_assert_eq!(group.time.len(), len);
                debug_assert_eq!(group.clock_time.len(), len);
                debug_assert_eq!(group.scroll_position.len(), len);
                debug_assert_eq!(group.value_data.len(), len);
                debug_assert_eq!(group.runtime_data.len(), len);
            }
            for (i, id) in time.note_id.iter().enumerate() {
                let loc = &self.note_id_lookup[id];
                let j = loc.index_in_scroll_sorted_group;
                debug_assert_eq!(loc.group_index, g);
                debug_assert_eq!(loc.index_in_time_sorted_group, i);
                debug_assert_eq!(scroll.note_id[j], *id);
                debug_assert_eq!(time.time[i], scroll.time[j]);
                debug_assert_eq!(time.clock_time[i], scroll.clock_time[j]);
                debug_assert_eq!(time.scroll_position[i], scroll.scroll_position[j]);
            }
        }
    }
}

impl<Note: ShortNoteData> ShortNoteGroup<Note> {
    fn remove(&mut self, i: usize) -> Note::RuntimeData {
        self.note_id.remove(i);
        self.time.remove(i);
        self.clock_time.remove(i);
        self.scroll_position.remove(i);
        self.value_data.remove(i);
        self.runtime_data.remove(i)
    }

    fn sort(&mut self, scroll: bool) {
        let mut order: Vec<_> = (0..self.note_id.len()).collect();
        if scroll {
            order.sort_by_key(|&i| (self.scroll_position[i], self.note_id[i]));
        }
        // Sort a permutation, then apply the same swaps to every component.
        let mut ranks = vec![0; order.len()];
        if !scroll {
            order.sort_by_key(|&i| (self.clock_time[i], self.note_id[i]));
        }
        for (rank, i) in order.into_iter().enumerate() {
            ranks[i] = rank;
        }
        sort_multi(&mut ranks, |i, j| {
            self.note_id.swap(i, j);
            self.time.swap(i, j);
            self.clock_time.swap(i, j);
            self.scroll_position.swap(i, j);
            self.value_data.swap(i, j);
            self.runtime_data.swap(i, j);
        });
    }
}

fn get_ids<'a, Note: ShortNoteData, G: NoteGroup>(
    note: &'a Note,
    note_groups: &'a NoteGroups<G>,
) -> (&'a GroupId, &'a TrackId, &'a TrackId) {
    let group_id = note.group_id();
    let group = note_groups
        .get_group(group_id)
        .unwrap_or(note_groups.main_group());
    let scroll_track_id = group.scroll_track_id();
    let tempo_track_id = group.tempo_track_id();
    (group_id, scroll_track_id, tempo_track_id)
}

fn apply_event_track_at<Note: ShortNoteData>(
    group: &mut ShortNoteGroup<Note>,
    index: usize,
    tempo_track: &crate::rhythm::TempoTrack,
    scroll_track: &crate::rhythm::ScrollSpeedTrack,
) {
    let mut runtime = group.runtime_data.remove(index);
    Note::apply_event_track(&mut runtime, tempo_track, scroll_track);
    group.runtime_data.push(runtime);
    let last = group.runtime_data.len() - 1;
    group.runtime_data.swap(index, last);
}

fn to_clock_time(tempo_tracks: &TempoTracks, time: Time, tempo_track_id: &TrackId) -> ClockTime {
    match time {
        Time::Clock(clock_time) => clock_time,
        Time::Beat(beat_time) => tempo_tracks
            .get_track_or_main(tempo_track_id)
            .beat_to_clock(beat_time),
    }
}

pub(super) fn soa_set_element<T: Soars>(soa: &mut Soa<T>, index: usize, data: T) {
    let len = soa.len();
    soa.swap_remove(index);
    soa.push(data);
    soa.swap(index, len - 1);
}