// The index-generic reduction kernel, and the **canonical** statement of the
// null / `NaN` rules every column reduction follows.
//
// `Series::sum` / `mean` / `min` / `max`, the grouped `agg`
// reducers, and the scoped `eval_agg` each used to carry their own
// null-skipping / `NaN`-handling fold. They now share the factories here: each
// resolves the dtype dispatch and (for the `Builtin` backend) the validity
// mask **once**, returning a total fold that reduces an index window — a count
// `n` plus an indexing function `at` mapping `0 ..< n` onto source rows — to a
// single `Scalar`. The whole-column callers fold `(len, k => k)`; the per-group
// caller hoists the factory once and folds each group `(g.length(), k => g[k])`,
// so the grouped reduction stays a single linear pass over the column rather
// than re-materialising the validity mask per group.
//
// The expression layer's aggregation family (`Expr::std` / `variance` / `median` /
// `n_unique` / `first` / `last`) reduces through the same kernel: each is a
// window reducer here, so it serves the whole-column, per-group, and per-scope
// surfaces with one fold and one set of null / `NaN` rules.
//
// The rules, canonically: `sum` / `mean` add `NaN` in (it is a
// value, only `Null` is missing) so it propagates — and `var` / `std`, built on
// the mean, propagate it too; the extremum fold skips it (treated as missing,
// matching Polars' regular `min` / `max` and `sort`), and `median`, an order
// statistic, skips it the same way. `n_unique` keys a `NaN` into one bucket (it
// is one value), and `first` / `last` are positional, so they skip nothing —
// a present `NaN` is returned verbatim.
//
// One qualification, since "canonical" is not "only": four of the `Series`
// reductions — `sum` / `mean` / `min` / `max` — dispatch a `Numeric`-backed
// column to `NumericColumn`'s bitmap-free fold (`internal/column/numeric.mbt`)
// before reaching here, and that fold implements the empty-column and `NaN`
// rules above a second time. So these are the rules to *read*, and to change
// first, but not the only code that encodes them: the two are held cell-for-cell
// in step by `series/backend_parity_test.mbt`, which runs every
// backend-sensitive `Series` op on both representations of the same column and
// asserts the results identical. The remaining reductions take no such fast
// path — this kernel already skips the validity materialisation for an all-valid
// `Numeric` column (`reduce_present` hands back a constant-`true` predicate),
// which is most of what the fast path buys.

///|
/// Per-row presence predicate, hoisted once per reduction. A `Builtin` column
/// materialises its validity mask; an all-valid `Numeric` column needs none
/// (every slot is present), so the predicate is constant `true` and no bitmap
/// is allocated. Read at a *source* row index, which the reducers reach through
/// their `at` mapping.
fn reduce_present(storage : @column.ColumnStorage) -> (Int) -> Bool {
  match storage.kind() {
    @column.StorageKind::Numeric => _ => true
    @column.StorageKind::Builtin => {
      let valid = storage.validity().to_bools()
      i => valid[i]
    }
  }
}

///|
/// `sum` / `mean` reducer over an index window. `sum` accumulates `Int` in
/// `Int64` (overflowing past 2^63 — the price of an exact integer sum) and
/// `Float` in `Double`; `mean`, whose output is `Double`, accumulates `Int` in
/// `Double` too, so it cannot wrap past 2^63 and flip the mean's sign (matching
/// Polars). `NaN` is summed in (so it propagates) — only a `Null`
/// cell is skipped. With `want_mean` the running total divides by the non-null
/// count and is `Scalar::Null` for an empty / all-null window (the surfaces
/// translate that: `Series::mean` to `InvalidOperation`, the grouped / scoped
/// reducers to a null cell); otherwise the result is the dtype-preserving sum,
/// the additive identity for an empty / all-null window. A non-numeric
/// (`Bool` / `String`) column `raise`s `TypeMismatch` here — at factory time,
/// before any window is folded — matching the previous per-surface reducers,
/// which rejected the dtype up front (and so reject it even for zero groups).
fn reduce_arith(
  storage : @column.ColumnStorage,
  want_mean~ : Bool,
) -> ((Int, (Int) -> Int) -> @types.Scalar) raise @types.DataError {
  let present = reduce_present(storage)
  match storage.data() {
    @column.ColumnData::Int(a) =>
      if want_mean {
        // The mean's output is `Double`, so accumulate the numerator in
        // `Double` from the start. An `Int64` accumulator would silently wrap
        // past 2^63 and flip the mean's sign — a result outside `[min, max]` —
        // even though `Double` represents the true mean (matching Polars).
        fn(n : Int, at : (Int) -> Int) {
          let mut acc = 0.0
          let mut cnt = 0
          for k in 0.. Int) {
          let mut acc : Int64 = 0
          for k in 0..
      fn(n : Int, at : (Int) -> Int) {
        let mut acc = 0.0
        let mut cnt = 0
        for k in 0.. {
      let verb = if want_mean { "mean" } else { "sum" }
      raise @types.DataError::TypeMismatch(
        @types.TypeMismatchDetail::Message(
          "\{verb} undefined for \{storage.dtype()}",
        ),
      )
    }
  }
}

///|
/// `min` / `max` reducer over an index window (`want_min` picks which),
/// keeping the source dtype. Total over every dtype — `Bool` orders
/// `false < true`, `String` by `compare_string_lex` — returning the dtype's
/// `Scalar` variant or `Scalar::Null` for an empty / all-null / all-`NaN`
/// (`Float`) window. `NaN` is treated as missing (skipped), matching
/// `Series::min` / `max` and `sort`; only a null slot is
/// otherwise skipped. Shares the single `fold_extremum` accumulation.
fn reduce_extremum(
  storage : @column.ColumnStorage,
  want_min~ : Bool,
) -> (Int, (Int) -> Int) -> @types.Scalar {
  let present = reduce_present(storage)
  match storage.data() {
    @column.ColumnData::Int(a) =>
      fn(n : Int, at : (Int) -> Int) {
        let get = k => a[at(k)]
        let skip = k => !present(at(k))
        let better = (x, y) => if want_min { x < y } else { x > y }
        match @numeric.fold_extremum(n, get, skip, better) {
          Some(b) => @types.Scalar::Int(b)
          None => @types.Scalar::Null
        }
      }
    @column.ColumnData::Float(a) =>
      fn(n : Int, at : (Int) -> Int) {
        let get = k => a[at(k)]
        // `NaN` is treated as missing (skipped), as in `Series::min`.
        let skip = k => !present(at(k)) || a[at(k)].is_nan()
        let better = (x, y) => if want_min { x < y } else { x > y }
        match @numeric.fold_extremum(n, get, skip, better) {
          Some(b) => @types.Scalar::Float(b)
          None => @types.Scalar::Null
        }
      }
    @column.ColumnData::Bool(a) =>
      fn(n : Int, at : (Int) -> Int) {
        let get = k => a[at(k)]
        let skip = k => !present(at(k))
        // `false < true`: under min `false` beats `true`, under max reverse.
        let better = (x, y) => if want_min { !x && y } else { x && !y }
        match @numeric.fold_extremum(n, get, skip, better) {
          Some(b) => @types.Scalar::Bool(b)
          None => @types.Scalar::Null
        }
      }
    @column.ColumnData::String(a) =>
      fn(n : Int, at : (Int) -> Int) {
        let get = k => a[at(k)]
        let skip = k => !present(at(k))
        let better = (x, y) => {
          if want_min {
            @text.compare_string_lex(x, y) < 0
          } else {
            @text.compare_string_lex(x, y) > 0
          }
        }
        match @numeric.fold_extremum(n, get, skip, better) {
          Some(b) => @types.Scalar::String(b)
          None => @types.Scalar::Null
        }
      }
  }
}

///|
/// `count` reducer over an index window: the number of non-null cells, as a
/// `Scalar::Int`. Dtype-independent (it reads only the presence predicate), so
/// it needs no per-dtype dispatch and is total over every column.
fn reduce_count(
  storage : @column.ColumnStorage,
) -> (Int, (Int) -> Int) -> @types.Scalar {
  let present = reduce_present(storage)
  fn(n : Int, at : (Int) -> Int) {
    let mut c : Int64 = 0
    for k in 0.. Double` cell
/// reader: an `Int` cell widens through `to_double`, a `Float` cell reads
/// directly. Shared by the numeric-only reductions `variance` / `std` and
/// `median`, which otherwise repeat this same three-arm match. A non-numeric
/// (`Bool` / `String`) column has no such reading, so it `raise`s `TypeMismatch`
/// here — naming the operation (`op`) that asked — at factory time, before any
/// window is folded, exactly as `reduce_arith` rejects a non-numeric column.
fn numeric_double_reader(
  storage : @column.ColumnStorage,
  op : String,
) -> ((Int) -> Double) raise @types.DataError {
  match storage.data() {
    @column.ColumnData::Int(a) => i => a[i].to_double()
    @column.ColumnData::Float(a) => i => a[i]
    @column.ColumnData::Bool(_) | @column.ColumnData::String(_) =>
      raise @types.DataError::TypeMismatch(
        @types.TypeMismatchDetail::Message(
          "\{op} undefined for \{storage.dtype()}",
        ),
      )
  }
}

///|
/// `var` / `std` reducer over an index window (`want_std` square-roots the
/// result): the **sample** statistics with `ddof = 1` (Polars' default). The
/// variance is `Σ(xᵢ − mean)² / (cnt − 1)` over the `cnt` non-null cells and
/// the standard deviation its square root, both `Scalar::Float` — `Int` widens
/// to `Double`, so the output always types `Float`. Computed by Welford's
/// online algorithm (a single pass that never forms the raw sum), so a
/// finite-variance window of near-`Double`-max values — whose sum would
/// overflow to `±inf` and poison a two-pass mean, reporting `+inf` for a true
/// variance of 0 — still yields its finite variance. When the true variance
/// itself exceeds `Double`'s range (finite cells more than `Double::MAX`
/// apart, e.g. opposite signs near `±1.8e308`), the result saturates to
/// `+inf` — never a negative variance or a `NaN` std from the internal
/// overflow — matching numpy / Polars. Numeric only: a
/// non-numeric (`Bool` / `String`) column `raise`s `TypeMismatch` here at
/// factory time, before any window is folded, exactly as `reduce_arith` does.
/// A `NaN` is a present value: when at least two non-null cells are present it
/// counts toward `cnt` and **propagates** — the mean is `NaN`, hence every
/// deviation and the result — matching `sum` / `mean`. The fewer-than-two rule
/// takes precedence, though: a window with fewer than two non-null cells has no
/// sample variance (the `cnt − 1` denominator is zero or negative), so it
/// reduces to `Scalar::Null` (the empty / single-value / all-null cases — so a
/// lone `NaN` cell is `Null`, not `NaN`), which the surfaces translate the way
/// they do `mean`'s null cell.
fn reduce_variance(
  storage : @column.ColumnStorage,
  want_std~ : Bool,
) -> ((Int, (Int) -> Int) -> @types.Scalar) raise @types.DataError {
  let present = reduce_present(storage)
  let read = numeric_double_reader(
    storage,
    if want_std {
      "std"
    } else {
      "variance"
    },
  )
  fn(n : Int, at : (Int) -> Int) {
    // Welford's online algorithm: one pass that never forms the raw sum, so a
    // finite-variance window of near-`Double`-max values (whose sum would
    // overflow to `±inf` and poison a two-pass mean, turning a true variance of
    // 0 into `+inf`) still yields its finite variance. `mean` is the running
    // mean, `m2` the running sum of squared deviations from it.
    let mut cnt = 0
    let mut mean = 0.0
    let mut m2 = 0.0
    let mut saw_nonfinite = false
    for k in 0..= 0` under finite in-range arithmetic, but `delta`
      // itself overflows to `±inf` when two finite cells sit more than
      // `Double::MAX` apart (opposite signs near `±1.8e308`), and the
      // artefact poisons `m2` to `-inf` or `NaN` — reporting a mathematically
      // impossible negative variance and a `NaN` std. With every *input*
      // finite, a negative/`NaN` `m2` can only be that artefact, so saturate
      // to `+inf`: the true variance exceeds `Double`'s range (numpy/Polars
      // report `+inf` here too), and `sqrt(+inf) = +inf` keeps std
      // consistent. Genuine `NaN` / `±inf` inputs skip the rescue and keep
      // IEEE propagation, matching `sum` / `mean`.
      let m2 = if !saw_nonfinite && !(m2 >= 0.0) {
        @double.infinity
      } else {
        m2
      }
      let variance = m2 / (cnt - 1).to_double()
      @types.Scalar::Float(if want_std { variance.sqrt() } else { variance })
    }
  }
}

///|
/// `median` reducer over an index window: the middle of the sorted non-null
/// cells, or the mean of the two middles for an even count, as a
/// `Scalar::Float` — `Int` widens to `Double`, so a median always types
/// `Float` (Polars). Numeric only: a non-numeric column `raise`s
/// `TypeMismatch` at factory time, as `reduce_arith` does. A `NaN` is
/// **skipped** (treated as missing), the `min` / `max` and `sort` rule rather
/// than `sum` / `mean` propagation — MoonFrame orders `NaN` as missing
/// throughout (see `sort`), so this order statistic drops it too. An empty /
/// all-null / all-`NaN` window reduces to `Scalar::Null`.
fn reduce_median(
  storage : @column.ColumnStorage,
) -> ((Int, (Int) -> Int) -> @types.Scalar) raise @types.DataError {
  let present = reduce_present(storage)
  let read = numeric_double_reader(storage, "median")
  fn(n : Int, at : (Int) -> Int) {
    // Collect the present, non-`NaN` values, then sort: the median is the
    // middle by order, and `NaN` (ordered as missing here) takes no part.
    let vals : Array[Double] = []
    for k in 0.. (Int, (Int) -> Int) -> Int {
  let present = reduce_present(storage)
  let key_of : (Int) -> KeyCell = match storage.data() {
    @column.ColumnData::Int(a) => i => KInt(a[i])
    // The `Float` arm shares `float_key_cell` with `key_cell` (NaN → KNaN,
    // -0.0 → +0.0) so it keys on the native `Double` without boxing a `Scalar`
    // per cell, like the other three arms — n_unique / group_by / join stay in
    // agreement on "distinct" by construction, not by hand-synced copies.
    @column.ColumnData::Float(a) => i => float_key_cell(a[i])
    @column.ColumnData::Bool(a) => i => KBool(a[i])
    @column.ColumnData::String(a) => i => KStr(a[i])
  }
  fn(n : Int, at : (Int) -> Int) {
    let seen : Map[KeyCell, Unit] = Map([])
    for k in 0.. (Int, (Int) -> Int) -> @types.Scalar {
  let counter = count_distinct(storage)
  fn(n : Int, at : (Int) -> Int) {
    @types.Scalar::Int(counter(n, at).to_int64())
  }
}

///|
/// `first` / `last` reducer over an index window (`want_first` picks which):
/// the cell at the window's first / last position, in the window's row order,
/// keeping the source dtype. Total over every dtype, and — unlike the other
/// reductions — it skips nothing: a present cell yields its value verbatim
/// (`NaN` included), and a *null* first / last cell yields `Scalar::Null`, as
/// does an empty window. (Polars' `first` / `last` are likewise positional.)
fn reduce_first_last(
  storage : @column.ColumnStorage,
  want_first~ : Bool,
) -> (Int, (Int) -> Int) -> @types.Scalar {
  let present = reduce_present(storage)
  // Resolve the dtype dispatch once: a reader that builds the source-dtype
  // `Scalar` at a source row. Only ever called on a present row (the null /
  // empty cases short-circuit to `Null` first), so it never reads a null slot.
  let read : (Int) -> @types.Scalar = match storage.data() {
    @column.ColumnData::Int(a) => i => @types.Scalar::Int(a[i])
    @column.ColumnData::Float(a) => i => @types.Scalar::Float(a[i])
    @column.ColumnData::Bool(a) => i => @types.Scalar::Bool(a[i])
    @column.ColumnData::String(a) => i => @types.Scalar::String(a[i])
  }
  fn(n : Int, at : (Int) -> Int) {
    if n == 0 {
      @types.Scalar::Null
    } else {
      let i = at(if want_first { 0 } else { n - 1 })
      if present(i) {
        read(i)
      } else {
        @types.Scalar::Null
      }
    }
  }
}

///|
/// The eleven reductions the shared kernel resolves. Lives here in `series` and
/// is `pub(all)` so the `frame`-side aggregations can construct it. `series`
/// sits below the expression layer and does not import `internal/ir`, so it
/// cannot name the AST aggregation tag `@ir.AggOp` — it carries its own
/// reduction vocabulary instead, and `frame`'s `reduce_op_of_agg` maps
/// `@ir.AggOp` into it for both the grouped (`agg`) and scoped (`eval_agg`)
/// reductions.
///
/// This variant list mirrors `@ir.AggOp` one-to-one. Keep the two in lockstep —
/// a reduction added here must also be added to `@ir.AggOp` and to that map.
#doc(hidden)
#internal(engine, "MoonFrame execution engine API")
pub(all) enum ReduceOp {
  Sum
  Mean
  Min
  Max
  Count
  Std
  Var
  Median
  NUnique
  First
  Last
}

///|
/// Resolve a `ReduceOp` to its window reducer plus the dtype probe that fixes
/// the output column's type: `Count` / `NUnique` are always `Int`, `Mean` /
/// `Std` / `Var` / `Median` always `Float`, and `Sum` / `Min` / `Max` /
/// `First` / `Last` keep the source dtype. Shared by the grouped (`agg`) and
/// scoped (`eval_agg`) reductions so the op → reducer mapping lives in one
/// place. `Sum` / `Mean` / `Std` / `Var` / `Median` `raise` `TypeMismatch`
/// here for a non-numeric column, before any window work.
///
/// The probe is a logical `DataType`, not a physical witness column: the
/// caller (`frame`'s aggregation layer) only forwards it back into
/// `scalars_to_series`, and forwarding a dtype keeps the storage
/// representation on this side of the seam.
#doc(hidden)
#internal(engine, "MoonFrame execution engine API")
pub fn reducer_for(
  column : Series,
  op : ReduceOp,
) -> ((Int, (Int) -> Int) -> @types.Scalar, @types.DataType) raise @types.DataError {
  let storage = column.storage
  match op {
    Count => (reduce_count(storage), @types.DataType::Int)
    Sum => (reduce_arith(storage, want_mean=false), storage.dtype())
    Mean => (reduce_arith(storage, want_mean=true), @types.DataType::Float)
    Min => (reduce_extremum(storage, want_min=true), storage.dtype())
    Max => (reduce_extremum(storage, want_min=false), storage.dtype())
    Std => (reduce_variance(storage, want_std=true), @types.DataType::Float)
    Var => (reduce_variance(storage, want_std=false), @types.DataType::Float)
    Median => (reduce_median(storage), @types.DataType::Float)
    NUnique => (reduce_nunique(storage), @types.DataType::Int)
    First => (reduce_first_last(storage, want_first=true), storage.dtype())
    Last => (reduce_first_last(storage, want_first=false), storage.dtype())
  }
}

///|
/// Project per-window reduction `Scalar`s into a typed output column. `probe`
/// fixes the output dtype, so an all-null result (every cell `Null`) still
/// types correctly; each cell is projected "matching variant → `Some`,
/// anything else (including `Null`) → `None`"; and an all-valid numeric result
/// then converges onto the `Numeric` fast path (`try_column_to_numeric`, a
/// no-op for a nullable, `Bool`, or `String` column, which keeps `Builtin`).
/// This is the single column builder behind the grouped (`agg`) and scoped
/// (`eval_agg`) reductions, so both agree on backend as well as on every cell
/// value: the catch-all `None` arm is reached by the nullable reductions
/// (`mean` / `min` / `max` / `std` / `var` / `median` over an empty / all-null
/// window, and `first` / `last` over a null cell), so no arm is dead.
///
/// `probe` is the logical dtype the caller resolved (`reducer_for`'s second
/// result, the cells' own inferred dtype, or an empty-scope evaluation's), so
/// no caller outside this package has to name a storage type to build a
/// column. `Null` is not a column dtype — it is a value-level concept with no
/// physical backend — so it `raise`s `Unsupported`, exactly as
/// `Series::empty_of` does.
#doc(hidden)
#internal(engine, "MoonFrame execution engine API")
pub fn scalars_to_series(
  probe : @types.DataType,
  name : String,
  cells : Array[@types.Scalar],
) -> Series raise @types.DataError {
  let built = match probe {
    @types.DataType::Int =>
      Series::from_int_options(
        name,
        cells.map(c => {
          match c {
            @types.Scalar::Int(v) => Some(v)
            _ => None
          }
        }),
      )
    @types.DataType::Float =>
      Series::from_float_options(
        name,
        cells.map(c => {
          match c {
            @types.Scalar::Float(v) => Some(v)
            _ => None
          }
        }),
      )
    @types.DataType::Bool =>
      Series::from_bool_options(
        name,
        cells.map(c => {
          match c {
            @types.Scalar::Bool(v) => Some(v)
            _ => None
          }
        }),
      )
    @types.DataType::String =>
      Series::from_string_options(
        name,
        cells.map(c => {
          match c {
            @types.Scalar::String(v) => Some(v)
            _ => None
          }
        }),
      )
    @types.DataType::Null =>
      raise @types.DataError::Unsupported(
        "cannot build a column of dtype Null: \{name}",
      )
  }
  try_column_to_numeric(built)
}