ommx 3.0.0-beta.4

mod approx;
mod parse;

use anyhow::{bail, Result};
use derive_more::{Deref, From};
use fnv::{FnvHashMap, FnvHashSet};
use std::{collections::BTreeSet, hash::Hash};

#[repr(transparent)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, From, Deref)]
pub struct SampleID(u64);

impl SampleID {
    pub fn into_inner(self) -> u64 {
        self.0
    }
}

pub type SampleIDSet = BTreeSet<SampleID>;

#[derive(Debug, Clone)]
pub struct Sampled<T> {
    offsets: FnvHashMap<SampleID, usize>,
    data: Vec<T>,
}

impl<T> Default for Sampled<T> {
    fn default() -> Self {
        Self {
            offsets: FnvHashMap::default(),
            data: Vec::new(),
        }
    }
}

impl<T> From<T> for Sampled<T> {
    fn from(value: T) -> Self {
        let mut offsets = FnvHashMap::default();
        offsets.insert(SampleID(0), 0);
        Self {
            offsets,
            data: vec![value],
        }
    }
}

impl<T> From<(SampleID, T)> for Sampled<T> {
    fn from((id, value): (SampleID, T)) -> Self {
        let mut offsets = FnvHashMap::default();
        offsets.insert(id, 0);
        Self {
            offsets,
            data: vec![value],
        }
    }
}

/// A sample ID already present in a collection or repeated in one append input.
///
/// Use [`Self::id`] to identify the conflicting ID before correcting the input
/// and retrying the atomic operation.
///
/// ```
/// let mut sampled = ommx::Sampled::from((ommx::SampleID::from(0), "existing"));
/// let error = sampled
///     .append([ommx::SampleID::from(0)], "replacement")
///     .unwrap_err();
/// assert_eq!(error.id(), ommx::SampleID::from(0));
/// ```
#[derive(Debug, thiserror::Error)]
#[error("Duplicated sample ID: {id:?}")]
pub struct DuplicatedSampleIDError {
    id: SampleID,
}

impl DuplicatedSampleIDError {
    /// Return the sample ID that caused the duplicate-ID failure.
    pub fn id(&self) -> SampleID {
        self.id
    }
}

impl<T> Sampled<T> {
    pub fn constants(ids: impl Iterator<Item = SampleID>, value: T) -> Self {
        let map = ids.map(|id| (id, 0)).collect();
        let data = vec![value];
        Self { offsets: map, data }
    }

    /// Append one value under every supplied sample ID.
    ///
    /// # Errors
    ///
    /// Returns [`DuplicatedSampleIDError`] if an ID already exists in this
    /// collection or occurs more than once in `ids`. All IDs are validated
    /// before mutation, so the collection is unchanged when this method fails.
    pub fn append(
        &mut self,
        ids: impl IntoIterator<Item = SampleID>,
        value: T,
    ) -> std::result::Result<(), DuplicatedSampleIDError> {
        let mut new_ids = Vec::new();
        let mut seen = FnvHashSet::default();
        for id in ids {
            if self.offsets.contains_key(&id) || !seen.insert(id) {
                return Err(DuplicatedSampleIDError { id });
            }
            new_ids.push(id);
        }

        let offset = self.data.len();
        self.data.push(value);
        for id in new_ids {
            self.offsets.insert(id, offset);
        }
        Ok(())
    }

    pub fn new<Iter, Inner>(ids: Iter, data: impl IntoIterator<Item = T>) -> Result<Self>
    where
        Iter: IntoIterator<Item = Inner>,
        Inner: IntoIterator<Item = SampleID>,
    {
        let mut out = Self::default();
        let mut ids_iter = ids.into_iter();
        let mut data_iter = data.into_iter();
        loop {
            match (ids_iter.next(), data_iter.next()) {
                (Some(ids), Some(data)) => out.append(ids, data)?,
                (None, None) => break,
                (Some(_), None) => bail!("Data length mismatch"),
                (None, Some(_)) => bail!("Sample IDs length mismatch"),
            }
        }
        Ok(out)
    }

    pub fn new_no_dedup(iter: impl Iterator<Item = (SampleID, T)>) -> Self {
        let mut offsets = FnvHashMap::default();
        let mut data = Vec::new();
        for (n, (id, value)) in iter.enumerate() {
            offsets.insert(id, n);
            data.push(value);
        }
        Self { offsets, data }
    }

    pub fn new_dedup<I>(iter: I) -> Self
    where
        I: Iterator<Item = (SampleID, T)>,
        T: Hash + Eq + Clone,
    {
        let mut offsets = FnvHashMap::default();
        let mut data = Vec::new();
        let mut value_to_offset: FnvHashMap<T, usize> = FnvHashMap::default();

        for (id, value) in iter {
            // Check if we already have this value using HashMap lookup (O(1))
            let offset = match value_to_offset.get(&value) {
                Some(&existing_offset) => {
                    // Reuse existing data
                    existing_offset
                }
                None => {
                    // Add new data
                    let new_offset = data.len();
                    value_to_offset.insert(value.clone(), new_offset);
                    data.push(value);
                    new_offset
                }
            };
            offsets.insert(id, offset);
        }

        Self { offsets, data }
    }

    pub fn iter(&self) -> impl Iterator<Item = (&SampleID, &T)> {
        self.offsets.iter().map(move |(id, offset)| {
            debug_assert!(*offset < self.data.len());
            (id, &self.data[*offset])
        })
    }

    /// Mutable iterator over the *unique* stored values.
    ///
    /// Each distinct `T` is yielded once, even when multiple [`SampleID`]s
    /// point at it. Mutating the value therefore mutates it for every sample
    /// that references it.
    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut T> {
        self.data.iter_mut()
    }

    pub fn ids(&self) -> SampleIDSet {
        self.offsets.keys().copied().collect()
    }

    pub fn has_same_ids(&self, ids: &SampleIDSet) -> bool {
        if self.offsets.len() != ids.len() {
            return false;
        }
        // Check that all IDs in the set are present in our offsets
        ids.iter().all(|id| self.offsets.contains_key(id))
    }

    pub fn map<U, F: FnMut(T) -> U>(self, f: F) -> Sampled<U> {
        Sampled {
            offsets: self.offsets,
            data: self.data.into_iter().map(f).collect(),
        }
    }

    /// Non-consuming, fallible variant of [`Self::map`].
    ///
    /// Applies `f` to each unique stored value once; sample-id grouping is
    /// preserved. Useful when evaluating per-sample-state like in
    /// `Evaluate::evaluate_samples`.
    pub fn try_map_ref<U, E>(
        &self,
        mut f: impl FnMut(&T) -> std::result::Result<U, E>,
    ) -> std::result::Result<Sampled<U>, E> {
        let data = self
            .data
            .iter()
            .map(&mut f)
            .collect::<std::result::Result<Vec<_>, E>>()?;
        Ok(Sampled {
            offsets: self.offsets.clone(),
            data,
        })
    }

    pub fn num_samples(&self) -> usize {
        self.offsets.len()
    }

    /// Get a reference to the value for a specific sample ID.
    ///
    /// Returns [`None`] if the sample ID is not known to this [`Sampled`].
    pub fn get(&self, sample_id: SampleID) -> Option<&T> {
        self.offsets.get(&sample_id).map(|&offset| {
            debug_assert!(offset < self.data.len());
            &self.data[offset]
        })
    }

    /// Gather up the sample ID for each sample.
    pub fn chunk(self) -> Vec<(T, FnvHashSet<SampleID>)> {
        let mut out = self
            .data
            .into_iter()
            .map(|data| (data, FnvHashSet::default()))
            .collect::<Vec<_>>();
        for (id, offset) in &self.offsets {
            out[*offset].1.insert(*id);
        }
        out
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_sampled() {
        let sampled = Sampled::new(
            [[SampleID(1), SampleID(2)], [SampleID(5), SampleID(7)]],
            [1, 2],
        )
        .unwrap();
        assert_eq!(sampled.num_samples(), 4);
        assert_eq!(
            sampled.iter().collect::<Vec<_>>(),
            vec![
                (&SampleID(5), &2),
                (&SampleID(7), &2),
                (&SampleID(1), &1),
                (&SampleID(2), &1),
            ]
        );

        // Size mismatch tests
        assert!(Sampled::new(
            [[SampleID(1), SampleID(2)], [SampleID(5), SampleID(7)]],
            [1, 2, 3],
        )
        .is_err());
        assert!(Sampled::new(
            [[SampleID(1), SampleID(2)], [SampleID(5), SampleID(7)]],
            [1],
        )
        .is_err());
    }

    #[test]
    fn test_sampled_get() {
        let sampled = Sampled::new(
            [[SampleID(1), SampleID(2)], [SampleID(5), SampleID(7)]],
            [10, 20],
        )
        .unwrap();

        // Test successful get
        assert_eq!(sampled.get(SampleID(1)).unwrap(), &10);
        assert_eq!(sampled.get(SampleID(2)).unwrap(), &10);
        assert_eq!(sampled.get(SampleID(5)).unwrap(), &20);
        assert_eq!(sampled.get(SampleID(7)).unwrap(), &20);

        // Test get with unknown sample ID
        assert!(sampled.get(SampleID(999)).is_none());
    }

    #[test]
    fn test_new_dedup() {
        let sampled = Sampled::new_dedup(
            [
                (SampleID(1), 10),
                (SampleID(2), 20),
                (SampleID(3), 10), // Same value as SampleID(1)
                (SampleID(4), 30),
                (SampleID(5), 20), // Same value as SampleID(2)
            ]
            .into_iter(),
        );

        // Should have 5 samples but only 3 data entries (deduplication occurred)
        assert_eq!(sampled.num_samples(), 5);
        assert_eq!(sampled.data.len(), 3); // Only 3 data entries stored due to deduplication

        // Test that same values point to the same data
        assert_eq!(sampled.get(SampleID(1)).unwrap(), &10);
        assert_eq!(sampled.get(SampleID(3)).unwrap(), &10); // Same value
        assert_eq!(sampled.get(SampleID(2)).unwrap(), &20);
        assert_eq!(sampled.get(SampleID(5)).unwrap(), &20); // Same value
        assert_eq!(sampled.get(SampleID(4)).unwrap(), &30);

        // Verify that samples with same values share the same offset
        let offset_1 = sampled.offsets[&SampleID(1)];
        let offset_3 = sampled.offsets[&SampleID(3)];
        assert_eq!(offset_1, offset_3); // Should point to same data

        let offset_2 = sampled.offsets[&SampleID(2)];
        let offset_5 = sampled.offsets[&SampleID(5)];
        assert_eq!(offset_2, offset_5); // Should point to same data
    }

    #[test]
    fn test_iter_mut_shared_value_propagation() {
        // Two sample IDs share one stored value via append([a, b], value).
        // Mutating through iter_mut must be visible from both IDs, and the
        // offsets table must not be disturbed.
        let mut sampled = Sampled::<Vec<i32>>::default();
        sampled
            .append([SampleID(1), SampleID(2)], vec![10])
            .unwrap();
        sampled.append([SampleID(3)], vec![20]).unwrap();

        // Distinct values only: iter_mut yields one reference per stored entry,
        // not per sample id.
        assert_eq!(sampled.iter_mut().count(), 2);

        for v in sampled.iter_mut() {
            v.push(99);
        }

        // Both shared IDs see the mutation through the same storage slot.
        assert_eq!(sampled.get(SampleID(1)).unwrap(), &vec![10, 99]);
        assert_eq!(sampled.get(SampleID(2)).unwrap(), &vec![10, 99]);
        assert_eq!(sampled.get(SampleID(3)).unwrap(), &vec![20, 99]);

        // Offsets remain stable: shared IDs still point at the same slot.
        assert_eq!(sampled.offsets[&SampleID(1)], sampled.offsets[&SampleID(2)]);
        assert_ne!(sampled.offsets[&SampleID(1)], sampled.offsets[&SampleID(3)]);
        assert_eq!(sampled.num_samples(), 3);
    }

    #[test]
    fn append_is_atomic_when_an_existing_id_follows_a_fresh_id() {
        let mut sampled = Sampled::from((SampleID(0), 10));

        let error = sampled.append([SampleID(1), SampleID(0)], 20).unwrap_err();

        assert_eq!(error.id(), SampleID(0));
        assert_eq!(sampled.ids(), BTreeSet::from([SampleID(0)]));
        assert_eq!(sampled.get(SampleID(0)), Some(&10));
        assert_eq!(sampled.get(SampleID(1)), None);
        assert_eq!(sampled.data, vec![10]);
    }

    #[test]
    fn append_is_atomic_for_duplicate_ids_in_one_input() {
        let mut sampled = Sampled::from((SampleID(0), 10));

        let error = sampled.append([SampleID(1), SampleID(1)], 20).unwrap_err();

        assert_eq!(error.id(), SampleID(1));
        assert_eq!(sampled.ids(), BTreeSet::from([SampleID(0)]));
        assert_eq!(sampled.get(SampleID(0)), Some(&10));
        assert_eq!(sampled.get(SampleID(1)), None);
        assert_eq!(sampled.data, vec![10]);
    }

    #[test]
    fn test_try_map_ref_preserves_offsets_and_propagates_errors() {
        let sampled = Sampled::new(
            [[SampleID(1), SampleID(2)], [SampleID(5), SampleID(7)]],
            [10, 20],
        )
        .unwrap();

        // Happy path: mapping preserves the ID → offset grouping exactly, so
        // IDs that shared storage before still share the mapped value.
        let mapped = sampled.try_map_ref(|v| anyhow::Ok(*v + 1)).unwrap();
        assert_eq!(mapped.num_samples(), 4);
        assert_eq!(mapped.get(SampleID(1)).unwrap(), &11);
        assert_eq!(mapped.get(SampleID(2)).unwrap(), &11);
        assert_eq!(mapped.get(SampleID(5)).unwrap(), &21);
        assert_eq!(mapped.get(SampleID(7)).unwrap(), &21);
        assert_eq!(mapped.offsets[&SampleID(1)], mapped.offsets[&SampleID(2)]);
        assert_eq!(mapped.offsets[&SampleID(5)], mapped.offsets[&SampleID(7)]);

        // Errors short-circuit: a failure on any stored value surfaces to the
        // caller instead of producing a partial Sampled<U>.
        let err = sampled
            .try_map_ref(|v| -> Result<i32> {
                if *v == 20 {
                    anyhow::bail!("boom");
                }
                Ok(*v)
            })
            .unwrap_err();
        assert_eq!(err.to_string(), "boom");
    }

    #[test]
    fn test_new_no_dedup_vs_new_dedup() {
        let data = [
            (SampleID(1), 10),
            (SampleID(2), 20),
            (SampleID(3), 10), // Duplicate value
            (SampleID(4), 20), // Duplicate value
        ];

        let no_dedup = Sampled::new_no_dedup(data.iter().copied());
        let dedup = Sampled::new_dedup(data.iter().copied());

        // Both should have the same number of samples
        assert_eq!(no_dedup.num_samples(), 4);
        assert_eq!(dedup.num_samples(), 4);

        // But different number of data entries
        assert_eq!(no_dedup.data.len(), 4); // No deduplication - each sample has its own data entry
        assert_eq!(dedup.data.len(), 2); // Deduplication applied - duplicate values share data entries

        // Values should be the same
        for (id, _) in data {
            assert_eq!(no_dedup.get(id).unwrap(), dedup.get(id).unwrap());
        }
    }
}