/**
 * The district-wide arrest total, taken from the posterior predictive draws.
 *
 * Why this exists: the time-series chart used to draw its band by summing each
 * student group's own `count_lower`/`count_upper`. That is not the interval of
 * the total. The sum of per-group 97.5th percentiles is the value you would see
 * if *every* group simultaneously landed at its own extreme in the same draw,
 * which is far less likely than any one group doing so — so the band came out
 * systematically too wide. Summing within each draw and taking quantiles of the
 * resulting totals answers the actual question: how many arrests does the model
 * think this district had?
 *
 * This is `build_state_summary()`'s method from crdc-arrests/R/summarize_draws.R
 * applied on a different axis — sum inside the draw, summarize across draws —
 * the same operation `poolBySex` performs for sex pooling.
 *
 * One caveat to carry into any caption. Upstream, `draw_id` is renumbered per
 * write batch, so draw *k* of one group is not the same posterior sample as
 * draw *k* of another; measured cross-group correlation is ≈ 0.02. Summing at a
 * draw index therefore convolves what are effectively independent draws. That
 * is a reasonable model here — these are *posterior predictive* draws, whose
 * observation noise is independent across groups by construction and dominates
 * the parameter-level covariance — but it does mean the result should be read
 * as a predictive total, not as a contrast that preserves parameter
 * correlation. It is still much closer to the truth than summing bounds.
 */

import { quantile } from './kde.js'
import { isCompleteDrawSet } from './drawGroups.js'

/** Matches the API's default interval, so the two are directly comparable. */
export const TOTAL_INTERVAL_MASS = 0.95

/**
 * District total at each draw index: the sum across every student group.
 *
 * Returns null unless *every* group has a complete draw set. A group silently
 * missing from the sum would undercount the total at every draw and shift the
 * whole interval down, which is indistinguishable on screen from a real result.
 *
 * @param {Record<string, number[]> | null | undefined} countsByGroup
 * @param {number} nDraws
 * @returns {number[] | null}
 */
export function totalPerDraw(countsByGroup, nDraws) {
  const groups = Object.values(countsByGroup || {})
  if (!groups.length || !(nDraws > 0)) return null
  if (!groups.every((counts) => isCompleteDrawSet(counts, nDraws))) return null

  const totals = new Array(nDraws).fill(0)
  for (const counts of groups) {
    for (let i = 0; i < nDraws; i++) totals[i] += counts[i]
  }
  return totals
}

/**
 * @param {number[] | null | undefined} totals - district total per draw
 * @returns {{lower: number, median: number, upper: number, nDraws: number} | null}
 */
export function totalInterval(totals) {
  if (!totals?.length) return null
  const tail = (1 - TOTAL_INTERVAL_MASS) / 2
  return {
    lower: quantile(totals, tail),
    median: quantile(totals, 0.5),
    upper: quantile(totals, 1 - tail),
    nDraws: totals.length,
  }
}