/**
 * The TypeScript implementation's side of the calculation contract.
 *
 * This is the trunk the strangler fig grew around: it computes nothing itself.
 * Every number here comes from the same `src/` modules the published site has
 * always used, and this file only writes them out in the contract's format so
 * that the R Markdown and Python implementations have something to be compared
 * against.
 *
 * Run: npm run calc:ts   (writes calc/out/typescript.tsv)
 */

import { mkdirSync, writeFileSync } from 'node:fs'
import { join } from 'node:path'

import { int, real, serialize, text, type Row } from '../src/contract.ts'
import { createCorpus, buildSchoolReport, type Corpus, type SchoolReport } from '../src/report.ts'
import { employeeKey, isSchoolLeader, isTeacher } from '../src/retention.ts'
import { K8_SCHOOLS, SECONDARY_SCHOOLS, type SchoolRef } from '../src/schools.ts'
import { keyDigest, loadSnapshots, snapshotDigest, YEARS } from '../src/snapshots.ts'
import { median, pearson } from '../src/stats.ts'

const ROOT = join(import.meta.dirname, '..')

/** Peer level as it appears in a contract key: no spaces, so keys stay sortable. */
const LEVELS: { level: string; token: string }[] = [
  { level: 'ELEMENTARY', token: 'ELEMENTARY' },
  { level: 'MIDDLE SCHOOL', token: 'MIDDLE_SCHOOL' },
  { level: 'HIGH SCHOOL', token: 'HIGH_SCHOOL' },
]
const token = (level: string): string => LEVELS.find((l) => l.level === level)?.token ?? level

const SCOPES = ['teacher', 'teacherWithAlt', 'all'] as const

/** Minimum teaching staff for a school to count as a peer. Matches buildSchoolReport. */
const MIN_TEACHERS = 15

/** Two digits, so that run 10 sorts after run 9 rather than after run 1. */
const idx = (i: number): string => String(i).padStart(2, '0')

const transitions = YEARS.slice(0, -1).map((from, i) => ({ from, to: YEARS[i + 1] as number }))
const span = (t: { from: number; to: number }): string => `${t.from}-${t.to}`

export function emitRows(): Row[] {
  const byYear = loadSnapshots(join(ROOT, 'data', 'extracted'))
  const corpus: Corpus = createCorpus(byYear)
  const rows: Row[] = []
  const add = (metric: string, key: string, value: string): void => {
    rows.push({ metric, key, value })
  }

  // ------------------------------------------------------------------ inputs
  // Emitted first so that a disagreement about the numbers can be told apart
  // from a disagreement about what was read off disk.
  for (const year of YEARS) {
    const staff = byYear.get(year) ?? []
    add('input.rows', String(year), int(staff.length))
    add('input.digest', String(year), text(snapshotDigest(staff)))
    add('input.keydigest', String(year), text(keyDigest(staff.map(employeeKey))))
    add('input.keys', String(year), int(new Set(staff.map(employeeKey)).size))
    add('input.teachers', String(year), int(staff.filter(isTeacher).length))
    add('input.principals', String(year), int(staff.filter(isSchoolLeader).length))
  }

  // -------------------------------------------------------------- peer sets
  for (const { level, token: tok } of LEVELS) {
    for (const year of YEARS) {
      add('peer.orgs', `${year}|${tok}`, int(corpus.peerOrgs(year, level, MIN_TEACHERS).length))
    }
    for (const t of transitions) {
      for (const scope of SCOPES) {
        const rates = [...corpus.peerRates(t.from, t.to, level, MIN_TEACHERS, scope)].sort((a, b) => a - b)
        const at = (q: number): number | null => rates[Math.floor(q * (rates.length - 1))] ?? null
        const key = `${span(t)}|${tok}|${scope}`
        add('peer.n', key, int(rates.length))
        add('peer.median', key, real(median(rates)))
        add('peer.q1', key, real(at(0.25)))
        add('peer.q3', key, real(at(0.75)))
      }
    }
  }

  // ----------------------------------------------------------- per school
  const schools: SchoolRef[] = [...K8_SCHOOLS, ...SECONDARY_SCHOOLS]
  const reports = new Map<string, SchoolReport>()

  for (const school of schools) {
    const level = school.level ?? 'ELEMENTARY'
    const r = buildSchoolReport(corpus, school, YEARS, level, MIN_TEACHERS)
    reports.set(school.org, r)
    const org = school.org

    add('school.level', org, text(token(level)))

    const seriesFor = { teacher: r.teacher, teacherWithAlt: r.teacherWithAlt, all: r.allStaff }
    for (const scope of SCOPES) {
      for (const y of seriesFor[scope]) {
        const key = `${org}|${y.from}-${y.to}`
        add(`series.${scope}.stayed`, key, int(y.stayed))
        add(`series.${scope}.start`, key, int(y.startCount))
        add(`series.${scope}.percentile`, key, real(y.percentile))
      }
    }

    // Pooled rates are published as their two integers, so the comparison of the
    // headline figure is a comparison of counts rather than of floating point.
    const pooled: [string, { stayed: number; startCount: number }[]][] = [
      ['teacher', r.teacher],
      ['teacherWithAlt', r.teacherWithAlt],
      ['all', r.allStaff],
      ['teacherSince2020', r.teacher.filter((t) => t.from >= 2020)],
    ]
    for (const [name, series] of pooled) {
      add(`pooled.${name}.stayed`, org, int(series.reduce((n, y) => n + y.stayed, 0)))
      add(`pooled.${name}.start`, org, int(series.reduce((n, y) => n + y.startCount, 0)))
    }

    for (const a of r.altShare) {
      add('alt.pft', `${org}|${a.year}`, int(a.pft))
      add('alt.alt', `${org}|${a.year}`, int(a.alt))
    }

    add('principal.runs', org, int(r.principalRuns.length))
    add('principal.changes', org, int(r.principalChanges))
    r.principalRuns.forEach((run, i) => {
      add('principal.run', `${org}|${idx(i)}`, text(`${run.from}-${run.to}|${run.names}`))
    })

    // The base year's own row counts the cohort against itself, so its
    // `remaining` is the size of the cohort.
    add('cohort2020.base', org, int(r.cohort2020[0]?.remaining ?? 0))
    for (const c of r.cohort2020) add('cohort2020.remaining', `${org}|${c.year}`, int(c.remaining))

    add('teacherCount', org, int(r.teacherCount))

    add('destination.stayed', org, int(r.destinations.stayed))
    add('destination.movedSchool', org, int(r.destinations.movedSchool))
    add('destination.changedRole', org, int(r.destinations.changedRole))
    add('destination.leftDistrict', org, int(r.destinations.leftDistrict))

    add('era.n', org, int(r.eras.length))
    r.eras.forEach((e, i) => {
      add('era.span', `${org}|${idx(i)}`, text(`${e.from}-${e.to}|${e.names}`))
      add('era.stayed', `${org}|${idx(i)}`, int(e.stayed))
      add('era.start', `${org}|${idx(i)}`, int(e.startCount))
    })
  }

  // ------------------------------------------- does principal churn predict retention?
  // The two input series are emitted alongside the coefficients: if a
  // correlation disagrees, this says whether the inputs or the arithmetic did.
  const rankable = K8_SCHOOLS.map((s) => reports.get(s.org) as SchoolReport).filter((r) => r.pooledTeacher !== null)
  const longestTenure = (r: SchoolReport): number => Math.max(...r.principalRuns.map((run) => run.to - run.from + 1))

  for (const r of rankable) {
    add('correlation.churn', r.org, int(r.principalChanges))
    add('correlation.tenure', r.org, int(longestTenure(r)))
  }
  add('correlation.n', '-', int(rankable.length))
  add('correlation.churnVsRetention', '-', real(pearson(rankable.map((r) => r.principalChanges), rankable.map((r) => r.pooledTeacher as number))))
  add('correlation.tenureVsRetention', '-', real(pearson(rankable.map(longestTenure), rankable.map((r) => r.pooledTeacher as number))))

  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, 'retention-typescript.tsv'), tsv)
  console.log(`wrote calc/out/retention-typescript.tsv (${tsv.split('\n').length - 2} facts)`)
}