// The reshape pair — `pivot` (long → wide: one output column per
// value of the `on` keys, one output row per distinct `index` combination)
// and `unpivot` (wide → long: the `on` columns fold into a
// `variable` / `value` pair repeated under the `index` columns).
//
// Both reshape whole frames, and both key on **columns** rather than
// computed expressions, but they differ in how a caller names those
// columns and the split follows the verbs' own shapes: `pivot`'s keys and
// values are evaluated expressions (the `group_by` / `sort` convention —
// a derived key is real here), while `unpivot`'s on / index lists are
// consulted by output name only (the `drop` / `drop_nulls` convention — a
// reshape selects columns, it does not compute them).
//
// Order is deterministic everywhere: `pivot`'s rows are the `index`
// combinations in first-appearance order and its columns the `on` values
// in first-appearance order, and `unpivot`'s rows walk the `on` columns
// outermost, one block per column holding its source rows in order —
// the verb's own docstring is the authority on the row order.

///|
/// How `pivot` collapses the cells that land in one output column. The
/// omitted-parameter case is the strict one: more than one cell in a slot
/// then `raise`s `InvalidOperation` rather than picking silently.
///
/// - `First` / `Last` — the first / last cell in row order, any dtype;
/// - `Sum` / `Min` / `Max` — the reduction family's rules (an all-null
///   slot sums to the additive identity; `min` / `max` skip `NaN`), the
///   value column's dtype kept;
/// - `Mean` / `Median` — numeric values columns only (`TypeMismatch`
///   otherwise, the reduction family's gate), the answer always `Float`,
///   an all-null slot a `Null` cell;
/// - `Len` — the **row count** that landed in the slot (nulls included),
///   as `Int`.
///
/// A slot no row landed in is a `Null` cell under every aggregation.
pub(all) enum PivotAgg {
  First
  Last
  Sum
  Min
  Max
  Mean
  Median
  Len
} derive(Eq, Debug)

///|
pub extend PivotAgg with Eq::{equal, not_equal}

///|
pub extend PivotAgg with Debug::{to_repr}

///|
/// Reshape long → wide: one output row per
/// distinct `index~` combination, one output column per distinct `on~`
/// value, each slot holding the `values~` cells of the rows that share
/// both — collapsed by `aggregate_function?` when more than one lands
/// there (`PivotAgg`'s rules; omitted, a multi-cell slot `raise`s
/// `InvalidOperation`).
///
/// `on~` / `index~` / `values~` are evaluated expressions (the `group_by`
/// convention): a bare `col(name)` keys by the column, a derived key by
/// the computed value, each named by its output name. One of
/// `index?` / `values?` must be given: `index?`
/// omitted takes every column neither `on~` nor `values~` names;
/// `values?` omitted takes every column neither `on~` nor `index~` names.
///
/// Output order is first-appearance on both axes. The `index` columns keep
/// their dtypes and head the frame;
/// each reshaped column is named by its `on` value alone when `values~`
/// holds one column, and `{value column}_{on value}` otherwise (a
/// multiple-`on` value joins its cells the same way). A name colliding
/// with another, or with an `index` column's, `raise`s `DuplicateColumn`
/// — caught at the card that renders it, before any slot is filled. A
/// `values~` resolution that comes up empty (an explicit empty list, or
/// defaults that name every column) `raise`s `InvalidOperation`: there
/// is nothing to reshape.
pub fn DataFrame::pivot(
  self : DataFrame,
  on~ : Array[@expr.Expr],
  index? : Array[@expr.Expr],
  values? : Array[@expr.Expr],
  aggregate_function? : PivotAgg,
) -> DataFrame raise @types.DataError {
  if on.is_empty() {
    raise @types.DataError::InvalidOperation("pivot needs at least one on key")
  }
  if index is None && values is None {
    raise @types.DataError::InvalidOperation(
      "pivot needs an index or a values list",
    )
  }
  // Materialise the `on` key columns (the `group_by` route: evaluate over
  // the whole frame, broadcast a length-1 key, name by output name, and
  // reject a duplicate name before any reshape work).
  let n = self.nrows()
  let scope = Array::makei(n, i => i)
  let on_cols = materialize_pivot_keys(on, self, n, scope)
  // An omitted side takes the columns the other two lists do not name —
  // resolved against the original expr lists, each into its own set.
  let skip_names = (extra : Array[@expr.Expr]) => {
    let m = names_of_exprs(on)
    for e in extra {
      m[e.output_name()] = ()
    }
    m
  }
  let index_cols = match index {
    Some(keys) => materialize_pivot_keys(keys, self, n, scope)
    None => columns_except(self, skip_names(values.unwrap()))
  }
  let values_cols = match values {
    Some(vs) => materialize_pivot_keys(vs, self, n, scope)
    None => columns_except(self, skip_names(index.unwrap()))
  }
  if values_cols.is_empty() {
    raise @types.DataError::InvalidOperation(
      "pivot resolved no values column to reshape",
    )
  }
  // Mean / Median reduce through a Float probe, so the values columns must
  // be numeric to begin with — the reduction family's gate. The projection
  // behind the probe would instead null every non-fitting cell silently,
  // which is how a String values column would lose its data whole.
  match aggregate_function {
    Some(Mean) | Some(Median) => {
      let label = if aggregate_function is Some(Mean) {
        "mean"
      } else {
        "median"
      }
      for c in values_cols {
        if !c.dtype().is_numeric() {
          raise @types.DataError::TypeMismatch(
            @types.TypeMismatchDetail::Message(
              "pivot \{label} requires a numeric values column, got \{c.dtype()} (column \{c.name()})",
            ),
          )
        }
      }
    }
    _ => ()
  }
  reshape_pivot(on_cols, index_cols, values_cols, aggregate_function)
}

///|
/// Evaluate the key / value expressions of a pivot: `group_by`'s
/// materialisation — evaluate over the whole frame, broadcast a length-1
/// key, rename by the output name, and reject a duplicate output name at
/// the second occurrence.
fn materialize_pivot_keys(
  keys : Array[@expr.Expr],
  df : DataFrame,
  n : Int32,
  scope : Array[Int32],
) -> Array[Series] raise @types.DataError {
  let seen : Map[String, Unit] = Map([])
  keys.map(key => {
    let column = @kernel.broadcast_series(eval_expr(key, df, scope), n)
    let name = key.output_name()
    if seen.contains(name) {
      raise @types.DataError::DuplicateColumn(name)
    }
    seen[name] = ()
    column.rename(name)
  })
}

///|
/// The set of output names a list of expressions produces.
fn names_of_exprs(exprs : Array[@expr.Expr]) -> Map[String, Unit] {
  let m : Map[String, Unit] = Map([])
  for e in exprs {
    m[e.output_name()] = ()
  }
  m
}

///|
/// Every column of `df` whose name is not in `skip`.
fn columns_except(df : DataFrame, skip : Map[String, Unit]) -> Array[Series] {
  df.columns.filter(c => !skip.contains(c.name()))
}

///|
/// The reshape itself: bucket every row by its (index combination, on
/// value) pair, then emit one row per index combination and one column per
/// (on value × value column), each slot reduced per the aggregation. Both
/// axes keep first-appearance order; an empty `index` list is the grand
/// total (every row one combination, `group_by([])`'s reading).
fn reshape_pivot(
  on_cols : Array[Series],
  index_cols : Array[Series],
  values_cols : Array[Series],
  agg : PivotAgg?,
) -> DataFrame raise @types.DataError {
  let n = on_cols[0].len()
  let on_scalars = on_cols.map(c => c.to_scalars())
  let on_key_cells = on_cols.map(c => c.key_cells())
  let index_key_cells = index_cols.map(c => c.key_cells())
  // The names a card's column will finally carry — the single-value shape
  // names it by itself, the multi-value shape prefixes each value column —
  // and the index names they must not collide with. Both checked as the
  // card renders, not at the schema the whole reshape would have to run
  // to reach.
  let index_names : Map[String, Unit] = Map([])
  for c in index_cols {
    index_names[c.name()] = ()
  }
  let seen_names : Map[String, Unit] = Map([])
  let multi_values = values_cols.length() > 1
  // Row → (group id, on-value id), both ids first-appearance; the group's
  // representative row and the card's rendered name record alongside.
  let gid_of : Map[Array[KeyCell], Int32] = Map([])
  let cid_of : Map[Array[KeyCell], Int32] = Map([])
  let gids : Array[Int32] = Array::makei(n, _ => 0)
  let cids : Array[Int32] = Array::makei(n, _ => 0)
  let rep_rows : Array[Int32] = []
  let on_names : Array[String] = []
  for i in 0.. {
          match col[i] {
            @types.Scalar::Null => "null"
            cell => cell.to_string()
          }
        })
        .join("_")
      if multi_values {
        for v in values_cols {
          let name = "\{v.name()}_\{card}"
          if index_names.contains(name) || seen_names.contains(name) {
            raise @types.DataError::DuplicateColumn(name)
          }
          seen_names[name] = ()
        }
      } else {
        if index_names.contains(card) || seen_names.contains(card) {
          raise @types.DataError::DuplicateColumn(card)
        }
        seen_names[card] = ()
      }
      on_names.push(card)
    }
    cids[i] = cid_of[ckey]
  }
  let groups = rep_rows.length()
  let cards = on_names.length()
  // Buckets: cells[group][card][value column], keyed sparsely — a
  // `(group, card)` pair that no row hits owns no array at all, so a
  // sparse long-table pivot (10⁴ groups × 10⁴ cards with 10⁵ rows) pays
  // for the pairs that exist, not the full grid. The slot reader below
  // answers an absent pair with the empty bucket.
  let buckets : Map[Int64, Array[Array[@types.Scalar]]] = Map([])
  fn bucket_at(g : Int32, c : Int32) -> Array[Array[@types.Scalar]] {
    let key = g.to_int64() * cards.to_int64() + c.to_int64()
    match buckets.get(key) {
      Some(bucket) => bucket
      None => {
        let fresh = Array::makei(values_cols.length(), _ => [])
        buckets[key] = fresh
        fresh
      }
    }
  }
  let value_scalars = values_cols.map(c => c.to_scalars())
  for i in 0.. c.gather_series(rep_rows))
  // One column per (card × value column), named per the naming rule.
  // Value-major: each value column's cards together.
  for v in 0.. @types.DataType::Float
        Some(Len) => @types.DataType::Int
        // The keepers: `First` / `Last` / `Sum` / `Min` / `Max` output the
        // source column's own dtype, and so does no aggregation at all.
        None | Some(First) | Some(Last) | Some(Sum) | Some(Min) | Some(Max) =>
          values_cols[v].dtype()
      }
      let cells : Array[@types.Scalar] = []
      for g in 0.. bucket[v]
          None => []
        }
        cells.push(reduce_pivot_slot(slot, agg, probe))
      }
      out_cols.push(@series.scalars_to_series(probe, name, cells))
    }
  }
  DataFrame::from_parts(out_cols, groups)
}

///|
/// Collapse one pivot slot's cells under the aggregation: `None` raises on
/// a multi-cell slot (the default strictness), `Len` counts rows, and
/// the reductions go through a same-dtype series — the caller gates the
/// dtypes `Mean` / `Median` need, since the probe here projects rather
/// than coerces — so every rule the reduction family documents applies
/// verbatim. An empty slot is a `Null` cell under every aggregation.
fn reduce_pivot_slot(
  cells : Array[@types.Scalar],
  agg : PivotAgg?,
  probe : @types.DataType,
) -> @types.Scalar raise @types.DataError {
  if cells.is_empty() {
    return @types.Scalar::Null
  }
  // A `Float` probe over `Int` cells (the mean / median / unified-value
  // shape) widens first — the probe projects, it does not coerce.
  let as_series = () => {
    @series.scalars_to_series(
      probe,
      "slot",
      cells.map(c => widen_unpivot_cell(c, probe)),
    )
  }
  match agg {
    None =>
      if cells.length() > 1 {
        raise @types.DataError::InvalidOperation(
          "pivot slot holds \{cells.length()} values; pass an aggregate_function",
        )
      } else {
        cells[0]
      }
    Some(Len) => @types.Scalar::Int(cells.length().to_int64())
    Some(First) => cells[0]
    Some(Last) => cells[cells.length() - 1]
    Some(Sum) => as_series().sum()
    Some(Min) => as_series().min()
    Some(Max) => as_series().max()
    Some(Mean) =>
      match as_series().mean_opt() {
        Some(m) => @types.Scalar::Float(m)
        None => @types.Scalar::Null
      }
    Some(Median) =>
      if as_series().count() == 0 {
        @types.Scalar::Null
      } else {
        @types.Scalar::Float(as_series().median())
      }
  }
}

///|
/// Reshape wide → long: each `on?`
/// column folds into a `variable_name?` / `value_name?` pair — one output
/// row per (source row, on column), the `index?` columns repeated beside.
/// `on?` and `index?` are consulted by output name only (the `drop`
/// convention); `on?` omitted takes every column `index?` does not name,
/// `index?` omitted keeps nothing, and an empty `on` resolution `raise`s
/// `InvalidOperation`. The output rows walk the `on` columns outermost —
/// one block per column, each holding its source rows in order — so the
/// index columns repeat within a block.
///
/// The value column carries one dtype: the `on` columns' own dtype when
/// they agree, `Float` when they are `Int` / `Float` in any mix (`Int`
/// cells widen), and any other mix `raise`s `TypeMismatch` — MoonFrame has
/// no stringly supertype to fall back on.
///
/// The two lists play disjoint roles — a named column cannot be both an
/// `on` and an `index` column (`InvalidOperation`) — and the generated
/// `variable_name` / `value_name` must differ from each other and from
/// every index name (`DuplicateColumn`); all three are entry facts,
/// checked before any row is reshaped.
pub fn DataFrame::unpivot(
  self : DataFrame,
  on? : Array[@expr.Expr],
  index? : Array[@expr.Expr],
  variable_name~ : String = "variable",
  value_name~ : String = "value",
) -> DataFrame raise @types.DataError {
  // Name-only resolution: the reshape selects columns, it does not compute.
  let index_cols = match index {
    Some(keys) => self.resolve_subset_columns(keys)
    None => []
  }
  let index_names : Map[String, Unit] = Map([])
  for c in index_cols {
    index_names[c.name()] = ()
  }
  let on_cols = match on {
    Some(keys) => self.resolve_subset_columns(keys)
    None => self.columns.filter(c => !index_names.contains(c.name()))
  }
  if on_cols.is_empty() {
    raise @types.DataError::InvalidOperation(
      "unpivot needs at least one on column",
    )
  }
  // A column folding into `value` while also standing beside it as an
  // index would appear twice over; the generated names land in the same
  // namespace as the index names.
  for c in on_cols {
    if index_names.contains(c.name()) {
      raise @types.DataError::InvalidOperation(
        "unpivot column \{c.name()} cannot be both an on and an index column",
      )
    }
  }
  if variable_name == value_name {
    raise @types.DataError::DuplicateColumn(variable_name)
  }
  for name in [variable_name, value_name] {
    if index_names.contains(name) {
      raise @types.DataError::DuplicateColumn(name)
    }
  }
  // The value dtype: agreement, or the Int / Float widening. The first
  // on-column's name rides along, so a mismatch on a wide frame names the
  // two columns that disagree instead of only their dtypes.
  let first_on = on_cols[0].name()
  let mut value_dtype = on_cols[0].dtype()
  for c in on_cols {
    let d = c.dtype()
    if d != value_dtype {
      if d.is_numeric() && value_dtype.is_numeric() {
        value_dtype = @types.DataType::Float
      } else {
        raise @types.DataError::TypeMismatch(
          @types.TypeMismatchDetail::Message(
            "unpivot on-columns must share a dtype, got \{value_dtype} (column \{first_on}) and \{d} (column \{c.name()})",
          ),
        )
      }
    }
  }
  let n = self.nrows()
  let k = on_cols.length()
  // The output height is `n * k`. A product past 2^31 would wrap, but the
  // wrap is unreachable before it: every output column below materialises
  // `n * k` cells first, and allocation fails long before the count
  // wraps — so there is deliberately no guard here, only the comment (a
  // guard would be an untestable dead arm).
  // Index columns repeat once per on-column block, in step with the value
  // blocks below: the row at output position `j * n + i` carries source
  // row `i`'s index cells, so the index and value columns stay aligned.
  let out_cols : Array[Series] = []
  for c in index_cols {
    let repeated : Array[@types.Scalar] = []
    let cells = c.to_scalars()
    for _ in 0.. @types.Scalar {
  match (cell, value_dtype) {
    (@types.Scalar::Int(i), @types.DataType::Float) =>
      @types.Scalar::Float(i.to_double())
    _ => cell
  }
}