henad-core 0.3.0

//! Stat series written as comma-separated values (CSV), one row per sampled tick.
//!
//! The column layout is fixed from the *first* sample and reused for every later row, so the
//! header and every row always agree. A model whose `stats()` shape changes mid-run is a
//! programming error, and [`StatsWriter::push`] reports it as one.

use std::fmt;
use std::io::{self, Write};

use crate::export::csv::{escape_field, fmt_f64};
use crate::view::{StatEntry, StatValue};

/// A stat series could not be written.
#[derive(Debug)]
pub enum StatsWriteError {
    /// Error of the underlying writer.
    Io(io::Error),
    /// A series changed shape after the first sample fixed the column layout.
    Shape(String),
}

impl fmt::Display for StatsWriteError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Io(err) => write!(f, "{err}"),
            Self::Shape(message) => f.write_str(message),
        }
    }
}

impl std::error::Error for StatsWriteError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            Self::Io(err) => Some(err),
            Self::Shape(_) => None,
        }
    }
}

impl From<io::Error> for StatsWriteError {
    fn from(err: io::Error) -> Self {
        Self::Io(err)
    }
}

/// Separator between a series label and a component suffix, e.g. `Average Velocity.x`.
const SUFFIX_SEP: char = '.';

/// One output column, and the part of the stat series feeding it.
#[derive(Debug, Clone, PartialEq, Eq)]
struct Column {
    /// Column name before CSV escaping.
    name: String,
    /// Header text, already escaped for CSV.
    header: String,
    /// Label of the series feeding the column.
    label: &'static str,
    /// Index into the `Vec<StatEntry>` returned by `stats()`.
    series: usize,
    part: Part,
}

impl Column {
    /// Returns this column's value in `stats`, the sample taken at `tick`.
    fn value(&self, tick: u64, stats: &[StatEntry]) -> Result<f64, StatsWriteError> {
        let Some(entry) = stats.get(self.series) else {
            return Err(StatsWriteError::Shape(format!(
                "stat series count changed mid-run: column '{}' needs series {} but tick {tick} has {}",
                self.header,
                self.series,
                stats.len()
            )));
        };
        part_value(&entry.value, self.part).ok_or_else(|| {
            StatsWriteError::Shape(format!(
                "stat series '{}' changed shape mid-run at tick {tick}: column '{}' no longer applies",
                entry.label, self.header
            ))
        })
    }
}

/// The scalar pulled out of a [`StatValue`] for one column.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Part {
    Scalar,
    VecX,
    VecY,
    /// Magnitude of a `Vector2D`, so a vector series is still usable without recombining lanes.
    VecMagnitude,
    /// Count in one histogram bucket.
    Bucket(usize),
    /// Total across all histogram buckets.
    BucketTotal,
}

/// Column layout of a stat series, planned from one sample.
///
/// A scalar series is one column, a vector is three (`.x`, `.y` and `.magnitude`), and a histogram
/// is one per bucket plus `.total`.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct StatColumns {
    columns: Vec<Column>,
}

impl StatColumns {
    /// Derives the column layout from the sample `stats`.
    pub fn plan(stats: &[StatEntry]) -> Self {
        let mut columns = Vec::new();
        for (series, entry) in stats.iter().enumerate() {
            let mut push = |suffix: Option<&str>, part: Part| {
                let name = match suffix {
                    Some(suffix) => format!("{}{SUFFIX_SEP}{suffix}", entry.label),
                    None => entry.label.to_owned(),
                };
                let header = escape_field(&name);
                columns.push(Column {
                    name,
                    header,
                    label: entry.label,
                    series,
                    part,
                });
            };
            match &entry.value {
                StatValue::Scalar(_) => push(None, Part::Scalar),
                StatValue::Vector2D { .. } => {
                    push(Some("x"), Part::VecX);
                    push(Some("y"), Part::VecY);
                    push(Some("magnitude"), Part::VecMagnitude);
                }
                StatValue::Histogram { edges, counts } => {
                    // Each bucket is named by its own range, and the columns can be read without the edge list.
                    // `edges` holds the boundaries, a lower and an upper edge per bucket. A bucket whose
                    // edges are missing is named by its index.
                    for bucket in 0..counts.len() {
                        let range = match (edges.get(bucket), edges.get(bucket + 1)) {
                            (Some(lo), Some(hi)) => format!("[{}, {})", fmt_f64(*lo), fmt_f64(*hi)),
                            _ => format!("bucket {bucket}"),
                        };
                        push(Some(&range), Part::Bucket(bucket));
                    }
                    push(Some("total"), Part::BucketTotal);
                }
            }
        }
        Self { columns }
    }

    /// Number of columns.
    pub fn len(&self) -> usize {
        self.columns.len()
    }

    /// Returns whether the layout has no columns.
    pub fn is_empty(&self) -> bool {
        self.columns.is_empty()
    }

    /// Returns the name of column `i`, before CSV escaping.
    ///
    /// # Panics
    ///
    /// Panics when `i` is not below [`Self::len`].
    pub fn name(&self, i: usize) -> &str {
        &self.columns[i].name
    }

    /// Returns the header of column `i`, escaped for CSV.
    ///
    /// # Panics
    ///
    /// Panics when `i` is not below [`Self::len`].
    pub fn header(&self, i: usize) -> &str {
        &self.columns[i].header
    }

    /// Returns whether column `i` counts a single histogram bucket.
    ///
    /// # Panics
    ///
    /// Panics when `i` is not below [`Self::len`].
    pub fn is_bucket(&self, i: usize) -> bool {
        matches!(self.columns[i].part, Part::Bucket(_))
    }

    /// Returns the index of the column called `name`.
    ///
    /// A bare vector or histogram label resolves to its magnitude or total column, the value
    /// [`StatValue::scalar`] returns. An exact column name wins over a bare label.
    pub fn resolve(&self, name: &str) -> Option<usize> {
        self.columns.iter().position(|column| column.name == name).or_else(|| {
            self.columns.iter().position(|column| {
                column.label == name && matches!(column.part, Part::VecMagnitude | Part::BucketTotal)
            })
        })
    }

    /// Replaces the contents of `out` with one value per column, read from the sample `stats` taken at `tick`.
    ///
    /// # Errors
    ///
    /// Returns [`StatsWriteError::Shape`] when `stats` no longer fits the planned layout.
    pub fn extract(&self, tick: u64, stats: &[StatEntry], out: &mut Vec<f64>) -> Result<(), StatsWriteError> {
        out.clear();
        for column in &self.columns {
            out.push(column.value(tick, stats)?);
        }
        Ok(())
    }
}

/// Streams a stat time series to a writer as CSV.
///
/// Call [`push`](Self::push) once per sampled tick, then [`finish`](Self::finish). The first `push`
/// writes the header, with the columns of its sample. A [`StatDescriptor`](crate::view::StatDescriptor)
/// does not say whether its value is a scalar, a vector or a histogram.
#[derive(Debug)]
pub struct StatsWriter<W: Write> {
    out: W,
    /// Column layout, `None` until the first `push` fixes it.
    columns: Option<StatColumns>,
    rows: u64,
}

impl<W: Write> StatsWriter<W> {
    /// Returns a writer over `out` that has written nothing.
    pub fn new(out: W) -> Self {
        Self {
            out,
            columns: None,
            rows: 0,
        }
    }

    /// Writes the row of the sample `stats`, taken at `tick`.
    ///
    /// The first call fixes the column layout and writes the header.
    ///
    /// # Errors
    ///
    /// Returns [`StatsWriteError::Io`] when a write fails, and [`StatsWriteError::Shape`] when `stats` does not fit
    /// the layout the first sample fixed.
    pub fn push(&mut self, tick: u64, stats: &[StatEntry]) -> Result<(), StatsWriteError> {
        if self.columns.is_none() {
            let columns = StatColumns::plan(stats);
            write!(self.out, "tick")?;
            for column in &columns.columns {
                write!(self.out, ",{}", column.header)?;
            }
            writeln!(self.out)?;
            self.columns = Some(columns);
        }
        let columns = self.columns.as_ref().map_or(&[][..], |planned| &planned.columns[..]);

        write!(self.out, "{tick}")?;
        for column in columns {
            let value = column.value(tick, stats)?;
            write!(self.out, ",{}", fmt_f64(value))?;
        }
        writeln!(self.out)?;
        self.rows += 1;
        Ok(())
    }

    /// Flushes the writer and returns the number of rows written.
    ///
    /// Note that a `BufWriter` dropped without a flush discards its write errors.
    ///
    /// # Errors
    ///
    /// Returns [`StatsWriteError::Io`] when the flush fails.
    pub fn finish(self) -> Result<u64, StatsWriteError> {
        Ok(self.into_inner()?.1)
    }

    /// Flushes the writer and returns it with the number of rows written.
    ///
    /// # Errors
    ///
    /// Returns [`StatsWriteError::Io`] when the flush fails.
    pub fn into_inner(mut self) -> Result<(W, u64), StatsWriteError> {
        self.out.flush()?;
        Ok((self.out, self.rows))
    }

    /// Number of rows written so far, not counting the header.
    pub fn rows(&self) -> u64 {
        self.rows
    }
}

/// Returns the scalar that `part` extracts from `value`, or `None` when `value` has no such part.
///
/// A `None` means the series changed shape after the layout was fixed.
fn part_value(value: &StatValue, part: Part) -> Option<f64> {
    match (value, part) {
        (StatValue::Scalar(v), Part::Scalar) => Some(*v),
        (StatValue::Vector2D { x, .. }, Part::VecX) => Some(*x),
        (StatValue::Vector2D { y, .. }, Part::VecY) => Some(*y),
        (StatValue::Vector2D { x, y }, Part::VecMagnitude) => Some(x.hypot(*y)),
        (StatValue::Histogram { counts, .. }, Part::Bucket(bucket)) => counts.get(bucket).map(|c| *c as f64),
        (StatValue::Histogram { counts, .. }, Part::BucketTotal) => Some(counts.iter().sum::<u64>() as f64),
        _ => None,
    }
}

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

    const C: [u8; 4] = [0, 0, 0, 255];

    fn scalar(label: &'static str, v: f64) -> StatEntry {
        StatEntry {
            label,
            value: StatValue::Scalar(v),
            color: C,
        }
    }

    fn vec2(label: &'static str, x: f64, y: f64) -> StatEntry {
        StatEntry {
            label,
            value: StatValue::Vector2D { x, y },
            color: C,
        }
    }

    fn hist(label: &'static str, edges: Vec<f64>, counts: Vec<u64>) -> StatEntry {
        StatEntry {
            label,
            value: StatValue::Histogram { edges, counts },
            color: C,
        }
    }

    /// Runs samples through a writer and returns the CSV text.
    fn render(samples: &[(u64, Vec<StatEntry>)]) -> String {
        let mut buf = Vec::new();
        let mut writer = StatsWriter::new(&mut buf);
        for (tick, stats) in samples {
            writer.push(*tick, stats).expect("push should succeed");
        }
        writer.finish().expect("finish should succeed");
        String::from_utf8(buf).expect("output should be utf8")
    }

    #[test]
    fn scalars_write_one_column_each() {
        let csv = render(&[
            (0, vec![scalar("Alive", 10.0), scalar("Dead", 5.0)]),
            (1, vec![scalar("Alive", 12.0), scalar("Dead", 3.0)]),
        ]);
        assert_eq!(csv, "tick,Alive,Dead\n0,10,5\n1,12,3\n");
    }

    #[test]
    fn vectors_split_into_x_y_and_magnitude() {
        let csv = render(&[(7, vec![vec2("Velocity", 3.0, 4.0)])]);
        assert_eq!(csv, "tick,Velocity.x,Velocity.y,Velocity.magnitude\n7,3,4,5\n");
    }

    #[test]
    fn histograms_label_buckets_by_range_and_add_a_total() {
        let csv = render(&[(0, vec![hist("Speed", vec![0.0, 1.0, 2.0], vec![4, 6])])]);
        assert_eq!(csv, "tick,\"Speed.[0, 1)\",\"Speed.[1, 2)\",Speed.total\n0,4,6,10\n");
    }

    /// A histogram with fewer edges than buckets still produces a column per bucket.
    #[test]
    fn histogram_falls_back_to_bucket_index_without_edges() {
        let csv = render(&[(0, vec![hist("H", vec![], vec![1, 2])])]);
        assert_eq!(csv, "tick,H.bucket 0,H.bucket 1,H.total\n0,1,2,3\n");
    }

    #[test]
    fn header_is_written_once_for_many_rows() {
        let csv = render(&[
            (0, vec![scalar("A", 1.0)]),
            (1, vec![scalar("A", 2.0)]),
            (2, vec![scalar("A", 3.0)]),
        ]);
        assert_eq!(csv.lines().filter(|l| l.starts_with("tick")).count(), 1);
        assert_eq!(csv.lines().count(), 4);
    }

    #[test]
    fn no_samples_writes_nothing() {
        let csv = render(&[]);
        assert!(csv.is_empty(), "expected empty output, got {csv:?}");
    }

    #[test]
    fn fractional_values_keep_precision() {
        let csv = render(&[(0, vec![scalar("A", 0.1 + 0.2)])]);
        // The sum keeps its round-trip digits past 0.3.
        assert!(csv.contains("0.30000000000000004"), "got {csv}");
    }

    #[test]
    fn non_finite_values_become_empty_cells() {
        let csv = render(&[(0, vec![scalar("A", f64::NAN), scalar("B", f64::INFINITY)])]);
        assert_eq!(csv, "tick,A,B\n0,,\n");
    }

    #[test]
    fn labels_containing_commas_are_quoted() {
        let csv = render(&[(0, vec![scalar("Susceptible, count", 1.0), scalar("B", 2.0)])]);
        assert_eq!(csv, "tick,\"Susceptible, count\",B\n0,1,2\n");
        // Three fields per row, so the comma did not shift the columns.
        assert_eq!(csv.lines().count(), 2);
    }

    #[test]
    fn finish_returns_the_row_count_excluding_the_header() {
        let mut buf = Vec::new();
        let mut writer = StatsWriter::new(&mut buf);
        writer.push(0, &[scalar("A", 1.0)]).expect("push");
        writer.push(1, &[scalar("A", 2.0)]).expect("push");
        assert_eq!(writer.finish().expect("finish"), 2);
    }

    #[test]
    fn a_series_disappearing_mid_run_is_an_error() {
        let mut buf = Vec::new();
        let mut writer = StatsWriter::new(&mut buf);
        writer.push(0, &[scalar("A", 1.0), scalar("B", 2.0)]).expect("push");
        let err = writer.push(1, &[scalar("A", 1.0)]).expect_err("should reject");
        assert!(err.to_string().contains("changed mid-run"), "got {err}");
    }

    #[test]
    fn a_series_changing_kind_mid_run_is_an_error() {
        let mut buf = Vec::new();
        let mut writer = StatsWriter::new(&mut buf);
        writer.push(0, &[scalar("A", 1.0)]).expect("push");
        let err = writer.push(1, &[vec2("A", 1.0, 2.0)]).expect_err("should reject");
        assert!(err.to_string().contains("changed shape"), "got {err}");
    }

    #[test]
    fn a_histogram_losing_buckets_mid_run_is_an_error() {
        let mut buf = Vec::new();
        let mut writer = StatsWriter::new(&mut buf);
        writer
            .push(0, &[hist("H", vec![0.0, 1.0, 2.0], vec![1, 2])])
            .expect("push");
        let err = writer
            .push(1, &[hist("H", vec![0.0, 1.0], vec![1])])
            .expect_err("should reject");
        assert!(err.to_string().contains("changed shape"), "got {err}");
    }

    #[test]
    fn mixed_kinds_keep_series_order() {
        let csv = render(&[(
            0,
            vec![
                scalar("S", 1.0),
                vec2("V", 0.0, 2.0),
                hist("H", vec![0.0, 1.0], vec![3]),
            ],
        )]);
        assert_eq!(csv, "tick,S,V.x,V.y,V.magnitude,\"H.[0, 1)\",H.total\n0,1,0,2,2,3,3\n");
    }

    #[test]
    fn a_model_with_no_stats_still_writes_ticks() {
        let csv = render(&[(0, vec![]), (5, vec![])]);
        assert_eq!(csv, "tick\n0\n5\n");
    }

    #[test]
    fn stat_columns_name_what_the_writer_writes() {
        let stats = vec![
            scalar("Susceptible, count", 7.0),
            vec2("V", 3.0, 4.0),
            hist("H", vec![0.0, 1.0, 2.0], vec![4, 6]),
        ];
        let columns = StatColumns::plan(&stats);
        let mut values = Vec::new();
        columns
            .extract(9, &stats, &mut values)
            .expect("a sample fits its own plan");

        let csv = render(&[(9, stats)]);
        let headers: Vec<&str> = (0..columns.len()).map(|i| columns.header(i)).collect();
        assert_eq!(csv.lines().next(), Some(format!("tick,{}", headers.join(",")).as_str()));

        let records = crate::export::csv::parse_records(&csv).expect("the writer writes valid CSV");
        let names: Vec<&str> = (0..columns.len()).map(|i| columns.name(i)).collect();
        assert_eq!(records[0][1..], names[..], "names are the headers unescaped");
        let row: Vec<String> = values.iter().map(|&value| fmt_f64(value)).collect();
        assert_eq!(records[1][1..], row[..]);
        assert_eq!(names[0], "Susceptible, count");

        let buckets: Vec<&str> = (0..columns.len())
            .filter(|&i| columns.is_bucket(i))
            .map(|i| columns.name(i))
            .collect();
        assert_eq!(buckets, ["H.[0, 1)", "H.[1, 2)"]);
    }

    #[test]
    fn a_bare_label_resolves_to_the_column_of_its_scalar() {
        let stats = vec![
            scalar("S", 1.0),
            vec2("V", 3.0, 4.0),
            hist("H", vec![0.0, 1.0, 2.0], vec![4, 6]),
        ];
        let columns = StatColumns::plan(&stats);
        let mut values = Vec::new();
        columns
            .extract(0, &stats, &mut values)
            .expect("a sample fits its own plan");
        for entry in &stats {
            let i = columns.resolve(entry.label).expect("every label resolves");
            assert_eq!(values[i], entry.value.scalar(), "{}", entry.label);
        }
        assert_eq!(columns.resolve("V"), columns.resolve("V.magnitude"));
        assert_eq!(columns.resolve("H"), columns.resolve("H.total"));
        assert_eq!(columns.resolve("V.y").map(|i| columns.name(i)), Some("V.y"));
        assert_eq!(columns.resolve("Missing"), None);

        // An exact name wins over a bare label.
        let shadowed = StatColumns::plan(&[vec2("V", 3.0, 4.0), scalar("V", 1.0)]);
        assert_eq!(shadowed.resolve("V"), Some(3));
    }

    /// The two paths by which a stat series reaches a file must produce the same file.
    ///
    /// The app writes through a [`StatsWriter`] while recording, and replays
    /// [`crate::view::StatsHistory`] otherwise. Both claim the column layout the headless runner
    /// writes, and only one of them holds structured values.
    mod parity {
        use super::super::StatsWriter;
        use crate::view::{StatDescriptor, StatEntry, StatValue, StatsHistory};

        const C: [u8; 4] = [9, 9, 9, 255];

        /// A scalar beside a vector, so the replay has to reach both stores.
        fn sample(i: u64) -> Vec<StatEntry> {
            vec![
                StatEntry {
                    label: "Alive",
                    value: StatValue::Scalar(i as f64 * 1.5),
                    color: C,
                },
                StatEntry {
                    label: "Average Velocity",
                    value: StatValue::Vector2D {
                        x: i as f64,
                        y: -(i as f64) * 0.25,
                    },
                    color: C,
                },
            ]
        }

        fn render(samples: &[(u64, Vec<StatEntry>)]) -> String {
            let mut writer = StatsWriter::new(Vec::new());
            for (tick, stats) in samples {
                writer.push(*tick, stats).expect("push should succeed");
            }
            let (csv, _) = writer.into_inner().expect("finish should succeed");
            String::from_utf8(csv).expect("output should be utf8")
        }

        #[test]
        fn a_replayed_history_writes_what_a_recording_writes() {
            let descriptors = vec![
                StatDescriptor::new("Alive", C),
                StatDescriptor::new("Average Velocity", C),
            ];
            let mut history = StatsHistory::new(descriptors, Some(64));
            let mut recorded = Vec::new();
            for tick in 0..20 {
                let stats = sample(tick);
                history.push_entries(&stats, tick);
                recorded.push((tick, stats));
            }

            let replayed: Vec<(u64, Vec<StatEntry>)> = (0..history.len())
                .map(|j| (history.tick(j).expect("tick"), history.entries(j).expect("entries")))
                .collect();

            assert_eq!(render(&replayed), render(&recorded));
        }

        /// A history that has wrapped writes the tail of what a recording holds, column for column.
        #[test]
        fn a_wrapped_history_writes_the_tail_of_the_recording() {
            let descriptors = vec![
                StatDescriptor::new("Alive", C),
                StatDescriptor::new("Average Velocity", C),
            ];
            let mut history = StatsHistory::new(descriptors, Some(5));
            let mut recorded = Vec::new();
            for tick in 0..20 {
                let stats = sample(tick);
                history.push_entries(&stats, tick);
                recorded.push((tick, stats));
            }

            let replayed: Vec<(u64, Vec<StatEntry>)> = (0..history.len())
                .map(|j| (history.tick(j).expect("tick"), history.entries(j).expect("entries")))
                .collect();

            let full = render(&recorded);
            let tail = render(&replayed);
            let header = full.lines().next().expect("a header");
            assert_eq!(
                tail.lines().next(),
                Some(header),
                "the columns must not depend on the path"
            );
            for line in tail.lines().skip(1) {
                assert!(full.contains(line), "replayed row {line:?} is not in the recording");
            }
            assert_eq!(tail.lines().count(), 6, "five samples plus the header");
        }
    }
}