///|
fn samples_by_query(
  samples : Array[NegativeSample],
) -> Map[String, Array[NegativeSample]] {
  let grouped : Map[String, Array[NegativeSample]] = Map([])
  for sample in samples {
    grouped.get_or_init(sample.query_id, fn() { [] }).push(sample)
  }
  grouped
}

///|
pub fn sample_with_budget(
  qrels : Array[JudgedDoc],
  pools : Array[CandidatePool],
  total_budget~ : Int,
) -> Array[NegativeSample] {
  let budget = Int::max(total_budget, 0)
  if budget == 0 || pools.is_empty() {
    return []
  }
  let candidates = sample_pool_negatives(
    qrels,
    pools,
    config=NegativeSampleConfig::new(per_query=budget, skip_judged=true),
  )
  let grouped = samples_by_query(candidates)
  let query_ids : Array[String] = []
  for query_id, _ in grouped {
    query_ids.push(query_id)
  }
  query_ids.sort()
  let result : Array[NegativeSample] = []
  let mut round = 0
  let mut active = true
  while active && result.length() < budget {
    active = false
    for query_id in query_ids {
      let bucket = grouped[query_id]
      if round < bucket.length() && result.length() < budget {
        result.push(bucket[round])
        active = true
      }
    }
    round += 1
  }
  result
}

///|
pub fn plan_sampling(
  qrels : Array[JudgedDoc],
  pools : Array[CandidatePool],
  total_budget~ : Int,
) -> String {
  let profile = profile_candidate_pools(pools)
  let coverage = candidate_recall(qrels, pools)
  let lines : Array[String] = [
    "budget=\{Int::max(total_budget, 0)}",
    "pools=\{profile.pool_count}",
    "candidates=\{profile.unique_candidate_count}",
    "duplicates=\{profile.duplicate_candidate_count}",
    "candidate_recall=\{format_metric(coverage)}",
    "allocation=round_robin_by_query",
  ]
  lines.join("\n")
}

///|
pub fn sample_hard_and_tail(
  qrels : Array[JudgedDoc],
  pools : Array[CandidatePool],
  per_query~ : Int,
  window~ : Int,
) -> Array[NegativeSample] {
  let head = sample_pool_negatives(
    qrels,
    pools,
    config=NegativeSampleConfig::new(
      per_query~,
      strategy=NegativeStrategy::hard_window(window),
    ),
  )
  let tail = sample_pool_negatives(
    qrels,
    pools,
    config=NegativeSampleConfig::new(
      per_query~,
      strategy=NegativeStrategy::tail(window),
    ),
  )
  let merged : Map[String, NegativeSample] = Map([])
  for sample in head {
    merged["\{sample.query_id}:\{sample.doc_id}"] = sample
  }
  for sample in tail {
    merged["\{sample.query_id}:\{sample.doc_id}"] = sample
  }
  let result : Array[NegativeSample] = []
  for _, sample in merged {
    result.push(sample)
  }
  result.sort_by(fn(a, b) {
    let by_query = a.query_id.compare(b.query_id)
    if by_query == 0 {
      a.source_rank - b.source_rank
    } else {
      by_query
    }
  })
  result
}

///|
pub fn sampling_yield(
  qrels : Array[JudgedDoc],
  pools : Array[CandidatePool],
  samples : Array[NegativeSample],
) -> Double {
  let possible = profile_candidate_pools(pools).unique_candidate_count
  let relevant_blocked = candidate_recall(qrels, pools)
  if possible == 0 {
    0.0
  } else {
    let adjusted = Double::from_int(samples.length()) /
      Double::from_int(possible)
    Double::min(adjusted * (1.0 - relevant_blocked), 1.0)
  }
}