///|
pub fn StratifiedFold::index(self : StratifiedFold) -> Int {
  self.fold_index
}

///|
pub fn StratifiedFold::train_indices(self : StratifiedFold) -> Array[Int] {
  self.training_indices.copy()
}

///|
pub fn StratifiedFold::test_indices(self : StratifiedFold) -> Array[Int] {
  self.testing_indices.copy()
}

///|
fn contains_index(indices : Array[Int], target : Int) -> Bool {
  for index in indices {
    if index == target {
      return true
    }
  }
  false
}

///|
/// Distributes each sorted class round-robin while retaining source row order.
pub fn stratified_folds(
  data : Dataset,
  fold_count : Int,
) -> Result[Array[StratifiedFold], SvmError] {
  if fold_count < 2 || fold_count > data.row_count() {
    return Err(InvalidFoldCount(fold_count))
  }
  let classes = data.classes()
  for label in classes {
    let count = data.class_count(label)
    if count < fold_count {
      return Err(InsufficientClassSamples(label, count, fold_count))
    }
  }
  let tests : Array[Array[Int]] = Array::makei(fold_count, fn(_) { [] })
  let labels = data.labels()
  for label in classes {
    let mut occurrence = 0
    for row, actual in labels {
      if actual == label {
        tests[occurrence % fold_count].push(row)
        occurrence = occurrence + 1
      }
    }
  }
  let folds : Array[StratifiedFold] = []
  for fold_index = 0; fold_index < fold_count; fold_index = fold_index + 1 {
    let train : Array[Int] = []
    for row = 0; row < data.row_count(); row = row + 1 {
      if !contains_index(tests[fold_index], row) {
        train.push(row)
      }
    }
    folds.push({
      fold_index,
      training_indices: train,
      testing_indices: tests[fold_index].copy(),
    })
  }
  Ok(folds)
}