///|
pub fn query_statistics(
  qrels : Array[JudgedDoc],
  run : Array[RetrievedDoc],
) -> Array[QueryStatistic] {
  let qrels_by_query = group_qrels(qrels)
  let runs_by_query = group_runs(run)
  let ids = sorted_query_ids(qrels_by_query, runs_by_query)
  let result : Array[QueryStatistic] = []
  for query_id in ids {
    let judged = qrels_by_query.get_or_default(query_id, [])
    let retrieved = runs_by_query.get_or_default(query_id, [])
    let relevance = build_relevance_map(judged)
    let mut relevant_count = 0
    for _, value in relevance {
      if value > 0 {
        relevant_count += 1
      }
    }
    let mut unjudged = 0
    for item in retrieved {
      if !relevance.contains(item.doc_id) {
        unjudged += 1
      }
    }
    result.push({
      query_id,
      judged_count: judged.length(),
      relevant_count,
      retrieved_count: retrieved.length(),
      unjudged_count: unjudged,
      mean_score: score_mean(retrieved),
    })
  }
  result
}

///|
pub fn document_statistics(
  qrels : Array[JudgedDoc],
) -> Array[DocumentStatistic] {
  let stats : Map[String, DocumentStatistic] = Map([])
  for item in qrels {
    let current = stats.get_or_default(item.doc_id, {
      doc_id: item.doc_id,
      query_frequency: 0,
      relevant_query_count: 0,
      maximum_relevance: 0,
    })
    let is_relevant = if item.relevance > 0 { 1 } else { 0 }
    stats[item.doc_id] = {
      doc_id: current.doc_id,
      query_frequency: current.query_frequency + 1,
      relevant_query_count: current.relevant_query_count + is_relevant,
      maximum_relevance: Int::max(current.maximum_relevance, item.relevance),
    }
  }
  let result : Array[DocumentStatistic] = []
  for _, item in stats {
    result.push(item)
  }
  result.sort_by(fn(a, b) { a.doc_id.compare(b.doc_id) })
  result
}

///|
pub fn candidate_recall(
  qrels : Array[JudgedDoc],
  pools : Array[CandidatePool],
) -> Double {
  let relevant : Map[String, Unit] = Map([])
  for item in qrels {
    if item.relevance > 0 {
      relevant["\{item.query_id}:\{item.doc_id}"] = ()
    }
  }
  if relevant.is_empty() {
    return 0.0
  }
  let covered : Map[String, Unit] = Map([])
  for pool in pools {
    for doc_id in pool.doc_ids {
      let key = "\{pool.query_id}:\{doc_id}"
      if relevant.contains(key) {
        covered[key] = ()
      }
    }
  }
  to_ratio(covered.length(), relevant.length())
}

///|
pub fn query_statistics_markdown(stats : Array[QueryStatistic]) -> String {
  let builder = StringBuilder()
  builder.write_string(
    "| Query | Judged | Relevant | Retrieved | Unjudged | Mean score |\n",
  )
  builder.write_string("| --- | ---: | ---: | ---: | ---: | ---: |\n")
  for item in stats {
    builder.write_string(
      "| \{item.query_id} | \{item.judged_count} | \{item.relevant_count} | \{item.retrieved_count} | \{item.unjudged_count} | \{format_metric(item.mean_score)} |\n",
    )
  }
  builder.to_string()
}