///|
pub(all) struct IndexStats {
  count : Int
  dim : Int
  extra_info : String
}

///|
pub impl Show for IndexStats with fn output(self, logger) {
  logger.write_string(
    "IndexStats{count: " +
    self.count.to_string() +
    ", dim: " +
    self.dim.to_string() +
    ", info: " +
    self.extra_info +
    "}",
  )
}

///|
pub fn FlatIndex::stats(self : FlatIndex) -> IndexStats {
  let count = self.documents.length()
  let dim = if count > 0 { self.documents[0].vector.length() } else { 0 }
  { count, dim, extra_info: "Index Type: FlatIndex" }
}

///|
pub fn IvfIndex::stats(self : IvfIndex) -> IndexStats {
  let mut count = 0
  for i = 0; i < self.k; i = i + 1 {
    count = count + self.inverted_lists[i].length()
  }
  let dim = if self.centroids.length() > 0 {
    self.centroids[0].length()
  } else {
    0
  }
  let extra_info = "Index Type: IvfIndex, Centroids(K): " + self.k.to_string()
  { count, dim, extra_info }
}

///|
pub fn KdTreeIndex::stats(self : KdTreeIndex) -> IndexStats {
  let (count, depth) = get_kd_stats(self.root)
  let dim = if count > 0 {
    match self.root {
      Node(_, doc, _, _) => doc.vector.length()
      Empty => 0
    }
  } else {
    0
  }
  let extra_info = "Index Type: KdTreeIndex, Tree Depth: " + depth.to_string()
  { count, dim, extra_info }
}

///|
fn get_kd_stats(node : KdNode) -> (Int, Int) {
  match node {
    Empty => (0, 0)
    Node(_, _, left, right) => {
      let (l_cnt, l_dep) = get_kd_stats(left)
      let (r_cnt, r_dep) = get_kd_stats(right)
      let cnt = 1 + l_cnt + r_cnt
      let dep = 1 + (if l_dep > r_dep { l_dep } else { r_dep })
      (cnt, dep)
    }
  }
}