// Analysis for the Vellora Research report "Geographic Variation in
// Cardiovascular Health Across Texas Counties". Pure and deterministic: the
// same inputs always give the same results (fixed permutation seed).
//
// Inputs are the validated files in public/data (see
// scripts/fetch-cardio-study.mjs and scripts/fetch-heart-mortality.mjs).
// The results are written to public/data/research/cardio-disparities-results.v1.json
// by `npm run research:cardio`; tests/cardio-study.test.ts fails if the file
// is stale, and scripts/verify-cardio-study.py recomputes them with SciPy.

import { jStat } from "jstat";
import { benjaminiHochberg, describe, moransI, quantile, rank, spearman } from "../stats";

export const STUDY_ID = "cardio-disparities-2026";
export const PERMUTATIONS = 9999;
export const SEED = 20261022;

export type VarId = "BPHIGH" | "DIABETES" | "OBESITY" | "CSMOKING" | "LPA" | "CHD" | "STROKE" | "HDMORT";
export const RISK_FACTORS: VarId[] = ["BPHIGH", "DIABETES", "OBESITY", "CSMOKING", "LPA"];
export const OUTCOMES: VarId[] = ["CHD", "STROKE", "HDMORT"];
export const VARS: VarId[] = [...RISK_FACTORS, ...OUTCOMES];

export const LABELS: Record<VarId, { name: string; short: string; unit: string }> = {
  BPHIGH: { name: "High blood pressure among adults", short: "High blood pressure", unit: "%" },
  DIABETES: { name: "Diagnosed diabetes among adults", short: "Diabetes", unit: "%" },
  OBESITY: { name: "Obesity among adults", short: "Obesity", unit: "%" },
  CSMOKING: { name: "Current cigarette smoking among adults", short: "Current smoking", unit: "%" },
  LPA: { name: "No leisure-time physical activity among adults", short: "No leisure-time activity", unit: "%" },
  CHD: { name: "Coronary heart disease among adults", short: "Coronary heart disease", unit: "%" },
  STROKE: { name: "Stroke among adults", short: "Stroke", unit: "%" },
  HDMORT: { name: "Heart disease death rate, adults 35+", short: "Heart disease deaths", unit: "per 100,000" },
};

interface Layer {
  values: (number | null)[];
  lo: (number | null)[];
  hi: (number | null)[];
}
export interface StudyInputs {
  countyOrder: string[];
  counties: { fips: string; name: string; rucc: number; metro: boolean; pop2020: number }[];
  measures: Record<Exclude<VarId, "HDMORT">, { aa: Layer; crude: Layer }>;
}
export interface MortalityInput {
  countyOrder: string[];
  values: (number | null)[];
}

type Series = (number | null)[];
const round = (v: number, d = 4) => Number(v.toFixed(d));

/** Pairs where both values are reported (complete-case). */
function complete(x: Series, y: Series) {
  const xs: number[] = [];
  const ys: number[] = [];
  x.forEach((v, i) => {
    if (v === null || y[i] === null) return;
    xs.push(v);
    ys.push(y[i]!);
  });
  return [xs, ys] as const;
}

/**
 * Spearman ρ with a 95% CI from the Fisher z transform using the
 * Bonett–Wright standard error √((1 + ρ²/2)/(n − 3)).
 */
export function spearmanCI(x: number[], y: number[]) {
  const s = spearman(x, y);
  const z = Math.atanh(s.r);
  const se = Math.sqrt((1 + (s.r * s.r) / 2) / (s.n - 3));
  const zc = jStat.normal.inv(0.975, 0, 1);
  return { rho: s.r, n: s.n, p: s.p, ci: [Math.tanh(z - zc * se), Math.tanh(z + zc * se)] as [number, number] };
}

/**
 * Two-sided Mann–Whitney U test, normal approximation with tie and continuity
 * corrections. `a` is the reference group: U counts pairs where a > b (ties ½),
 * so prob = U / (na·nb) is the probability that a random county in `a` has the
 * higher value.
 */
export function mannWhitney(a: number[], b: number[]) {
  const na = a.length;
  const nb = b.length;
  const n = na + nb;
  const r = rank([...a, ...b]);
  const ra = r.slice(0, na).reduce((s, v) => s + v, 0);
  const U = ra - (na * (na + 1)) / 2;
  const counts = new Map<number, number>();
  for (const v of [...a, ...b]) counts.set(v, (counts.get(v) ?? 0) + 1);
  const ties = [...counts.values()].reduce((s, t) => s + t ** 3 - t, 0);
  const mu = (na * nb) / 2;
  const sigma = Math.sqrt(((na * nb) / 12) * (n + 1 - ties / (n * (n - 1))));
  const d = U - mu;
  const z = d === 0 ? 0 : (d - 0.5 * Math.sign(d)) / sigma;
  const p = Math.min(1, 2 * (1 - jStat.normal.cdf(Math.abs(z), 0, 1)));
  return { U, z, p, prob: U / (na * nb) };
}

export interface Descriptive {
  n: number;
  withheld: string[];
  mean: number;
  sd: number;
  median: number;
  q1: number;
  q3: number;
  p10: number;
  p90: number;
  ratio90to10: number;
  min: number;
  max: number;
  minCounties: string[];
  maxCounties: string[];
  medianCIWidth: number | null;
}

export function analyze(inputs: StudyInputs, mortality: MortalityInput, neighbors: Map<string, string[]>) {
  if (mortality.countyOrder.join() !== inputs.countyOrder.join()) throw new Error("county order differs between inputs");
  const fips = inputs.countyOrder;
  const nameOf = new Map(inputs.counties.map((c) => [c.fips, c.name]));
  const series = (id: VarId, type: "aa" | "crude" = "aa"): Series => (id === "HDMORT" ? mortality.values : inputs.measures[id][type].values);
  const metro = inputs.counties.map((c) => c.metro);

  const descriptives = {} as Record<VarId, Descriptive>;
  for (const id of VARS) {
    const s = series(id);
    const reported = s.flatMap((v) => (v === null ? [] : [v]));
    const d = describe(reported)!;
    const p10 = quantile(reported, 0.1);
    const p90 = quantile(reported, 0.9);
    const at = (v: number) => fips.filter((_, i) => s[i] === v).map((f) => nameOf.get(f)!);
    const layer = id === "HDMORT" ? null : inputs.measures[id].aa;
    const widths = layer ? layer.hi.flatMap((h, i) => (h === null ? [] : [h - layer.lo[i]!])) : null;
    descriptives[id] = {
      n: d.n,
      withheld: fips.filter((_, i) => s[i] === null).map((f) => nameOf.get(f)!),
      mean: round(d.mean),
      sd: round(d.sd),
      median: round(d.median),
      q1: round(d.q1),
      q3: round(d.q3),
      p10: round(p10),
      p90: round(p90),
      ratio90to10: round(p90 / p10),
      min: d.min,
      max: d.max,
      minCounties: at(d.min),
      maxCounties: at(d.max),
      medianCIWidth: widths ? round(quantile(widths, 0.5)) : null,
    };
  }

  const spatial = {} as Record<VarId, { I: number; expected: number; p: number; n: number; permutations: number }>;
  for (const id of VARS) {
    const s = series(id);
    const m = moransI(new Map(fips.flatMap((f, i) => (s[i] === null ? [] : [[f, s[i]!] as [string, number]]))), neighbors, PERMUTATIONS, SEED)!;
    spatial[id] = { I: round(m.I), expected: round(m.expected), p: m.p, n: m.n, permutations: m.permutations };
  }

  const mw = VARS.map((id) => {
    const s = series(id);
    const non = s.flatMap((v, i) => (v === null || metro[i] ? [] : [v]));
    const met = s.flatMap((v, i) => (v === null || !metro[i] ? [] : [v]));
    const t = mannWhitney(non, met);
    const [mn, mm] = [quantile(non, 0.5), quantile(met, 0.5)];
    return { id, nNonmetro: non.length, nMetro: met.length, medianNonmetro: round(mn), medianMetro: round(mm), difference: round(mn - mm), U: t.U, z: round(t.z), p: t.p, probNonmetroHigher: round(t.prob) };
  });
  const mwAdj = benjaminiHochberg(mw.map((m) => m.p));
  const metroComparison = mw.map((m, i) => ({ ...m, pAdj: mwAdj[i]! }));

  const correlate = (pairs: [VarId, VarId][], type: "aa" | "crude") => {
    const rows = pairs.map(([x, y]) => {
      const [xs, ys] = complete(series(x, type), series(y, type));
      const c = spearmanCI(xs, ys);
      return { x, y, n: c.n, rho: round(c.rho), ci: [round(c.ci[0]), round(c.ci[1])] as [number, number], p: c.p! };
    });
    const adj = benjaminiHochberg(rows.map((r) => r.p));
    return rows.map((r, i) => ({ ...r, pAdj: adj[i]! }));
  };
  const grid = RISK_FACTORS.flatMap((x) => OUTCOMES.map((y) => [x, y] as [VarId, VarId]));

  const pop = inputs.counties.reduce((s, c) => s + c.pop2020, 0);
  const metroPop = inputs.counties.filter((c) => c.metro).reduce((s, c) => s + c.pop2020, 0);

  return {
    study: STUDY_ID,
    parameters: {
      permutations: PERMUTATIONS,
      seed: SEED,
      alpha: 0.05,
      multipleTesting: "Benjamini–Hochberg within each table",
      prevalenceType: "Age-adjusted prevalence (AgeAdjPrv); crude prevalence (CrdPrv) as a sensitivity analysis",
    },
    counties: {
      total: fips.length,
      metro: metro.filter(Boolean).length,
      nonmetro: metro.filter((m) => !m).length,
      metroPopulationShare: round(metroPop / pop),
      population2020: pop,
    },
    descriptives,
    spatial,
    metroComparison,
    correlations: correlate(grid, "aa"),
    correlationsCrude: correlate(grid, "crude"),
    outcomePairs: correlate([["CHD", "STROKE"], ["CHD", "HDMORT"], ["STROKE", "HDMORT"]], "aa"),
    /** Per-county values behind the charts: age-adjusted prevalence and the death rate. */
    values: {
      fips,
      names: inputs.counties.map((c) => c.name),
      metro,
      ...Object.fromEntries(VARS.map((id) => [id, series(id)])),
    } as { fips: string[]; names: string[]; metro: boolean[] } & Record<VarId, Series>,
  };
}

export type StudyResults = ReturnType<typeof analyze>;