vsrg 0.3.0

use crate::collections::sort_multi;
use std::{collections::HashMap, sync::Arc};

use crate::rhythm::{ClockTime, TempoTracks, Time};

/// Identifier of an event track: either the main track or a named track.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum TrackId {
    /// The main track.
    Main,
    /// A track identified by its name.
    Named(Arc<str>),
}

impl TrackId {
    /// Gets the name of the track, or [`None`] if this is the main track.
    pub fn as_deref(&self) -> Option<&str> {
        match self {
            TrackId::Main => None,
            TrackId::Named(name) => Some(name),
        }
    }
}

pub trait TrackEvent {
    type Cumulated: Default + Clone + PartialEq + std::fmt::Debug;

    fn time(&self) -> Time;
    fn cumulate(
        current_cumulated: Self::Cumulated,
        delta: ClockTime,
        full_duration: ClockTime,
        current: &Self,
        next: &Self,
    ) -> Self::Cumulated;
    fn initial_cumulate(first_time: ClockTime, event: &Self) -> Self::Cumulated;
}

/// Storage of event tracks keyed by Id.
#[derive(Clone, Debug)]
pub struct EventTracks<T> {
    tracks: HashMap<Arc<str>, T>,
    main_track: T,
}

impl<T> EventTracks<T> {
    /// Create an empty storage.
    pub fn new() -> Self
    where
        T: Default,
    {
        Self {
            tracks: Default::default(),
            main_track: T::default(),
        }
    }

    /// Create storage from an initial list of tracks.
    pub fn with_tracks(main: T, tracks: impl Iterator<Item = (Arc<str>, T)>) -> Self {
        let tracks = tracks.into_iter().collect();
        Self {
            tracks,
            main_track: main,
        }
    }

    /// List all tracks.
    pub fn tracks(&self) -> impl Iterator<Item = &T> {
        std::iter::chain(std::iter::once(&self.main_track), self.tracks.values())
    }

    /// List all tracks, yielding mut reference.
    pub fn tracks_mut(&mut self) -> impl Iterator<Item = &mut T> {
        std::iter::chain(
            std::iter::once(&mut self.main_track),
            self.tracks.values_mut(),
        )
    }

    /// Add a new track with a given Id.
    ///
    /// Does nothing if there's an existing track with the given Id.
    pub fn add_track(&mut self, id: Arc<str>, track: T) {
        if self.tracks.contains_key(&id) {
            return;
        }

        self.tracks.insert(id, track);
    }

    /// Remove the track with a given Id, if it exists.
    ///
    /// Returns the removed track if it exists.
    pub fn remove_track(&mut self, id: &str) -> Option<T> {
        self.tracks.remove(id)
    }

    /// Inserts or replace a track with a given Id.
    ///
    /// # Parameters
    /// - `id`: If [`TrackId::Main`], then the main track is replaced, otherwise it inserts or
    ///   replace into the given Id.
    /// - `track`: The track to insert or replace with.
    pub fn insert_or_replace_track(&mut self, id: TrackId, track: T) -> Option<T> {
        match id {
            TrackId::Named(id) => self.tracks.insert(id, track),
            TrackId::Main => Some(std::mem::replace(&mut self.main_track, track)),
        }
    }

    /// Get a mut reference to a track with the given Id.
    ///
    /// # Parameters
    /// - `id`: If [`TrackId::Main`], then the main track is returned, else it returns the track
    ///   with the given Id if exists.
    pub fn get_track_mut(&mut self, id: &TrackId) -> Option<&mut T> {
        match id.as_deref() {
            Some(id) => self.tracks.get_mut(id),
            None => Some(&mut self.main_track),
        }
    }

    /// Get a reference to a track with the given Id.
    ///
    /// # Parameters
    /// - `id`: If [`TrackId::Main`], then the main track is returned, else it returns the track
    ///   with the given Id if exists.
    pub fn get_track(&self, id: &TrackId) -> Option<&T> {
        match id.as_deref() {
            Some(id) => self.tracks.get(id),
            None => Some(&self.main_track),
        }
    }

    /// Get a reference to a track with the given Id, falling back to main track if the id is
    /// invalid.
    ///
    /// # Parameters
    /// - `id`: If [`TrackId::Main`], or if no track has the given id, then the main track is returned, else
    ///   it returns the track with the given id.
    pub fn get_track_or_main(&self, id: &TrackId) -> &T {
        match id.as_deref() {
            Some(id) => self.tracks.get(id).unwrap_or(&self.main_track),
            None => &self.main_track,
        }
    }

    /// Get a mut reference to a track with the given Id, falling back to main track if the id is
    /// invalid.
    ///
    /// # Parameters
    /// - `id`: If [`TrackId::Main`], or if no track has the given id, then the main track is returned, else
    ///   it returns the track with the given id.
    pub fn get_track_or_main_mut(&mut self, id: &TrackId) -> &mut T {
        match id.as_deref() {
            Some(id) => self.tracks.get_mut(id).unwrap_or(&mut self.main_track),
            None => &mut self.main_track,
        }
    }

    /// Get a reference to the main track.
    pub fn main_track(&self) -> &T {
        &self.main_track
    }

    /// Get a mut reference to the main track.
    pub fn main_track_mut(&mut self) -> &mut T {
        &mut self.main_track
    }
}

impl<T: Default> Default for EventTracks<T> {
    fn default() -> Self {
        Self::new()
    }
}

impl<Ev: TrackEvent> Default for EventTrack<Ev> {
    fn default() -> Self {
        Self::new(TrackId::Main)
    }
}

/// An event track which contains different events at points in time.
#[derive(Debug, Clone)]
pub struct EventTrack<Ev: TrackEvent> {
    tempo_track_id: TrackId,
    id_counter: u64,
    id_lookup: HashMap<u64, usize>,

    ids: Vec<u64>,
    time: Vec<Time>,
    clock_time: Vec<ClockTime>,
    content: Vec<Ev>,
    cumulated: Vec<Ev::Cumulated>,
}

impl<Ev: TrackEvent> EventTrack<Ev> {
    /// Create a new empty track.
    pub fn new(tempo_track_id: TrackId) -> Self {
        Self {
            tempo_track_id,
            id_counter: 0,
            id_lookup: Default::default(),
            ids: Default::default(),
            time: Default::default(),
            clock_time: Default::default(),
            content: Default::default(),
            cumulated: Default::default(),
        }
    }

    /// Create a track from a list of events with the same time unit.
    pub fn with_events(
        tempo_track_id: TrackId,
        mut events: Vec<Ev>,
        tempo_tracks: &TempoTracks,
    ) -> Self {
        let len = events.len();

        let tempo_track = tempo_tracks.get_track_or_main(&tempo_track_id);
        let mut time = events.iter().map(|ev| ev.time()).collect::<Vec<_>>();
        let mut clock_time = events
            .iter()
            .map(|ev| match ev.time() {
                Time::Clock(clock_time) => clock_time,
                Time::Beat(beat_time) => tempo_track.beat_to_clock(beat_time),
            })
            .collect::<Vec<_>>();
        sort_multi(&mut clock_time, |i, j| {
            time.swap(i, j);
            events.swap(i, j);
        });

        let mut s = Self {
            tempo_track_id,
            id_counter: len as u64,
            ids: (0..len as u64).collect(),
            time,
            clock_time,
            content: events,
            id_lookup: Default::default(),
            cumulated: Default::default(),
        };
        s.recalc_cumulative_and_lookup();
        s
    }

    /// Resizes and recalculate cumulated values, should be called after every edit to the event list.
    fn recalc_cumulative_and_lookup(&mut self) {
        self.id_lookup.clear();
        while self.cumulated.len() > self.ids.len() {
            self.cumulated.pop();
        }
        while self.cumulated.len() < self.ids.len() {
            self.cumulated.push(Ev::Cumulated::default());
        }
        if self.ids.is_empty() {
            return;
        }

        let mut cumulated = Ev::initial_cumulate(self.clock_time[0], &self.content[0]);
        self.cumulated[0] = cumulated.clone();
        self.id_lookup.insert(self.ids[0], 0);
        for i in 1..self.ids.len() {
            self.id_lookup.insert(self.ids[i], i);
            let delta = self.clock_time[i] - self.clock_time[i - 1];
            let current = &self.content[i - 1];
            let next = &self.content[i];
            cumulated = Ev::cumulate(
                cumulated.clone(),
                delta,
                self.clock_time[i] - self.clock_time[i - 1],
                current,
                next,
            );
            self.cumulated[i] = cumulated.clone();
        }
    }

    fn add_event_impl(&mut self, event: Ev, tempo_tracks: &TempoTracks) {
        let time = event.time();
        let tempo_track = tempo_tracks.get_track_or_main(&self.tempo_track_id);
        let clock_time = match time {
            Time::Clock(clock_time) => clock_time,
            Time::Beat(beat_time) => tempo_track.beat_to_clock(beat_time),
        };
        let insert_index = self.clock_time.partition_point(|x| *x <= clock_time);
        let id = self.id_counter;
        self.id_counter += 1;
        self.ids.insert(insert_index, id);
        self.time.insert(insert_index, time);
        self.clock_time.insert(insert_index, clock_time);
        self.content.insert(insert_index, event);
    }

    /// Add an event.
    pub fn add_event(&mut self, event: Ev, tempo_tracks: &TempoTracks) {
        self.add_event_impl(event, tempo_tracks);
        self.recalc_cumulative_and_lookup();
        self.health_check();
    }

    /// Add multiple events.
    pub fn add_events(&mut self, events: impl Iterator<Item = Ev>, tempo_tracks: &TempoTracks) {
        for ev in events {
            self.add_event_impl(ev, tempo_tracks);
        }
        self.recalc_cumulative_and_lookup();
        self.health_check();
    }

    /// Remove an event from the track.
    pub fn remove_event(&mut self, id: u64) {
        let Some(&index) = self.id_lookup.get(&id) else {
            return;
        };

        self.ids.remove(index);
        self.time.remove(index);
        self.clock_time.remove(index);
        self.content.remove(index);
        self.recalc_cumulative_and_lookup();
        self.health_check();
    }

    /// Remove multiple events from the track.
    pub fn remove_events(&mut self, ids: impl IntoIterator<Item = u64>) {
        let mut removing_indices = ids
            .into_iter()
            .flat_map(|id| self.id_lookup.get(&id))
            .copied()
            .collect::<Vec<_>>();
        removing_indices.sort_unstable();
        removing_indices.dedup();
        for i in removing_indices.into_iter().rev() {
            self.ids.remove(i);
            self.time.remove(i);
            self.clock_time.remove(i);
            self.content.remove(i);
        }
        self.recalc_cumulative_and_lookup();
        self.health_check();
    }

    /// Replace an event in the track.
    pub fn replace_event(&mut self, id: u64, replace_with: Ev, tempo_tracks: &TempoTracks) {
        let time = replace_with.time();
        let tempo_track = tempo_tracks.get_track_or_main(&self.tempo_track_id);
        let clock_time = match time {
            Time::Clock(clock_time) => clock_time,
            Time::Beat(beat_time) => tempo_track.beat_to_clock(beat_time),
        };

        let Some(&index) = self.id_lookup.get(&id) else {
            return;
        };

        let old_time = self.time[index];
        self.time[index] = time;
        self.clock_time[index] = clock_time;
        self.content[index] = replace_with;

        let len = self.ids.len();
        let mut j = index;
        if self.time[j] > old_time {
            while j + 1 < len && self.time[j] > self.time[j + 1] {
                self.ids.swap(j, j + 1);
                self.time.swap(j, j + 1);
                self.clock_time.swap(j, j + 1);
                self.content.swap(j, j + 1);
                j += 1;
            }
        } else {
            while j > 0 && self.time[j - 1] > self.time[j] {
                self.ids.swap(j - 1, j);
                self.time.swap(j - 1, j);
                self.clock_time.swap(j - 1, j);
                self.content.swap(j - 1, j);
                j -= 1;
            }
        }
        self.recalc_cumulative_and_lookup();
        self.health_check();
    }

    /// Replace multiple events in the track.
    pub fn replace_events(
        &mut self,
        events: impl Iterator<Item = (u64, Ev)>,
        tempo_tracks: &TempoTracks,
    ) {
        let tempo_track = tempo_tracks.get_track_or_main(&self.tempo_track_id);
        for (id, ev) in events {
            let Some(&index) = self.id_lookup.get(&id) else {
                continue;
            };
            let time = ev.time();
            self.time[index] = time;
            let clock_time = match time {
                Time::Clock(clock_time) => clock_time,
                Time::Beat(beat_time) => tempo_track.beat_to_clock(beat_time),
            };
            self.clock_time[index] = clock_time;
            self.content[index] = ev;
        }
        sort_multi(&mut self.clock_time, |i, j| {
            self.ids.swap(i, j);
            self.time.swap(i, j);
            self.content.swap(i, j);
        });
        self.recalc_cumulative_and_lookup();
        self.health_check();
    }

    /// Rebuild the event list after changes in tempo.
    ///
    /// # Remarks
    /// After changes in tempo, events with beat time may now have different clock time,
    /// requiring event tracks to reindex.
    pub fn rebuild_after_tempo_change(&mut self, tempo_tracks: &TempoTracks) {
        let tempo_track = tempo_tracks.get_track_or_main(&self.tempo_track_id);
        for i in 0..self.clock_time.len() {
            let time = self.time[i];
            let clock_time = match time {
                Time::Clock(clock_time) => clock_time,
                Time::Beat(beat_time) => tempo_track.beat_to_clock(beat_time),
            };
            self.clock_time[i] = clock_time;
        }

        sort_multi(&mut self.clock_time, |i, j| {
            self.ids.swap(i, j);
            self.time.swap(i, j);
            self.content.swap(i, j);
        });
        self.recalc_cumulative_and_lookup();
        self.health_check();
    }

    /// Change the tempo track that this event track follow then rebuild.
    pub fn change_tempo_track(&mut self, tempo_track_id: TrackId, tempo_tracks: &TempoTracks) {
        self.tempo_track_id = tempo_track_id;
        self.rebuild_after_tempo_change(tempo_tracks);
    }

    /// Check if all invariants are intact, but only in debug mode.
    ///
    /// Does nothing in release mode.
    pub(crate) fn health_check(&self) {
        debug_assert!(self.clock_time.is_sorted());

        let len = self.ids.len();
        debug_assert_eq!(self.time.len(), len);
        debug_assert_eq!(self.content.len(), len);
        debug_assert_eq!(self.cumulated.len(), len);
        debug_assert_eq!(self.id_lookup.len(), len);
        for i in 0..len {
            if i == 0 {
                debug_assert_eq!(
                    self.cumulated[0],
                    Ev::initial_cumulate(self.clock_time[0], &self.content[0])
                );
            } else {
                let delta = self.clock_time[i] - self.clock_time[i - 1];
                let expected = Ev::cumulate(
                    self.cumulated[i - 1].clone(),
                    delta,
                    self.clock_time[i] - self.clock_time[i - 1],
                    &self.content[i - 1],
                    &self.content[i],
                );
                debug_assert_eq!(self.cumulated[i], expected);
            }
        }
    }

    /// Gets the tempo track that this event track follows.
    pub fn tempo_track_id(&self) -> &TrackId {
        &self.tempo_track_id
    }

    pub fn ids(&self) -> &[u64] {
        &self.ids
    }

    pub fn times(&self) -> &[Time] {
        &self.time
    }

    pub fn clock_times(&self) -> &[ClockTime] {
        &self.clock_time
    }

    pub fn content(&self) -> &[Ev] {
        &self.content
    }

    pub fn cumulated(&self) -> &[Ev::Cumulated] {
        &self.cumulated
    }
}

#[cfg(test)]
mod test {
    use crate::{
        math::Easing,
        rhythm::{BeatTime, ClockTime, Tempo, TempoChangeEvent, TempoTrack},
    };
    use proptest::prelude::*;

    use super::*;

    #[derive(Debug, Clone)]
    struct TestEvent {
        time: Time,
    }

    impl TrackEvent for TestEvent {
        type Cumulated = ();

        fn time(&self) -> Time {
            self.time
        }

        // None of the test checks cumulated
        fn cumulate(
            _: Self::Cumulated,
            _: ClockTime,
            _: ClockTime,
            _: &Self,
            _: &Self,
        ) -> Self::Cumulated {
        }

        fn initial_cumulate(_: ClockTime, _: &Self) -> Self::Cumulated {}
    }
    type TestTrack = EventTrack<TestEvent>;

    fn test_tempo_tracks() -> TempoTracks {
        TempoTracks::with_tracks(
            crate::rhythm::TempoTrack::with_events(vec![TempoChangeEvent {
                time: Time::Clock(ClockTime::ZERO),
                tempo: crate::rhythm::Tempo::from_bpm(60.0).unwrap(),
                ease: crate::math::Easing::InConst,
            }])
            .unwrap(),
            std::iter::empty(),
        )
    }

    fn arb_event() -> impl Strategy<Value = TestEvent> {
        (any::<u16>()).prop_map(|t| TestEvent {
            time: Time::Clock(ClockTime::from_seconds(t as f64 / 1000.0).unwrap()),
        })
    }

    fn arb_events(max_len: usize) -> impl Strategy<Value = Vec<TestEvent>> {
        prop::collection::vec(arb_event(), 0..=max_len)
    }

    fn arb_track(max_len: usize) -> impl Strategy<Value = TestTrack> {
        arb_events(max_len)
            .prop_map(|v| TestTrack::with_events(TrackId::Main, v, &test_tempo_tracks()))
    }

    fn arb_track_select(max_len: usize) -> impl Strategy<Value = (TestTrack, u64)> {
        prop::collection::vec(arb_event(), 1..=max_len)
            .prop_map(|events| TestTrack::with_events(TrackId::Main, events, &test_tempo_tracks()))
            .prop_flat_map(|track| {
                let len = track.ids.len();
                let ids = track.ids.clone();
                (Just(track), (0..len).prop_map(move |index| ids[index]))
            })
    }

    fn arb_track_select_multi(
        max_len: usize,
        max_select_len: usize,
    ) -> impl Strategy<Value = (TestTrack, Vec<u64>)> {
        prop::collection::vec(arb_event(), 1..=max_len)
            .prop_map(|events| TestTrack::with_events(TrackId::Main, events, &test_tempo_tracks()))
            .prop_flat_map(move |track| {
                let len = track.ids.len();
                let ids = track.ids.clone();
                (
                    Just(track),
                    proptest::collection::hash_set(0..len, 0..=(max_select_len.min(len)))
                        .prop_map(move |indices| indices.into_iter().map(|i| ids[i]).collect()),
                )
            })
    }

    fn arb_track_replace_multi(
        max_len: usize,
        max_replace_len: usize,
    ) -> impl Strategy<Value = (TestTrack, Vec<u64>, Vec<TestEvent>)> {
        arb_track_select_multi(max_len, max_replace_len).prop_flat_map(|(track, ids)| {
            let len = ids.len();
            (
                Just(track),
                Just(ids),
                proptest::collection::vec(arb_event(), len),
            )
        })
    }

    proptest! {
        #[test]
        fn add_track_stays_consistent(mut track in arb_track(100), event in arb_event()) {
            let len = track.ids.len();
            track.add_event(event, &test_tempo_tracks());
            assert_eq!(track.ids.len(), len + 1);
        }
    }

    proptest! {
        #[test]
        fn add_multi_track_stays_consistent(mut track in arb_track(100), events in arb_events(100)) {
            let len = track.ids.len();
            let add_len = events.len();
            track.add_events(events.into_iter(), &test_tempo_tracks());
            assert_eq!(track.ids.len(), len + add_len);
        }
    }

    proptest! {
        #[test]
        fn remove_track_stays_consistent((mut track, index) in arb_track_select(100)) {
            let len = track.ids.len();
            track.remove_event(index);
            assert_eq!(track.ids.len(), len - 1);
        }
    }

    proptest! {
        #[test]
        fn remove_multi_track_stays_consistent((mut track, ids) in arb_track_select_multi(100, 50)) {
            let len = track.ids.len();
            let remove_len = ids.len();
            track.remove_events(ids);
            assert_eq!(track.ids.len(), len - remove_len);
        }
    }

    proptest! {
        #[test]
        fn replace_track_stays_consistent((mut track, index) in arb_track_select(100), repl in arb_event()) {
            let len = track.ids.len();
            track.replace_event(index, repl, &test_tempo_tracks());
            assert_eq!(track.ids.len(), len);
        }
    }

    proptest! {
        #[test]
        fn replace_multi_track_stays_consistent((mut track, ids, events) in arb_track_replace_multi(100, 50)) {
            let len = track.ids.len();
            track.replace_events(std::iter::zip(ids, events), &test_tempo_tracks());
            assert_eq!(track.ids.len(), len);
        }
    }

    #[test]
    fn change_tempo_track() {
        let tempo_tracks = TempoTracks::with_tracks(
            TempoTrack::with_events(vec![TempoChangeEvent {
                time: Time::Clock(ClockTime::ZERO),
                tempo: Tempo::from_bpm(60.0).unwrap(),
                ease: Easing::InConst,
            }])
            .unwrap(),
            vec![(
                "SecondTrack".into(),
                TempoTrack::with_events(vec![TempoChangeEvent {
                    time: Time::Clock(ClockTime::ZERO),
                    tempo: Tempo::from_bpm(120.0).unwrap(),
                    ease: Easing::InConst,
                }])
                .unwrap(),
            )]
            .into_iter(),
        );
        let mut track = TestTrack::with_events(
            TrackId::Main,
            vec![
                // 2 beat at 60bpm -> 2s
                // 2 beat at 120bpm -> 1s
                TestEvent {
                    time: Time::Beat(BeatTime::new(2.0).unwrap()),
                },
                TestEvent {
                    time: Time::Clock(ClockTime::from_seconds(1.5).unwrap()),
                },
            ],
            &tempo_tracks,
        );

        assert_eq!(track.clock_time[0], ClockTime::from_seconds(1.5).unwrap());
        assert_eq!(track.clock_time[1], ClockTime::from_seconds(2.0).unwrap());

        track.change_tempo_track(TrackId::Named("SecondTrack".into()), &tempo_tracks);

        assert_eq!(track.clock_time[0], ClockTime::from_seconds(1.0).unwrap());
        assert_eq!(track.clock_time[1], ClockTime::from_seconds(1.5).unwrap());
    }
}