/**
 * The TypeScript implementation's side of the SEDA contract.
 *
 * As with the other two emitters, nothing is computed here that the site does
 * not already compute: `src/seda.ts` does the join and the percentiles, and this
 * writes the result out in the shared format.
 *
 * Run: npm run calc:ts   (writes calc/out/seda-typescript.tsv)
 */

import { createHash } from 'node:crypto'
import { existsSync, mkdirSync, readFileSync, writeFileSync } from 'node:fs'
import { join } from 'node:path'

import { int, NA, real, serialize, text, type Row } from '../src/contract.ts'
import { parseCsv } from '../src/csv.ts'
import { joinSeda, num, nycSedaId, NYC_SEDA_LEA, type SedaRow } from '../src/seda.ts'
import { K8_SCHOOLS, NYC_SCHOOLS } from '../src/schools.ts'
import { median, pearson } from '../src/stats.ts'
import { currentEnrollment, type CcdEnrollmentRow } from '../src/demographics.ts'

const ROOT = join(import.meta.dirname, '..')
const SEDA = join(ROOT, 'seda')
const DEMOGRAPHICS = join(ROOT, 'demographics')

/** LEAs the current federal enrolment counts are read for. */
const CCD_LEAS = ['4218990', '0103360', '3600091', '3600092']

const sha256 = (s: string): string => createHash('sha256').update(s, 'utf8').digest('hex')

/** Where the low-income measure runs out of room; see CONTRACT.md. */
const POVERTY_CEILING = 0.95

const PHILA_LEA = '4218990'
const TUSC_CITY_LEA = '103360'
const TUSC_COUNTY_LEA = '103390'

const DISTRICT_NAMES: Record<string, string> = {
  [PHILA_LEA]: 'Philadelphia City SD',
  [TUSC_CITY_LEA]: 'Tuscaloosa City',
  [TUSC_COUNTY_LEA]: 'Tuscaloosa County',
}

/** Each source, and the columns every implementation reads out of it. */
const SOURCES: { name: string; file: string; columns: string[] }[] = [
  {
    name: 'achievement',
    file: 'seda_school_pool_cs_4.1.csv',
    columns: ['sedasch', 'sedaschname', 'cs_mn_avg_eb', 'cs_mn_grd_eb'],
  },
  {
    name: 'covariates',
    file: 'seda_cov_school_pool_4.1.csv',
    columns: ['sedasch', 'sedalea', 'schcity', 'schnam', 'level', 'perblk', 'totenrl', 'perecd'],
  },
  {
    name: 'districts',
    file: 'seda2022_dist.csv',
    columns: ['sedaadmin', 'sedaadminname', 'subject', 'subgroup', 'ys_mn19_eb', 'ys_mn22_eb', 'ys_chg_eb'],
  },
]

const read = (file: string): SedaRow[] => {
  const path = join(SEDA, file)
  if (!existsSync(path)) throw new Error(`missing ${path} — run ./fetch-seda.sh`)
  return parseCsv(readFileSync(path, 'latin1'))
}

export function emitRows(): Row[] {
  const rows: Row[] = []
  const add = (metric: string, key: string, value: string): void => {
    rows.push({ metric, key, value })
  }

  const parsed = new Map<string, SedaRow[]>()
  for (const src of SOURCES) {
    const p = read(src.file)
    parsed.set(src.name, p)
    add('input.rows', src.name, int(p.length))
    add('input.digest', src.name, text(sha256(p.map((r) => src.columns.map((c) => r[c] ?? '').join('\x1f')).join('\n'))))
  }

  // Current demographics, so no decade-old racial composition is emitted alone.
  const ccdRows: CcdEnrollmentRow[] = []
  for (const lea of CCD_LEAS) {
    const path = join(DEMOGRAPHICS, `enrollment_${lea}.json`)
    if (!existsSync(path)) throw new Error(`missing ${path} — run ./fetch-demographics.sh`)
    const body = readFileSync(path, 'utf8')
    add('input.filedigest', `enrollment_${lea}`, text(sha256(body)))
    ccdRows.push(...((JSON.parse(body) as { results?: CcdEnrollmentRow[] }).results ?? []))
  }
  const nowByNces = currentEnrollment(ccdRows)
  add('input.rows', 'ccd_enrollment', int(ccdRows.length))
  add('input.schools', 'ccd_enrollment', int(nowByNces.size))

  const achievement = parsed.get('achievement') as SedaRow[]
  const covariates = parsed.get('covariates') as SedaRow[]
  const distRows = parsed.get('districts') as SedaRow[]

  // ------------------------------------------------------------ the schools
  const labels: Record<string, string> = {}
  for (const s of K8_SCHOOLS) if (s.nces) labels[s.nces] = s.label
  const wanted = new Set(K8_SCHOOLS.map((s) => s.nces).filter((v): v is string => v !== undefined))

  // SEDA files New York under one synthetic citywide LEA rather than the
  // geographic districts the federal collections use, so the ids crosswalk.
  const nycLabels: Record<string, string> = {}
  const nycWanted = new Set<string>()
  for (const s of NYC_SCHOOLS) {
    const id = nycSedaId(s.nces)
    nycLabels[id] = s.label
    nycWanted.add(id)
  }

  const nycFederalId = new Map(NYC_SCHOOLS.map((s) => [nycSedaId(s.nces), s.nces]))

  /**
   * The current federal counts for one SEDA school, or null.
   *
   * SEDA drops Alabama's leading zero and files New York under a synthetic
   * citywide LEA; both have to be undone to reach the federal counts.
   */
  const nowOf = (id: string): { enrollment: number | null; black: number | null } | null =>
    nowByNces.get(nycFederalId.get(id) ?? (id.length === 11 ? `0${id}` : id)) ?? null

  /**
   * One line per school, carrying exactly the values the per-school facts do.
   *
   * Brooklyn is in the contract twice: the ten schools somebody named, and the
   * whole borough they sit in. The borough is four hundred schools, and writing
   * it out would add nearly five thousand lines and bury everything else, so it
   * is checked by hashing this dump instead. One differing field in one school
   * changes the hash.
   */
  const schoolDump = (schools: ReturnType<typeof joinSeda>): string =>
    schools
      .map((s, i) => {
        const now = nowOf(s.id)
        return [
          s.id,
          s.name,
          s.district,
          s.gradeLevel,
          real(s.level),
          real(s.growth),
          s.levelPct === null ? NA : int(s.levelPct),
          s.growthPct === null ? NA : int(s.growthPct),
          real(s.pctBlack),
          real(s.enrollment),
          now === null || now.enrollment === null ? NA : int(now.enrollment),
          now === null || now.black === null ? NA : int(now.black),
          int(i),
        ].join('\x1f')
      })
      .join('\n')

  const isBrooklyn = (c: SedaRow): boolean =>
    c['sedalea'] === NYC_SEDA_LEA && (c['schcity'] ?? '').toUpperCase() === 'BROOKLYN'

  const groups: { group: string; schools: ReturnType<typeof joinSeda>; perSchool: boolean }[] = [
    { group: 'phila', schools: joinSeda(achievement, covariates, (c) => wanted.has(c['sedasch'] ?? ''), labels), perSchool: true },
    { group: 'tuscaloosa', schools: joinSeda(achievement, covariates, (c) => c['sedalea'] === TUSC_CITY_LEA), perSchool: true },
    { group: 'brooklyn', schools: joinSeda(achievement, covariates, (c) => nycWanted.has(c['sedasch'] ?? ''), nycLabels), perSchool: true },
    { group: 'allBrooklyn', schools: joinSeda(achievement, covariates, isBrooklyn), perSchool: false },
  ]

  for (const { group, schools, perSchool } of groups) {
    add('seda.n', group, int(schools.length))
    // The page draws a median rule on every row and plots growth for whichever
    // schools have it, so both are contract facts rather than page arithmetic.
    add('seda.median', group, real(median(schools.map((s) => s.level))))
    add('seda.withGrowth', group, int(schools.filter((s) => s.growth !== null).length))
    add('seda.digest', group, text(sha256(schoolDump(schools))))
    if (!perSchool) continue
    schools.forEach((s, i) => {
      const key = `${group}|${s.id}`
      add('seda.level', key, real(s.level))
      add('seda.growth', key, real(s.growth))
      add('seda.levelPct', key, s.levelPct === null ? 'NA' : int(s.levelPct))
      add('seda.growthPct', key, s.growthPct === null ? 'NA' : int(s.growthPct))
      add('seda.pctBlack', key, real(s.pctBlack))
      add('seda.enrollment', key, real(s.enrollment))
      add('seda.name', key, text(s.name))
      add('seda.district', key, text(s.district))
      add('seda.gradeLevel', key, text(s.gradeLevel))
      // The page lists schools by level, descending. Emitting the position
      // checks the ordering too, not just the values.
      add('seda.rank', key, int(i))
      const now = nowOf(s.id)
      add('seda.nowEnrollment', key, now === null || now.enrollment === null ? NA : int(now.enrollment))
      add('seda.nowBlack', key, now === null || now.black === null ? NA : int(now.black))
    })
  }

  // The two national distributions percentiles are taken against.
  const levels = achievement.map((a) => num(a['cs_mn_avg_eb'])).filter((v): v is number => v !== null)
  const growths = achievement.map((a) => num(a['cs_mn_grd_eb'])).filter((v): v is number => v !== null)
  add('national.levels', '-', int(levels.length))
  add('national.growths', '-', int(growths.length))

  // How much of each district SEDA can say anything about. New York is the
  // reason this is in the contract: it has a growth estimate for about a third
  // of its schools, against most of Philadelphia's and all of Tuscaloosa's.
  const covIds = (lea: string): Set<string> =>
    new Set(covariates.filter((c) => c['sedalea'] === lea).map((c) => c['sedasch'] as string))
  for (const [place, lea] of [
    ['Philadelphia', PHILA_LEA],
    ['Tuscaloosa_City', TUSC_CITY_LEA],
    ['New_York_City', NYC_SEDA_LEA],
  ] as [string, string][]) {
    const ids = covIds(lea)
    const rows2 = achievement.filter((a) => ids.has(a['sedasch'] ?? ''))
    add('coverage.schools', place, int(rows2.length))
    add('coverage.level', place, int(rows2.filter((a) => num(a['cs_mn_avg_eb']) !== null).length))
    add('coverage.growth', place, int(rows2.filter((a) => num(a['cs_mn_grd_eb']) !== null).length))
  }

  // New York is absent from the 2022 district release entirely.
  add('district.newYorkRows', '-', int(distRows.filter((r) => r['stateabb'] === 'NY').length))

  // ---------------------------------------------------------- the districts
  for (const r of distRows) {
    const id = r['sedaadmin'] ?? ''
    if (r['subgroup'] !== 'all' || !(id in DISTRICT_NAMES)) continue
    const subject = r['subject'] === 'mth' ? 'math' : 'reading'
    add('district.level2019', `${id}|${subject}`, real(num(r['ys_mn19_eb'])))
    add('district.level2022', `${id}|${subject}`, real(num(r['ys_mn22_eb'])))
    add('district.change', `${id}|${subject}`, real(num(r['ys_chg_eb'])))
  }
  for (const [id, name] of Object.entries(DISTRICT_NAMES)) add('district.name', id, text(name))

  for (const subj of ['mth', 'rla']) {
    const subject = subj === 'mth' ? 'math' : 'reading'
    const forSubject = distRows.filter((r) => r['subgroup'] === 'all' && r['subject'] === subj)
    add('usMedian.level2019', subject, real(median(forSubject.map((r) => num(r['ys_mn19_eb'])).filter((v): v is number => v !== null))))
    add('usMedian.change', subject, real(median(forSubject.map((r) => num(r['ys_chg_eb'])).filter((v): v is number => v !== null))))
    // Every district row for the subject, including the ones with no estimate:
    // this is the denominator the page prints.
    add('usMedian.n', subject, int(forSubject.length))
  }

  // ------------------------------------------------- how much is just intake
  const povertyOf = new Map<string, number>()
  for (const c of covariates) {
    const id = c['sedasch']
    const p = num(c['perecd'])
    if (id !== undefined && p !== null && p >= 0 && p <= 1) povertyOf.set(id, p)
  }
  add('poverty.schools', '-', int(povertyOf.size))

  // Where `perecd` stops measuring anything. It is built from free-and-reduced
  // -lunch eligibility, so a district on universal free meals books nearly
  // every school at the maximum whatever its families earn. The share of each
  // set at the ceiling is the only way to see which rows have run out of room,
  // so it is a contract fact rather than something the page works out.
  add('poverty.ceiling', '-', int([...povertyOf.values()].filter((p) => p >= POVERTY_CEILING).length))
  add('poverty.ceilingAt', '-', real(POVERTY_CEILING))

  for (const [name, column] of [
    ['level', 'cs_mn_avg_eb'],
    ['growth', 'cs_mn_grd_eb'],
  ] as [string, string][]) {
    const xs: number[] = []
    const ys: number[] = []
    for (const a of achievement) {
      const id = a['sedasch']
      const v = num(a[column])
      const pov = id === undefined ? undefined : povertyOf.get(id)
      if (v === null || pov === undefined) continue
      xs.push(pov)
      ys.push(v)
    }
    add(`poverty.${name}.r`, '-', real(pearson(xs, ys)))
    add(`poverty.${name}.n`, '-', int(xs.length))
  }

  // ------------------------------------------- the sets the page compares
  //
  // The intake table reads five sets off one national line. Its numbers were
  // the last part of the page not covered here, which is exactly the part a
  // reader is most likely to want checked.
  const levelOf = new Map<string, number>()
  for (const a of achievement) {
    const id = a['sedasch']
    const v = num(a['cs_mn_avg_eb'])
    if (id !== undefined && v !== null) levelOf.set(id, v)
  }

  const setStats = (set: string, ids: string[]): void => {
    const levels: number[] = []
    const povs: number[] = []
    for (const id of ids) {
      const l = levelOf.get(id)
      if (l !== undefined) levels.push(l)
      const p = povertyOf.get(id)
      if (p !== undefined) povs.push(p)
    }
    add('set.n', set, int(levels.length))
    add('set.level', set, real(median(levels)))
    add('set.poverty', set, real(median(povs)))
    // The denominator is the schools with a share at all — the same set the
    // median is taken over, not `set.n`, which counts level estimates.
    add('set.ceiling', set, int(povs.filter((p) => p >= POVERTY_CEILING).length))
    add('set.ceilingOf', set, int(povs.length))
  }

  const covIdsWhere = (keep: (c: SedaRow) => boolean): string[] =>
    covariates.filter(keep).map((c) => c['sedasch'] as string).filter((id) => id !== undefined)

  setStats('philaEleven', K8_SCHOOLS.map((s) => s.nces).filter((v): v is string => v !== undefined))
  setStats('philaAll', covIdsWhere((c) => c['sedalea'] === PHILA_LEA))
  setStats('brooklynTen', NYC_SCHOOLS.map((s) => nycSedaId(s.nces)))
  setStats('brooklynAll', covIdsWhere(isBrooklyn))
  setStats('tuscaloosaCity', covIdsWhere((c) => c['sedalea'] === TUSC_CITY_LEA))

  // Least squares of level on the low-income share, all US schools — the line
  // the table's "predicted" column is read off.
  {
    const xs: number[] = []
    const ys: number[] = []
    for (const a of achievement) {
      const id = a['sedasch']
      const v = num(a['cs_mn_avg_eb'])
      const p = id === undefined ? undefined : povertyOf.get(id)
      if (v === null || p === undefined) continue
      xs.push(p)
      ys.push(v)
    }
    const n = xs.length
    let sx = 0
    let sy = 0
    for (let i = 0; i < n; i++) {
      sx += xs[i] as number
      sy += ys[i] as number
    }
    const mx = sx / n
    const my = sy / n
    let sxy = 0
    let sxx = 0
    for (let i = 0; i < n; i++) {
      const dx = (xs[i] as number) - mx
      sxy += dx * ((ys[i] as number) - my)
      sxx += dx * dx
    }
    const slope = sxy / sxx
    add('povertyFit.slope', '-', real(slope))
    add('povertyFit.intercept', '-', real(my - slope * mx))
    add('povertyFit.n', '-', int(n))
  }

  return rows
}

if (import.meta.filename === process.argv[1]) {
  const outDir = join(ROOT, 'calc', 'out')
  mkdirSync(outDir, { recursive: true })
  const tsv = serialize(emitRows())
  writeFileSync(join(outDir, 'seda-typescript.tsv'), tsv)
  console.log(`wrote calc/out/seda-typescript.tsv (${tsv.split('\n').length - 2} facts)`)
}