///|
fn query_id_set_from_qrels(qrels : Array[JudgedDoc]) -> Map[String, Unit] {
  let ids : Map[String, Unit] = Map([])
  for item in qrels {
    ids[item.query_id] = ()
  }
  ids
}

///|
fn query_id_set_from_run(run : Array[RetrievedDoc]) -> Map[String, Unit] {
  let ids : Map[String, Unit] = Map([])
  for item in run {
    ids[item.query_id] = ()
  }
  ids
}

///|
fn distinct_qrels_key_count(qrels : Array[JudgedDoc]) -> Int {
  let keys : Map[String, Unit] = Map([])
  for item in qrels {
    keys["\{item.query_id}:\{item.doc_id}"] = ()
  }
  keys.length()
}

///|
pub fn profile_dataset(
  qrels : Array[JudgedDoc],
  run : Array[RetrievedDoc],
) -> DatasetProfile {
  let qrel_queries = query_id_set_from_qrels(qrels)
  let run_queries = query_id_set_from_run(run)
  let all_queries : Map[String, Unit] = Map([])
  for query_id, _ in qrel_queries {
    all_queries[query_id] = ()
  }
  for query_id, _ in run_queries {
    all_queries[query_id] = ()
  }
  let relevance_map = build_query_relevance_map(qrels)
  let mut relevant_count = 0
  for _, relevance in relevance_map {
    if relevance > 0 {
      relevant_count += 1
    }
  }
  let mut unjudged_retrievals = 0
  for item in run {
    let key = "\{item.query_id}:\{item.doc_id}"
    if !relevance_map.contains(key) {
      unjudged_retrievals += 1
    }
  }
  let mut empty_query_count = 0
  for query_id, _ in all_queries {
    let has_qrels = qrel_queries.contains(query_id)
    let has_run = run_queries.contains(query_id)
    if !has_qrels || !has_run {
      empty_query_count += 1
    }
  }
  let query_count = all_queries.length()
  {
    query_count,
    judged_count: qrels.length(),
    retrieved_count: run.length(),
    relevant_count,
    run_query_coverage: if qrel_queries.length() == 0 {
      0.0
    } else {
      to_ratio(
        query_id_intersection_count(qrel_queries, run_queries),
        qrel_queries.length(),
      )
    },
    qrels_query_coverage: if run_queries.length() == 0 {
      0.0
    } else {
      to_ratio(
        query_id_intersection_count(qrel_queries, run_queries),
        run_queries.length(),
      )
    },
    mean_run_length: if run_queries.length() == 0 {
      0.0
    } else {
      Double::from_int(run.length()) / Double::from_int(run_queries.length())
    },
    mean_score: score_mean(run),
    score_stddev: score_stddev(run),
    unjudged_retrievals,
    empty_query_count,
    duplicate_query_count: qrels.length() - distinct_qrels_key_count(qrels),
  }
}

///|
fn query_id_intersection_count(
  left : Map[String, Unit],
  right : Map[String, Unit],
) -> Int {
  let mut count = 0
  for key, _ in left {
    if right.contains(key) {
      count += 1
    }
  }
  count
}

///|
pub fn classify_queries(
  report : BenchmarkReport,
  cutoff~ : Int,
) -> Array[QueryBucket] {
  let buckets : Array[QueryBucket] = []
  let metric = metric_name("ndcg", cutoff)
  for query in report.queries {
    let score = query.metrics.get_or_default(metric, 0.0)
    let label = if query.relevant_total == 0 {
      "unjudged"
    } else if score >= 0.8 {
      "strong"
    } else if score >= 0.5 {
      "mixed"
    } else {
      "weak"
    }
    buckets.push({ query_id: query.query_id, label, score })
  }
  buckets.sort_by(fn(a, b) { a.query_id.compare(b.query_id) })
  buckets
}

///|
pub fn render_profile_text(profile : DatasetProfile) -> String {
  let lines : Array[String] = [
    "MoonRAGBench dataset profile",
    "queries=\{profile.query_count} judged=\{profile.judged_count} retrieved=\{profile.retrieved_count}",
    "relevant=\{profile.relevant_count} unjudged_retrievals=\{profile.unjudged_retrievals}",
    "run_query_coverage=\{format_metric(profile.run_query_coverage)}",
    "qrels_query_coverage=\{format_metric(profile.qrels_query_coverage)}",
    "mean_run_length=\{format_metric(profile.mean_run_length)}",
    "mean_score=\{format_metric(profile.mean_score)} score_stddev=\{format_metric(profile.score_stddev)}",
    "empty_query_count=\{profile.empty_query_count} duplicate_query_count=\{profile.duplicate_query_count}",
  ]
  lines.join("\n")
}

///|
pub fn render_profile_json(profile : DatasetProfile) -> String {
  ToJson::to_json(profile).stringify(indent=2)
}