///|
pub fn deduplicate_qrels(qrels : Array[JudgedDoc]) -> Array[JudgedDoc] {
  let merged : Map[String, JudgedDoc] = Map([])
  for item in qrels {
    let key = "\{item.query_id}:\{item.doc_id}"
    match merged.get(key) {
      Some(existing) =>
        if item.relevance > existing.relevance {
          merged[key] = item
        }
      None => merged[key] = item
    }
  }
  let result : Array[JudgedDoc] = []
  for _, item in merged {
    result.push(item)
  }
  result.sort_by(fn(a, b) {
    let query_order = a.query_id.compare(b.query_id)
    if query_order == 0 {
      a.doc_id.compare(b.doc_id)
    } else {
      query_order
    }
  })
  result
}

///|
pub fn deduplicate_run(run : Array[RetrievedDoc]) -> Array[RetrievedDoc] {
  let grouped = group_runs(run)
  let result : Array[RetrievedDoc] = []
  let query_ids : Array[String] = []
  for query_id, _ in grouped {
    query_ids.push(query_id)
  }
  query_ids.sort()
  for query_id in query_ids {
    let seen : Map[String, Unit] = Map([])
    for item in grouped[query_id] {
      if !seen.contains(item.doc_id) {
        seen[item.doc_id] = ()
        result.push(item)
      }
    }
  }
  result
}

///|
pub fn filter_qrels_by_queries(
  qrels : Array[JudgedDoc],
  query_ids : Array[String],
) -> Array[JudgedDoc] {
  let allowed : Map[String, Unit] = Map([])
  for query_id in query_ids {
    allowed[query_id] = ()
  }
  qrels.filter(fn(item) { allowed.contains(item.query_id) })
}

///|
pub fn filter_run_by_queries(
  run : Array[RetrievedDoc],
  query_ids : Array[String],
) -> Array[RetrievedDoc] {
  let allowed : Map[String, Unit] = Map([])
  for query_id in query_ids {
    allowed[query_id] = ()
  }
  run.filter(fn(item) { allowed.contains(item.query_id) })
}

///|
pub fn partition_qrels(
  qrels : Array[JudgedDoc],
  partition_count : Int,
) -> Array[QueryPartition] {
  let count = Int::max(partition_count, 1)
  let ids : Array[String] = []
  let groups : Map[String, Array[JudgedDoc]] = Map([])
  for item in qrels {
    let group = groups.get_or_init(item.query_id, fn() { [] })
    group.push(item)
  }
  for query_id, _ in groups {
    ids.push(query_id)
  }
  ids.sort()
  let partitions : Array[QueryPartition] = []
  for partition in 0.. Int {
  let mut total = 0
  for partition in partitions {
    total += partition.qrels.length()
  }
  total
}

///|
pub fn join_qrels_with_pool(
  qrels : Array[JudgedDoc],
  pools : Array[CandidatePool],
) -> Array[JudgedDoc] {
  let judged : Map[String, Int] = build_query_relevance_map(qrels)
  let result : Array[JudgedDoc] = []
  for pool in pools {
    for doc_id in pool.doc_ids {
      let key = "\{pool.query_id}:\{doc_id}"
      let relevance = judged.get_or_default(key, 0)
      if judged.contains(key) {
        result.push({ query_id: pool.query_id, doc_id, relevance })
      } else {
        result.push({ query_id: pool.query_id, doc_id, relevance: 0 })
      }
    }
  }
  result
}