// Analysis for the Vellora Research report "Rural vs. Metropolitan Health
// Disparities Across Texas Counties". Pure and deterministic. It implements
// the prespecified plan in public/data/research/rural-metro-analysis-plan.v1.md
// (committed before any group difference was computed).
//
// Results are written to public/data/research/rural-metro-results.v1.json by
// `npm run research:rural`; tests/rural-study.test.ts fails if that file is
// stale, and scripts/verify-rural-study.py recomputes it with SciPy.

import { jStat } from "jstat";
import { benjaminiHochberg, describe, quantile, spearman } from "../stats";
import { mannWhitney } from "./cardio-study";

export const STUDY_ID = "rural-metro-2026";

export type IndId = "HDMORT" | "LUNG" | "CHD" | "BPHIGH" | "DIABETES" | "OBESITY" | "CSMOKING" | "LPA";
export type SecondaryId = Exclude<IndId, "HDMORT" | "LUNG">;
export const PRIMARY: IndId[] = ["HDMORT", "LUNG"];
export const SECONDARY: SecondaryId[] = ["CHD", "BPHIGH", "DIABETES", "OBESITY", "CSMOKING", "LPA"];
export const INDICATORS: IndId[] = [...PRIMARY, ...SECONDARY];

export const META: Record<IndId, { name: string; short: string; unit: string; source: string; period: string; adjustment: string }> = {
  HDMORT: { name: "Heart disease death rate, adults 35+", short: "Heart disease deaths", unit: "per 100,000", source: "CDC heart disease mortality (DHDSP)", period: "Deaths 2022–2024", adjustment: "Age-standardized (2000 U.S. standard); spatially smoothed" },
  LUNG: { name: "Lung and bronchus cancer incidence", short: "Lung cancer incidence", unit: "per 100,000", source: "State Cancer Profiles (Texas Cancer Registry)", period: "Diagnoses 2019–2023", adjustment: "Age-adjusted (2000 U.S. standard)" },
  CHD: { name: "Coronary heart disease among adults", short: "Coronary heart disease", unit: "%", source: "CDC PLACES 2025", period: "BRFSS 2023", adjustment: "Age-adjusted prevalence, adults 18+ (model-based)" },
  BPHIGH: { name: "High blood pressure among adults", short: "High blood pressure", unit: "%", source: "CDC PLACES 2025", period: "BRFSS 2023", adjustment: "Age-adjusted prevalence, adults 18+ (model-based)" },
  DIABETES: { name: "Diagnosed diabetes among adults", short: "Diabetes", unit: "%", source: "CDC PLACES 2025", period: "BRFSS 2023", adjustment: "Age-adjusted prevalence, adults 18+ (model-based)" },
  OBESITY: { name: "Obesity among adults", short: "Obesity", unit: "%", source: "CDC PLACES 2025", period: "BRFSS 2023", adjustment: "Age-adjusted prevalence, adults 18+ (model-based)" },
  CSMOKING: { name: "Current cigarette smoking among adults", short: "Current smoking", unit: "%", source: "CDC PLACES 2025", period: "BRFSS 2023", adjustment: "Age-adjusted prevalence, adults 18+ (model-based)" },
  LPA: { name: "No leisure-time physical activity among adults", short: "No leisure-time activity", unit: "%", source: "CDC PLACES 2025", period: "BRFSS 2023", adjustment: "Age-adjusted prevalence, adults 18+ (model-based)" },
};

type Series = (number | null)[];
export interface Cell {
  values: Series;
  reasons?: Record<string, string>;
}
export interface RuralStudyInputs {
  countyOrder: string[];
  counties: { fips: string; name: string; rucc: number; metro: boolean; pop2020: number }[];
  /** Age-adjusted values (for HDMORT and LUNG, the only published values). */
  primary: Record<IndId, Cell>;
  /** PLACES crude prevalence for the secondary measures (sensitivity analysis E1). */
  crude: Record<SecondaryId, Cell>;
}

const round = (v: number, d = 4) => Number(v.toFixed(d));

/**
 * Hodges–Lehmann shift estimate (median of all a − b differences) with the
 * distribution-free CI from the order statistics of those differences:
 * C = ⌊n_a·n_b/2 − z·√(n_a·n_b·(n_a+n_b+1)/12)⌋, CI = [D(C), D(N+1−C)] (1-indexed).
 * Normal approximation for C, no tie correction.
 */
export function hodgesLehmann(a: number[], b: number[], level = 0.95) {
  const d: number[] = [];
  for (const x of a) for (const y of b) d.push(x - y);
  d.sort((p, q) => p - q);
  const N = d.length;
  const z = jStat.normal.inv(1 - (1 - level) / 2, 0, 1);
  const C = Math.floor((a.length * b.length) / 2 - z * Math.sqrt((a.length * b.length * (a.length + b.length + 1)) / 12));
  return { estimate: quantile(d, 0.5), ci: C >= 1 ? ([d[C - 1], d[N - C]] as [number, number]) : null };
}

const lnFact = (n: number) => jStat.gammaln(n + 1);
/** Two-sided Fisher's exact test for [[a, b], [c, d]]: sums tables no more probable than the observed one (relative tolerance 1e-7, as SciPy). */
export function fisherExact(a: number, b: number, c: number, d: number) {
  const r1 = a + b;
  const r2 = c + d;
  const c1 = a + c;
  const n = r1 + r2;
  const prob = (x: number) => Math.exp(lnFact(r1) + lnFact(r2) + lnFact(c1) + lnFact(n - c1) - lnFact(n) - lnFact(x) - lnFact(r1 - x) - lnFact(c1 - x) - lnFact(r2 - c1 + x));
  const pObs = prob(a);
  let p = 0;
  for (let x = Math.max(0, c1 - r2); x <= Math.min(r1, c1); x++) {
    const px = prob(x);
    if (px <= pObs * (1 + 1e-7)) p += px;
  }
  return Math.min(1, p);
}

function groupStats(v: number[]) {
  const d = describe(v)!;
  return { n: d.n, median: round(d.median), q1: round(d.q1), q3: round(d.q3), mean: round(d.mean), sd: round(d.sd), min: d.min, max: d.max };
}

/** Group a (nonmetropolitan) against group b (metropolitan): medians, HL shift with CI, probability of superiority, Mann–Whitney. */
export function compare(a: number[], b: number[]) {
  const t = mannWhitney(a, b);
  const hl = hodgesLehmann(a, b);
  return {
    nNonmetro: a.length,
    nMetro: b.length,
    medianDifference: round(quantile(a, 0.5) - quantile(b, 0.5)),
    hl: round(hl.estimate),
    hlCI: hl.ci && ([round(hl.ci[0]), round(hl.ci[1])] as [number, number]),
    probSuperiority: round(t.prob),
    rankBiserial: round(2 * t.prob - 1),
    U: t.U,
    z: round(t.z),
    p: t.p,
  };
}

export function analyze(inp: RuralStudyInputs) {
  const { countyOrder: fips, counties } = inp;
  if (counties.map((c) => c.fips).join() !== fips.join()) throw new Error("county order differs between inputs");
  if (!counties.every((c) => c.metro === c.rucc <= 3)) throw new Error("metro flag disagrees with RUCC 1–3");
  const metro = counties.map((c) => c.metro);
  const nMetro = metro.filter(Boolean).length;
  const split = (s: Series) => ({
    non: s.flatMap((v, i) => (v === null || metro[i] ? [] : [v])),
    met: s.flatMap((v, i) => (v === null || !metro[i] ? [] : [v])),
  });

  const rows = INDICATORS.map((id) => {
    const cell = inp.primary[id];
    const { non, met } = split(cell.values);
    const suppressed = fips.flatMap((f, i) => (cell.values[i] === null ? [{ fips: f, name: counties[i].name, metro: metro[i], pop2020: counties[i].pop2020, reason: cell.reasons?.[f] ?? "No value published" }] : []));
    const sNon = suppressed.filter((s) => !s.metro).length;
    const sMet = suppressed.length - sNon;
    return {
      id,
      family: PRIMARY.includes(id) ? ("primary" as const) : ("secondary" as const),
      nonmetro: { ...groupStats(non), suppressed: sNon },
      metro: { ...groupStats(met), suppressed: sMet },
      comparison: { ...compare(non, met), pAdj: NaN },
      suppressed,
      missingness: suppressed.length
        ? {
            fisherP: fisherExact(sNon, fips.length - nMetro - sNon, sMet, nMetro - sMet),
            medianPop2020Suppressed: quantile(suppressed.map((s) => s.pop2020), 0.5),
            medianPop2020Published: quantile(counties.filter((_, i) => cell.values[i] !== null).map((c) => c.pop2020), 0.5),
            maxPop2020Suppressed: Math.max(...suppressed.map((s) => s.pop2020)),
          }
        : null,
    };
  });
  for (const fam of [PRIMARY, SECONDARY as IndId[]]) {
    const idx = rows.flatMap((r, i) => (fam.includes(r.id) ? [i] : []));
    const adj = benjaminiHochberg(idx.map((i) => rows[i].comparison.p));
    idx.forEach((i, k) => (rows[i].comparison.pAdj = adj[k]!));
  }

  const crude = SECONDARY.map((id) => {
    const { non, met } = split(inp.crude[id].values);
    return { id, nonmetroMedian: round(quantile(non, 0.5)), metroMedian: round(quantile(met, 0.5)), ...compare(non, met) };
  });
  const gradient = INDICATORS.map((id) => {
    const s = inp.primary[id].values;
    const idx = s.flatMap((v, i) => (v === null ? [] : [i]));
    const r = spearman(
      idx.map((i) => counties[i].rucc),
      idx.map((i) => s[i]!),
    );
    return { id, n: r.n, rho: round(r.r), p: r.p! };
  });
  const isAdjacent = (i: number) => [4, 6, 8].includes(counties[i].rucc);
  const adjacency = INDICATORS.map((id) => {
    const s = inp.primary[id].values;
    const adj = s.flatMap((v, i) => (v === null || metro[i] || !isAdjacent(i) ? [] : [v]));
    const non = s.flatMap((v, i) => (v === null || metro[i] || isAdjacent(i) ? [] : [v]));
    const c = compare(non, adj); // nonadjacent minus adjacent
    return { id, nAdjacent: adj.length, nNonadjacent: non.length, medianAdjacent: round(quantile(adj, 0.5)), medianNonadjacent: round(quantile(non, 0.5)), hl: c.hl, hlCI: c.hlCI, probSuperiority: c.probSuperiority, p: c.p };
  });

  const group = (m: boolean) => counties.filter((c) => c.metro === m);
  return {
    study: STUDY_ID,
    parameters: { alpha: 0.05, direction: "nonmetropolitan minus metropolitan", multipleTesting: "Benjamini–Hochberg FDR within the primary family (2 tests) and within the secondary family (6 tests)", ciLevel: 0.95 },
    counties: {
      total: fips.length,
      metro: group(true).length,
      nonmetro: group(false).length,
      byRucc: Object.fromEntries([1, 2, 3, 4, 5, 6, 7, 8, 9].map((k) => [k, counties.filter((c) => c.rucc === k).length])),
      pop2020Metro: group(true).reduce((s, c) => s + c.pop2020, 0),
      pop2020Nonmetro: group(false).reduce((s, c) => s + c.pop2020, 0),
      medianPop2020Metro: quantile(group(true).map((c) => c.pop2020), 0.5),
      medianPop2020Nonmetro: quantile(group(false).map((c) => c.pop2020), 0.5),
    },
    indicators: rows,
    exploratory: { crude, gradient, adjacency },
    /** Per-county values behind the figures. */
    values: {
      fips,
      names: counties.map((c) => c.name),
      metro,
      rucc: counties.map((c) => c.rucc),
      ...Object.fromEntries(INDICATORS.map((id) => [id, inp.primary[id].values])),
    } as { fips: string[]; names: string[]; metro: boolean[]; rucc: number[] } & Record<IndId, Series>,
  };
}

export type RuralStudyResults = ReturnType<typeof analyze>;