// Load FLDOE Membership by School by Grade by Race → pct_black, pct_hispanic,
// pct_white, pct_asian, pct_two_or_more, pct_amer_indian, pct_hawaiian_pac_isl
// as shares of the same year's October enrollment (enrollment_oct).
//
// The file has no total column and hides cells under 10 students ('*', or
// '**' in 2018-19 and 2021-22). Blank cells are grades not served (zero).
// The school's unallocated count U = enrollment_oct − visible cells is
// exact, so the question is only how to spread U over the hidden cells.
// Rules, chosen by measuring against the unsuppressed NCES CCD copy
// (docs/cities/miami/probe/race_imputation.py and race_imputation_prior.py;
// findings F.15):
//   1. hidden cells in a grade whose by-grade total is already fully visible
//      are zeros;
//   2. in a year whose tokens can only hide students (no school has more
//      hidden cells S than hidden students U), U is split equally over the
//      remaining hidden cells — mean error vs CCD 0.8–0.9 pp on the Black and
//      Hispanic shares in 2018-19 and 2023-24;
//   3. in a year where some school has S > U — the 2022-23 file marks zeros
//      with '*' too (257 such schools) — U is split over the hidden groups in
//      proportion to the school's mean share in the adjacent years that pass
//      rule 2 (their rule-2 shares), so groups the school never has get
//      nothing, and no group receives more than 9 students per hidden cell
//      (the publisher's own floor); the overflow is re-spread over the other
//      hidden groups by prior, then equally. Measured 2022-23 error vs CCD:
//      Black 0.89 pp (equal split 2.39, uncapped prior 1.21), Hispanic 0.66
//      (2.15), White 0.84 (1.73), American Indian 0.08 (1.64); every school-
//      year inside the process bound; the small-group medians return to
//      0.000 as in every other year. Applying the prior in rule-2 years was
//      measured worse (Black 1.10 vs 0.91) and is not done. Schools with no
//      adjacent rule-2 year keep the equal split.
//      The rule was chosen among five candidates by its mean error against the
//      CCD 2022-23 race subtotals (docs/cities/miami/probe/race_imputation_prior.py:
//      Black 2.39 → 0.89 pp, Hispanic 2.15 → 0.66, White 1.73 → 0.84), the same
//      reference the harness reports it against, so that number is in-sample.
//      The probe's split-half check (2026-09-16b) chooses the same rule on either
//      half of the schools and its held-out error equals the in-sample error of
//      the other half (0.78 / 0.83 pp), so the selection optimism is negligible.
// Schools whose October total is itself suppressed get suppressed shares.
// The harness re-measures the CCD agreement every run
// (pipeline_miami/verify/cross_publisher_race.py).
// Run:  npx tsx scripts/loaders/miami/race_membership.ts

import fs from "node:fs";
import { prisma, requireSource, startDataLoad, finishDataLoad, padSchoolNumber, parseCell, schoolUniverse, upsertFacts, sheetRows, runMain, SOURCES_ROOT, type Fact, type Provenance } from "./_lib";

const SCRIPT_PATH = "scripts/loaders/miami/race_membership.ts";
const norm = (s: unknown) => String(s ?? "").toLowerCase().replace(/[^a-z0-9]+/g, " ").trim();
const RACE_KEYS: [RegExp, string][] = [
  [/^white/, "pct_white"], [/^black/, "pct_black"], [/^hispanic/, "pct_hispanic"], [/^asian/, "pct_asian"],
  [/^native hawaiian/, "pct_hawaiian_pac_isl"], [/^american indian/, "pct_amer_indian"], [/^two or more/, "pct_two_or_more"],
];
const KEYS = RACE_KEYS.map(([, k]) => k);

/** Grade labels differ by file: "KG", "PK", "01".."12" or "1".."12". */
function normGrade(g: string): string {
  const t = g.trim().toUpperCase();
  if (t === "PK" || t === "KG") return t;
  return /^\d+$/.test(t) ? String(Number(t)) : t;
}

function quantile(xs: number[], q: number): number {
  if (!xs.length) return NaN;
  const s = [...xs].sort((a, b) => a - b);
  return s[Math.min(s.length - 1, Math.floor(q * s.length))];
}

type Acc = { visible: Record<string, number>; hidden: Record<string, number>; hiddenZero: number; hiddenCells: number; cells: number };
type SchoolYear = { total: number; U: number; S: number; acc: Acc; equal: Record<string, number> };   // equal = rule-2 shares (0–100)
type YearData = { sourceId: string; year: string; filePath: string; prov: Provenance; dade: number; sums: Map<string, Acc>; enr: Map<string, number>; cells: Map<string, SchoolYear>; noTotal: string[]; zerosMarked: number };

async function parseYear(sourceId: string, universe: Set<string>): Promise<YearData | null> {
  const year = sourceId.slice(-7);
  const { filePath, prov } = requireSource(sourceId);
  const { rows } = sheetRows(filePath, "School");
  const h = rows.findIndex((r) => norm(r?.[0]) === "district");
  if (h < 0) { console.warn(`[race ${year}] header not found; skipped`); return null; }
  const header = rows[h].map(norm);
  const iSch = header.indexOf("school");
  const cols = RACE_KEYS.map(([re, key]) => ({ key, i: header.findIndex((c) => re.test(c)) })).filter((c) => c.i >= 0);
  if (cols.length !== RACE_KEYS.length) { console.warn(`[race ${year}] race columns not found in ${JSON.stringify(header)}; skipped`); return null; }
  const iGrade = header.indexOf("grade");
  const enrRows = await prisma.miamiSchoolYearMetric.findMany({ where: { year, subgroup: "ALL", suppressed: false, OR: [{ metricKey: "enrollment_oct" }, { metricKey: { startsWith: "enrollment_oct_g" } }] }, select: { schoolNumber: true, metricKey: true, value: true } });
  const enr = new Map<string, number>();
  const enrGrade = new Map<string, number>();
  for (const r of enrRows) {
    if (r.metricKey === "enrollment_oct") enr.set(r.schoolNumber, r.value ?? 0);
    else enrGrade.set(`${r.schoolNumber}|${normGrade(r.metricKey.slice("enrollment_oct_g".length))}`, r.value ?? 0);
  }
  const sums = new Map<string, Acc>();
  let dade = 0;
  for (const r of rows.slice(h + 1)) {
    if (String(r[0] ?? "").trim().replace(/\.0$/, "") !== "13") continue;
    dade++;
    const num = padSchoolNumber(r[iSch]);
    if (!num || !universe.has(num)) continue;
    const s = sums.get(num) ?? { visible: {}, hidden: {}, hiddenZero: 0, hiddenCells: 0, cells: 0 };
    // Per grade: visible sum and the hidden cells; if the grade's published total equals the visible sum, the hidden cells are zeros (rule 1).
    let vis = 0;
    const hiddenHere: string[] = [];
    for (const c of cols) {
      const raw = String(r[c.i] ?? "").trim();
      s.cells++;
      if (raw === "") continue;                       // grade not served → zero
      const p = parseCell(raw);
      if (p.suppressed) hiddenHere.push(c.key);
      else { s.visible[c.key] = (s.visible[c.key] ?? 0) + (p.value ?? 0); vis += p.value ?? 0; }
    }
    const gTotal = iGrade >= 0 ? enrGrade.get(`${num}|${normGrade(String(r[iGrade] ?? ""))}`) : undefined;
    const gradeFullyVisible = gTotal !== undefined && Math.abs(gTotal - vis) < 0.5;
    for (const k of hiddenHere) {
      if (gradeFullyVisible) s.hiddenZero++;
      else { s.hidden[k] = (s.hidden[k] ?? 0) + 1; s.hiddenCells++; }
    }
    sums.set(num, s);
  }
  // Rule-2 (equal split) shares and the token diagnostic per school.
  const cells = new Map<string, SchoolYear>();
  const noTotal: string[] = [];
  let zerosMarked = 0;
  for (const [num, acc] of sums) {
    const total = enr.get(num);
    if (!total) { noTotal.push(num); continue; }
    const numeric = Object.values(acc.visible).reduce((a, b) => a + b, 0);
    const U = Math.max(0, total - numeric);
    const S = acc.hiddenCells;
    if (S > U) zerosMarked++;
    const perCell = S > 0 ? U / S : 0;
    const equal: Record<string, number> = {};
    for (const k of KEYS) equal[k] = (((acc.visible[k] ?? 0) + perCell * (acc.hidden[k] ?? 0)) / total) * 100;
    cells.set(num, { total, U, S, acc, equal });
  }
  return { sourceId, year, filePath, prov, dade, sums, enr, cells, noTotal, zerosMarked };
}

async function main() {
  const dirs = fs.readdirSync(SOURCES_ROOT).filter((d) => /^fldoe_membership_by_school_grade_race_\d{4}-\d{2}$/.test(d)).sort();
  const universe = await schoolUniverse();
  const years: YearData[] = [];
  for (const sourceId of dirs) { const y = await parseYear(sourceId, universe); if (y) years.push(y); }
  for (let i = 0; i < years.length; i++) {
    const yd = years[i];
    const trustworthy = yd.zerosMarked === 0;
    // Rule 3 prior: the loaded years one CALENDAR year either side (not the list neighbours, so a year that failed to
    // parse cannot make a non-adjacent year the prior) whose tokens hide students only.
    const spring = (y: string) => parseInt(y.slice(0, 4), 10);
    const adjacent = years.filter((a) => Math.abs(spring(a.year) - spring(yd.year)) === 1 && a.zerosMarked === 0);
    const facts: Fact[] = [];
    const unallocated: number[] = [];
    let hiddenZeroTotal = 0, hiddenCellsTotal = 0, priorUsed = 0, priorFallback = 0, capViolations = 0;
    for (const [num, c] of yd.cells) {
      unallocated.push((c.U / c.total) * 100);
      hiddenZeroTotal += c.acc.hiddenZero; hiddenCellsTotal += c.S;
      let share = c.equal;
      if (!trustworthy && c.S > 0 && c.U > 0) {
        const priors = adjacent.map((a) => a.cells.get(num)?.equal).filter((p): p is Record<string, number> => !!p);
        const hiddenKeys = KEYS.filter((k) => (c.acc.hidden[k] ?? 0) > 0);
        const w: Record<string, number> = {};
        for (const k of hiddenKeys) w[k] = priors.length ? priors.reduce((s, p) => s + p[k], 0) / priors.length : NaN;
        const wSum = hiddenKeys.reduce((s, k) => s + (Number.isFinite(w[k]) ? w[k] : 0), 0);
        if (priors.length && wSum > 0) {
          // Prior-proportional allocation, capped at 9 students per hidden cell; overflow re-spread by prior, then equally.
          const alloc: Record<string, number> = Object.fromEntries(hiddenKeys.map((k) => [k, 0]));
          const limit: Record<string, number> = Object.fromEntries(hiddenKeys.map((k) => [k, 9 * (c.acc.hidden[k] ?? 0)]));
          let remaining = c.U;
          let active = hiddenKeys.filter((k) => w[k] > 0);
          for (let iter = 0; iter <= hiddenKeys.length && remaining > 1e-9 && active.length; iter++) {
            const aSum = active.reduce((s, k) => s + w[k], 0);
            for (const k of active) { const take = Math.min(remaining * w[k] / aSum, limit[k] - alloc[k]); alloc[k] += take; }
            remaining = c.U - hiddenKeys.reduce((s, k) => s + alloc[k], 0);
            active = active.filter((k) => alloc[k] < limit[k] - 1e-9);
          }
          // Overflow (every prior>0 cell at its cap): to the still-uncapped hidden cells — those with a zero prior — by
          // cell count, iterating against the cap; only if every hidden cell is capped (U > 9 × S, a file inconsistency)
          // does the rest go over all hidden cells by cell count, and that case is counted.
          for (let iter = 0; iter <= hiddenKeys.length && remaining > 1e-9; iter++) {
            const open = hiddenKeys.filter((k) => alloc[k] < limit[k] - 1e-9);
            if (!open.length) break;
            const hs = open.reduce((s, k) => s + (c.acc.hidden[k] ?? 0), 0);
            for (const k of open) { const take = Math.min(remaining * (c.acc.hidden[k] ?? 0) / hs, limit[k] - alloc[k]); alloc[k] += take; }
            remaining = c.U - hiddenKeys.reduce((s, k) => s + alloc[k], 0);
          }
          if (remaining > 1e-9) { capViolations++; const hs = hiddenKeys.reduce((s, k) => s + (c.acc.hidden[k] ?? 0), 0); for (const k of hiddenKeys) alloc[k] += remaining * (c.acc.hidden[k] ?? 0) / hs; }
          share = {};
          for (const k of KEYS) share[k] = (((c.acc.visible[k] ?? 0) + (alloc[k] ?? 0)) / c.total) * 100;
          priorUsed++;
        } else priorFallback++;
      }
      for (const k of KEYS) facts.push({ schoolNumber: num, year: yd.year, metricKey: k, value: Math.min(100, share[k]), denominator: c.total });
    }
    // Schools whose October total is suppressed: shares cannot be computed; serve them suppressed (never a stale CCD-derived value).
    for (const num of yd.noTotal) for (const k of KEYS) facts.push({ schoolNumber: num, year: yd.year, metricKey: k, value: null, suppressed: true });
    const rule = trustworthy
      ? "equal split of U over hidden cells (no school has more hidden cells than hidden students)"
      : `adjacent-year prior capped at 9 per hidden cell (${adjacent.map((a) => a.year).join(", ") || "none available"}) — ${yd.zerosMarked} schools have more hidden cells than hidden students, so '*' marks zeros too; prior applied to ${priorUsed} schools, equal-split fallback ${priorFallback}, schools whose unallocated count exceeds 9 × hidden cells (cap unavoidable) ${capViolations}`;
    const note = `Dade rows ${yd.dade}; schools ${yd.sums.size}; without same-year enrollment_oct (served suppressed) ${yd.noTotal.length}; hidden cells ${hiddenCellsTotal} (+${hiddenZeroTotal} resolved as zeros from grade totals) of ${[...yd.sums.values()].reduce((a, s) => a + s.cells, 0)}; unallocated share pp median ${quantile(unallocated, 0.5).toFixed(2)} p90 ${quantile(unallocated, 0.9).toFixed(2)} max ${quantile(unallocated, 1).toFixed(2)}; rule: ${rule}`;
    const load = await startDataLoad({ sourceId: yd.sourceId, sourceUrl: yd.prov.url, sha256: yd.prov.sha256, bytes: yd.prov.bytes, yearCovered: yd.year, scriptPath: SCRIPT_PATH, notes: note });
    const res = await upsertFacts(facts, load.id);
    await finishDataLoad(load.id, { inserted: res.written, updated: 0, suppressed: res.suppressed });
    console.log(`[race ${yd.year}] ${res.written} facts; ${note}`);
  }
}

runMain(main);