ggplot-rs 0.34.0

use std::borrow::Cow;
use std::collections::HashMap;

use crate::data::{DataFrame, Value};

use super::{Position, PositionParams};

/// Stack bars/areas on top of each other (ggplot2's `position_stack(vjust
/// = )`).
///
/// `vjust` places the row's `y` within its stacked segment: `1` (the
/// default) puts it at the segment's top — the pre-B3-4 behaviour, and
/// byte-identical (no `ymax` column is written). Anything else writes
/// `y = (1 − vjust)·ymin + vjust·ymax` and an `ymax` column holding the
/// segment's top, so e.g. `vjust(0.5)` centres a label in its segment.
///
/// `PositionStack` is also usable as a bare value (`vjust` 1, as before
/// the struct gained fields) via the [`PositionStack`](const@PositionStack)
/// const.
#[derive(Clone, Debug, PartialEq)]
pub struct PositionStack {
    pub vjust: f64,
    /// Baseline of each x's stack ([`StackOffset`]); `Zero` by default.
    pub offset: StackOffset,
}

/// Where each x's stack starts (d3-shape's stack offsets). Non-`Zero`
/// offsets are meant for non-negative series (streamgraphs).
#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
pub enum StackOffset {
    /// Stacks rise from 0 (ggplot2's `position_stack`).
    #[default]
    Zero,
    /// Stacks are centred on 0 (ThemeRiver).
    Silhouette,
    /// Byron & Wattenberg's wiggle-minimising baseline (d3's
    /// `stackOffsetWiggle`), re-centred on 0 — the classic streamgraph.
    Wiggle,
}

impl Default for PositionStack {
    fn default() -> Self {
        PositionStack {
            vjust: 1.0,
            offset: StackOffset::Zero,
        }
    }
}

/// The default stack (`vjust` 1): lets `.position(PositionStack)` keep
/// compiling now that the type has fields.
#[allow(non_upper_case_globals)]
pub const PositionStack: PositionStack = PositionStack {
    vjust: 1.0,
    offset: StackOffset::Zero,
};

impl PositionStack {
    /// Place `y` within each stacked segment (`0` = bottom, `0.5` =
    /// middle, `1` = top, the default).
    pub fn with_vjust(mut self, vjust: f64) -> Self {
        self.vjust = vjust;
        self
    }

    /// Shift each x's stack to the given baseline ([`StackOffset`]).
    pub fn with_offset(mut self, offset: StackOffset) -> Self {
        self.offset = offset;
        self
    }
}

/// Per-row baseline shift for a non-`Zero` [`StackOffset`]: every x's stack
/// (rows in data order, the first row on top) moves by the same amount.
fn baseline_shift(
    offset: StackOffset,
    data: &DataFrame,
    x_col: &[Value],
    y_col: &[Value],
) -> Vec<f64> {
    // Distinct x in increasing order (numeric x by value, otherwise by
    // first appearance), and each row's x index.
    let mut xs: Vec<(String, Option<f64>)> = Vec::new();
    let mut x_index: HashMap<String, usize> = HashMap::new();
    for x in x_col {
        let k = x.key_str().into_owned();
        if !x_index.contains_key(&k) {
            x_index.insert(k.clone(), xs.len());
            xs.push((k, x.as_f64()));
        }
    }
    let mut order: Vec<usize> = (0..xs.len()).collect();
    if xs.iter().all(|x| x.1.is_some()) {
        order.sort_by(|&a, &b| xs[a].1.unwrap().total_cmp(&xs[b].1.unwrap()));
    }
    let mut slot = vec![0; xs.len()];
    for (j, &k) in order.iter().enumerate() {
        slot[k] = j;
    }
    let row_slot: Vec<usize> = x_col
        .iter()
        .map(|x| slot[x_index[x.key_str().as_ref()]])
        .collect();
    let m = xs.len();

    // Series (bottom → top): the stack puts the first-seen series on top.
    let series_col = ["group", "fill", "color"]
        .iter()
        .find_map(|c| data.column(c));
    let mut series: Vec<String> = Vec::new();
    let mut s_index: HashMap<String, usize> = HashMap::new();
    let row_series: Vec<usize> = (0..x_col.len())
        .map(|i| {
            let k = series_col
                .map(|c| c[i].key_str().into_owned())
                .unwrap_or_default();
            *s_index.entry(k.clone()).or_insert_with(|| {
                series.push(k);
                series.len() - 1
            })
        })
        .collect();
    let n = series.len();
    let mut v = vec![vec![0.0; m]; n];
    for i in 0..x_col.len() {
        if let Some(y) = super::stack_row(&y_col[i]).value() {
            v[n - 1 - row_series[i]][row_slot[i]] += y.max(0.0);
        }
    }
    let totals: Vec<f64> = (0..m).map(|j| (0..n).map(|s| v[s][j]).sum()).collect();
    let base: Vec<f64> = match offset {
        StackOffset::Zero => vec![0.0; m],
        StackOffset::Silhouette => totals.iter().map(|t| -t / 2.0).collect(),
        StackOffset::Wiggle => {
            let mut g = vec![0.0; m];
            let mut y = 0.0;
            for j in 1..m {
                let (mut s1, mut s2) = (0.0, 0.0);
                for i in 0..n {
                    let mut s3 = (v[i][j] - v[i][j - 1]) / 2.0;
                    for row in v.iter().take(i) {
                        s3 += row[j] - row[j - 1];
                    }
                    s1 += v[i][j];
                    s2 += s3 * v[i][j];
                }
                if s1 != 0.0 {
                    y -= s2 / s1;
                }
                g[j] = y;
            }
            // Re-centre: the stream's mid-line averages 0.
            let mid =
                g.iter().zip(&totals).map(|(g, t)| g + t / 2.0).sum::<f64>() / m.max(1) as f64;
            g.iter().map(|g| g - mid).collect()
        }
    };
    row_slot.iter().map(|&j| base[j]).collect()
}

/// `position_stack()` / `position_stack().with_vjust(v)` — see
/// [`PositionStack`](struct@PositionStack).
pub fn position_stack() -> PositionStack {
    PositionStack::default()
}

impl Position for PositionStack {
    fn compute(&self, data: &mut DataFrame, _params: &PositionParams) {
        // Group by x, accumulate y values
        let x_col = match data.column("x") {
            Some(c) => c.to_vec(),
            None => return,
        };
        let y_col = match data.column("y") {
            Some(c) => c.to_vec(),
            None => return,
        };
        super::preserve_raw_y(data, &y_col);

        let n = y_col.len();
        let mut tops = vec![0.0_f64; n];
        let mut bottoms = vec![0.0_f64; n];

        // ggplot2 stacks a group's positive and negative values
        // independently within each x (#154, ggplot2's internal
        // `pos`/`neg` split in `collide`/`pos_stack`): run the same top-
        // down, legend-order accumulation twice, once over the rows
        // with y >= 0 and once over y < 0, so e.g. a bar chart mixing a
        // group's positive and negative segments at the same x gets two
        // independent stacks (one above zero, one below) instead of a
        // single pooled running total that lets a negative row cancel
        // into a positive one's segment. ggplot2 stacks the first group
        // at the top of its bucket (so top-to-bottom matches the
        // legend), so each bucket still accumulates downward from its
        // own total rather than upward from 0 — "top"/"bottom" here are
        // in that accumulation order, not value order (a negative
        // bucket's "top", nearer 0, has the smaller `tops[i]`), so
        // `ymin`/`ymax`/`y` below always take the min/max of the pair
        // rather than assume `tops[i] >= bottoms[i]`.
        // Per-x accumulators keyed by the borrowed x key: O(1) per row
        // instead of a linear scan over the distinct x values.
        //
        // A missing row (`super::stack_row`, review #167 finding 2 on
        // #171) joins the positive bucket but contributes 0 to its
        // total, and its own top/bottom are left NaN below — so it
        // doesn't poison any other row's stack, matching ggplot2's
        // `pos_stack` (`y <- ifelse(is.na(df$y), 0, df$y)`).
        for positive in [true, false] {
            let mut totals: HashMap<Cow<'_, str>, f64> = HashMap::new();
            for (x, y) in x_col.iter().zip(y_col.iter()) {
                let row = super::stack_row(y);
                if row.is_positive_bucket() != positive {
                    continue;
                }
                *totals.entry(x.key_str()).or_insert(0.0) += row.contribution();
            }

            let mut consumed: HashMap<Cow<'_, str>, f64> = HashMap::new();
            for (i, (x, y)) in x_col.iter().zip(y_col.iter()).enumerate() {
                let row = super::stack_row(y);
                if row.is_positive_bucket() != positive {
                    continue;
                }
                let x_key = x.key_str();
                let total = totals.get(&x_key).copied().unwrap_or(0.0);
                let run = consumed.get(&x_key).copied().unwrap_or(0.0);

                let Some(y_val) = row.value() else {
                    // Missing: draw nothing for this row (NA in
                    // ggplot2), but don't advance `consumed` — it
                    // contributed 0, so the next row in this bucket
                    // still starts from the same `run`.
                    tops[i] = f64::NAN;
                    bottoms[i] = f64::NAN;
                    continue;
                };

                // This group occupies [total - run - y, total - run]
                // within its own (same-sign) bucket.
                tops[i] = total - run;
                bottoms[i] = total - run - y_val;

                *consumed.entry(x_key).or_insert(0.0) += y_val;
            }
        }

        if self.offset != StackOffset::Zero {
            let shift = baseline_shift(self.offset, data, &x_col, &y_col);
            for i in 0..n {
                tops[i] += shift[i];
                bottoms[i] += shift[i];
            }
        }

        // A non-finite vjust is validated (and warned on) separately; here
        // it falls back to 1 (the byte-identical top-anchored path).
        let vjust = if self.vjust.is_finite() {
            self.vjust
        } else {
            1.0
        };

        let write_ymax = (vjust - 1.0).abs() > f64::EPSILON;
        let lo_hi: Vec<(f64, f64)> = tops
            .iter()
            .zip(bottoms.iter())
            .map(|(&t, &b)| (t.min(b), t.max(b)))
            .collect();
        let new_y: Vec<Value> = if write_ymax {
            lo_hi
                .iter()
                .map(|&(lo, hi)| Value::Float((1.0 - vjust) * lo + vjust * hi))
                .collect()
        } else {
            lo_hi.iter().map(|&(_, hi)| Value::Float(hi)).collect()
        };
        let ymin_vals: Vec<Value> = lo_hi.iter().map(|&(lo, _)| Value::Float(lo)).collect();

        if let Some(col) = data.column_mut("y") {
            *col = new_y;
        }
        if !data.has_column("ymin") {
            data.add_column("ymin".to_string(), ymin_vals);
        } else if let Some(col) = data.column_mut("ymin") {
            *col = ymin_vals;
        }
        if write_ymax {
            let ymax_vals: Vec<Value> = lo_hi.iter().map(|&(_, hi)| Value::Float(hi)).collect();
            if !data.has_column("ymax") {
                data.add_column("ymax".to_string(), ymax_vals);
            } else if let Some(col) = data.column_mut("ymax") {
                *col = ymax_vals;
            }
        }
    }

    fn name(&self) -> &str {
        "stack"
    }

    fn per_panel(&self) -> bool {
        true
    }

    fn validate(&self, warnings: &mut Vec<String>) {
        if !self.vjust.is_finite() {
            warnings.push(format!(
                "position_stack: vjust({}) must be finite; falling back to 1",
                self.vjust
            ));
        }
    }
}

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

    fn frame(x: &[f64], g: &[&str], y: &[f64]) -> DataFrame {
        let mut df = DataFrame::new();
        df.add_column("x".into(), x.iter().map(|v| Value::Float(*v)).collect());
        df.add_column(
            "fill".into(),
            g.iter().map(|s| Value::Str(s.to_string())).collect(),
        );
        df.add_column("y".into(), y.iter().map(|v| Value::Float(*v)).collect());
        df
    }

    fn col(df: &DataFrame, c: &str) -> Vec<f64> {
        df.column(c)
            .unwrap()
            .iter()
            .map(|v| v.as_f64().unwrap())
            .collect()
    }

    /// R parity (ggplot2 4.0.3, `layer_data()`): x=1, fill groups a/b/c/d
    /// with y = 3, -2, 4, -1 and `position_stack(vjust = 0.5)` gives
    /// y/ymin/ymax = (5.5, 4, 7), (-2, -3, -1), (2, 0, 4), (-0.5, -1, 0)
    /// (#154) — positives (a, c) stack above 0 using only their own sum
    /// (7) as the total; negatives (b, d) stack below 0 using only their
    /// own sum (-3), independently.
    #[test]
    fn negative_and_positive_values_stack_independently() {
        let mut df = frame(&[1.0; 4], &["a", "b", "c", "d"], &[3.0, -2.0, 4.0, -1.0]);
        PositionStack::default()
            .with_vjust(0.5)
            .compute(&mut df, &PositionParams::default());
        let y = col(&df, "y");
        let ymin = col(&df, "ymin");
        let ymax = col(&df, "ymax");
        for (i, (want_y, want_min, want_max)) in [
            (5.5, 4.0, 7.0),
            (-2.0, -3.0, -1.0),
            (2.0, 0.0, 4.0),
            (-0.5, -1.0, 0.0),
        ]
        .into_iter()
        .enumerate()
        {
            assert!((y[i] - want_y).abs() < 1e-9, "y[{i}]={}", y[i]);
            assert!((ymin[i] - want_min).abs() < 1e-9, "ymin[{i}]={}", ymin[i]);
            assert!((ymax[i] - want_max).abs() < 1e-9, "ymax[{i}]={}", ymax[i]);
        }
    }

    /// All-positive data (the common case) is byte-identical to the
    /// pre-#154 pooled algorithm: the negative bucket is simply empty.
    #[test]
    fn all_positive_values_are_unaffected() {
        let mut df = frame(&[1.0, 1.0], &["a", "b"], &[3.0, 5.0]);
        PositionStack.compute(&mut df, &PositionParams::default());
        assert_eq!(col(&df, "y"), vec![8.0, 5.0]);
        assert_eq!(col(&df, "ymin"), vec![5.0, 0.0]);
    }

    /// Review #167 finding 2 on sipemu/ggplot-rs#171 (originally #170's
    /// NaN bucket fix, which got the *bucket* right but the *arithmetic*
    /// wrong): a missing `y` joins the positive bucket, like ggplot2,
    /// but contributes **0** to its running total — not the raw NaN —
    /// so the *other* row sharing that bucket (`a`) stacks exactly as if
    /// `b` weren't there at all, instead of turning into NaN too.
    /// ggplot2 4.0.3 (`layer_data()`, `pos_stack`'s `y <-
    /// ifelse(is.na(df$y), 0, df$y)`): a=3, b=NaN, c=-2 gives a 0..3,
    /// c -2..0, b's own y/ymin left `NA`. `c` (the negative bucket) was
    /// already unaffected either way — it shares no bucket with `b`.
    #[test]
    fn missing_y_contributes_zero_and_does_not_poison_its_bucket() {
        let mut df = frame(&[1.0, 1.0, 1.0], &["a", "b", "c"], &[3.0, f64::NAN, -2.0]);
        PositionStack.compute(&mut df, &PositionParams::default());
        let y = col(&df, "y");
        let ymin = col(&df, "ymin");
        assert_eq!(
            (y[0], ymin[0]),
            (3.0, 0.0),
            "a must stack as if b's row weren't there at all: y={y:?} ymin={ymin:?}"
        );
        assert!(
            y[1].is_nan() && ymin[1].is_nan(),
            "b's own y/ymin must stay NaN (not drawn): y={y:?} ymin={ymin:?}"
        );
        assert_eq!(
            (y[2], ymin[2]),
            (0.0, -2.0),
            "c (negative bucket) must be unaffected: y={y:?} ymin={ymin:?}"
        );
    }

    /// The pre-existing bug finding 2 also fixes: a `Value::Na` `y` (not
    /// just a `Float(NaN)`) used to fall back to `0.0` via
    /// `as_f64().unwrap_or(0.0)` and get stacked as a *real* point at
    /// the top of its bucket, drawing a line straight through a
    /// supposedly-removed row. It must now be treated exactly like a
    /// `Float(NaN)` — contributing 0, with its own output left NaN.
    #[test]
    fn value_na_is_treated_like_float_nan_not_zero() {
        let mut df = DataFrame::new();
        df.add_column(
            "x".into(),
            vec![Value::Float(1.0), Value::Float(1.0), Value::Float(1.0)],
        );
        df.add_column(
            "fill".into(),
            vec![
                Value::Str("a".into()),
                Value::Str("b".into()),
                Value::Str("c".into()),
            ],
        );
        df.add_column(
            "y".into(),
            vec![Value::Float(3.0), Value::Na, Value::Float(-2.0)],
        );
        PositionStack.compute(&mut df, &PositionParams::default());
        let y = col(&df, "y");
        let ymin = col(&df, "ymin");
        assert_eq!(
            (y[0], ymin[0]),
            (3.0, 0.0),
            "a must be unaffected by b's Value::Na: y={y:?} ymin={ymin:?}"
        );
        assert!(
            y[1].is_nan() && ymin[1].is_nan(),
            "a Value::Na row must not become a real point at the stack top: \
             y={y:?} ymin={ymin:?}"
        );
    }

    #[test]
    fn silhouette_centres_each_stack_on_zero() {
        let mut df = frame(
            &[1.0, 1.0, 2.0, 2.0],
            &["a", "b", "a", "b"],
            &[2.0, 4.0, 1.0, 5.0],
        );
        position_stack()
            .with_offset(StackOffset::Silhouette)
            .compute(&mut df, &PositionParams::default());
        // x=1: total 6 → [-3, 3]; "a" on top: [1, 3], "b" [-3, 1].
        assert_eq!(col(&df, "ymin"), vec![1.0, -3.0, 2.0, -3.0]);
        assert_eq!(col(&df, "y"), vec![3.0, 1.0, 3.0, 2.0]);
    }

    #[test]
    fn wiggle_is_centred_and_keeps_thickness() {
        let x = [1.0, 1.0, 2.0, 2.0, 3.0, 3.0];
        let mut df = frame(
            &x,
            &["a", "b", "a", "b", "a", "b"],
            &[1.0, 3.0, 4.0, 2.0, 2.0, 2.0],
        );
        position_stack()
            .with_offset(StackOffset::Wiggle)
            .compute(&mut df, &PositionParams::default());
        let (lo, hi) = (col(&df, "ymin"), col(&df, "y"));
        for i in 0..6 {
            let want = [1.0, 3.0, 4.0, 2.0, 2.0, 2.0][i];
            assert!((hi[i] - lo[i] - want).abs() < 1e-12, "row {i}");
        }
        // Mid-lines of the three stacks average 0.
        let mid: f64 = [0usize, 2, 4]
            .iter()
            .map(|&i| (hi[i] + lo[i + 1]) / 2.0)
            .sum::<f64>();
        assert!(mid.abs() < 1e-9, "{mid}");
    }
}