#!/usr/bin/env python3
"""
The Python implementation of the SDP retention calculation.

This is a second, independent implementation of the same calculation the
published site runs, written against the same eleven October extracts. It shares
no code with either the TypeScript original or the R Markdown implementation —
only the contract in ``calc/CONTRACT.md``, which says what facts to emit and how
to write them down. If all three files come out byte-identical, the calculation
does not depend on who implemented it.

Standard library only, on purpose: pandas would make some of this shorter, but
the point of the exercise is that three implementations agree, and a
column-oriented library agreeing with two row-oriented ones proves less than it
appears to. What is worth checking here is the definitions, not the dataframes.

Run: npm run calc:py   (writes calc/out/retention-python.tsv)
"""

from __future__ import annotations

import csv
import math
import os
import re
import sys
from collections import OrderedDict

from contract import fmt_int, fmt_real, fmt_text, serialize, sha256

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
DATA = os.path.join(ROOT, "data", "extracted")

# --------------------------------------------------------------------- inputs

SNAPSHOTS = [
    (2015, "employee_information_Oct2015.csv"),
    (2016, "employee_information_Oct2016.csv"),
    (2017, "employee_information_Oct2017.csv"),
    (2018, "employee_information_Oct2018.csv"),
    (2019, "employee_information_Oct2019.csv"),
    (2020, "employee_information_Oct2020.csv"),
    (2021, "employee_information_Oct2021.csv"),
    (2022, "employee_information_Oct2022.csv"),
    (2023, "Employee_Data_2023-2024/employee_information_October2023.csv"),
    (2024, "employee_information_Oct2024.csv"),
    (2025, "Oct 2025/employee_information_Oct2025.csv"),
]

YEARS = [y for y, _ in SNAPSHOTS]
TRANSITIONS = list(zip(YEARS, YEARS[1:]))

LEVELS = [("ELEMENTARY", "ELEMENTARY"), ("MIDDLE SCHOOL", "MIDDLE_SCHOOL"), ("HIGH SCHOOL", "HIGH_SCHOOL")]
SCOPES = ["teacher", "teacherWithAlt", "all"]
MIN_TEACHERS = 15

# org code -> peer level. Anything not listed is ranked against elementary peers.
K8_ORGS = ["1300", "1280", "1340", "1260", "1390", "1470", "2380", "2590", "2470", "2480", "2210"]
SECONDARY_ORGS = [("2050", "MIDDLE SCHOOL"), ("2650", "HIGH SCHOOL"), ("2680", "HIGH SCHOOL"), ("6040", "HIGH SCHOOL")]
SCHOOLS = [(org, "ELEMENTARY") for org in K8_ORGS] + SECONDARY_ORGS

# The six fields every implementation reads, in the order they are hashed.
FIELDS = ("LAST_NAME", "FIRST_NAME", "TITLE_DESCRIPTION", "HOME_ORGANIZATION",
          "ORGANIZATION_LEVEL", "TYPE_OF_REPRESENTATION")


class Employee:
    """One row of an extract, reduced to the fields the calculation uses."""

    __slots__ = ("last", "first", "title", "org", "level", "rep")

    def __init__(self, row):
        self.last = row.get("LAST_NAME") or ""
        self.first = row.get("FIRST_NAME") or ""
        self.title = row.get("TITLE_DESCRIPTION") or ""
        self.org = row.get("HOME_ORGANIZATION") or ""
        self.level = row.get("ORGANIZATION_LEVEL") or ""
        self.rep = row.get("TYPE_OF_REPRESENTATION") or ""

    def fields(self):
        return (self.last, self.first, self.title, self.org, self.level, self.rep)


def load(path):
    """Read one October extract. The extracts are latin-1, not UTF-8."""
    with open(path, encoding="latin-1", newline="") as fh:
        return [Employee(r) for r in csv.DictReader(fh)]


# ----------------------------------------------------------------- definitions

_NOT_LETTER = re.compile(r"[^A-Z ]")
_SPACES = re.compile(r"\s+")


def employee_key(e):
    """
    Identity key for matching a person across two snapshots.

    The extracts carry no employee ID, so name is all there is. Case,
    punctuation and internal spacing are normalized away so that a cosmetic
    spelling change does not read as somebody leaving.
    """
    def norm(s):
        return _SPACES.sub(" ", _NOT_LETTER.sub("", s.upper())).strip()
    return norm(e.last) + "|" + norm(e.first)


# `\b` is deliberately ASCII-only, to match the JavaScript regular expression
# this definition was first written as; Python's default is Unicode-aware.
_TEACHER_TITLE = re.compile(r"^TEACHER\b", re.ASCII)


def is_teacher(e):
    """Classroom teachers: a PFT teaching local *and* a TEACHER title."""
    return e.rep.upper().startswith("PFT") and _TEACHER_TITLE.match(e.title.upper()) is not None


def is_alt_instructional(e):
    """Student teachers and bought-in contracted placements: not district hires."""
    t = e.title.upper()
    return "STUDENT TEACHER" in t or "CONTRACTED TEACHER" in t


def is_school_leader(e):
    """The building principal — not assistant principals or retired fill-ins."""
    return e.title.strip().upper() == "PRINCIPAL"


SCOPE_FILTERS = {
    "teacher": is_teacher,
    "teacherWithAlt": lambda e: is_teacher(e) or is_alt_instructional(e),
    "all": lambda e: True,
}


# --------------------------------------------------------------------- corpus

class Corpus:
    """Indexed view over the yearly extracts, mirroring src/report.ts."""

    def __init__(self, by_year):
        self.by_year = by_year
        # year -> org -> [employee], with orgs in the order they first appear.
        self.index = {}
        for year, staff in by_year.items():
            by_org = OrderedDict()
            for e in staff:
                by_org.setdefault(e.org, []).append(e)
            self.index[year] = by_org

        # Person-level indexes, for telling a transfer from an exit. A duplicated
        # name is resolved the same way a dictionary assignment resolves it: the
        # last row in the file wins.
        self.teacher_org = {}
        self.payroll = {}
        for year, staff in by_year.items():
            t, p = {}, set()
            for e in staff:
                k = employee_key(e)
                p.add(k)
                if is_teacher(e):
                    t[k] = e.org
            self.teacher_org[year] = t
            self.payroll[year] = p

        self._peer_orgs = {}
        self._peer_rates = {}

    def staff_at(self, year, org, scope):
        keep = SCOPE_FILTERS[scope]
        return [e for e in self.index.get(year, {}).get(org, []) if keep(e)]

    def leaders_at(self, year, org):
        return [e for e in self.index.get(year, {}).get(org, []) if is_school_leader(e)]

    def peer_orgs(self, year, level, min_teachers):
        cached = self._peer_orgs.get((year, level, min_teachers))
        if cached is not None:
            return cached
        out = []
        for org, staff in self.index.get(year, {}).items():
            # The level is read off the org's first row, as the original does:
            # a handful of orgs carry rows labelled with more than one level.
            if not staff or level not in staff[0].level.upper():
                continue
            if sum(1 for e in staff if is_teacher(e)) >= min_teachers:
                out.append(org)
        self._peer_orgs[(year, level, min_teachers)] = out
        return out

    def peer_rates(self, frm, to, level, min_teachers, scope):
        cached = self._peer_rates.get((frm, to, level, min_teachers, scope))
        if cached is not None:
            return cached
        rates = []
        for org in self.peer_orgs(frm, level, min_teachers):
            r = retention_between(self.staff_at(frm, org, scope), self.staff_at(to, org, scope))
            if r[2] is not None:
                rates.append(r[2])
        self._peer_rates[(frm, to, level, min_teachers, scope)] = rates
        return rates


def retention_between(before, after):
    """(stayed, start_count, rate) for one school across consecutive Octobers."""
    b = {employee_key(e) for e in before}
    a = {employee_key(e) for e in after}
    stayed = len(b & a)
    start = len(b)
    return stayed, start, (None if start == 0 else stayed / start)


# ----------------------------------------------------------------------- stats

def median(values):
    if not values:
        return None
    s = sorted(values)
    mid = len(s) // 2
    return s[mid] if len(s) % 2 == 1 else (s[mid - 1] + s[mid]) / 2


def percentile_rank(values, target):
    """Share of `values` at or below `target`."""
    if not values:
        return None
    return sum(1 for v in values if v <= target) / len(values)


def quantile_at(sorted_rates, q):
    """The original's quartile: index into the sorted rates, no interpolation."""
    if not sorted_rates:
        return None
    return sorted_rates[math.floor(q * (len(sorted_rates) - 1))]


def pearson(xs, ys):
    """
    Pearson's r, accumulated in an explicit left-to-right loop.

    Written as a loop rather than with sum() so that all three implementations
    add their terms up in the same order, and land on the same bits.
    """
    n = min(len(xs), len(ys))
    if n < 2:
        return None
    sx = 0.0
    sy = 0.0
    for i in range(n):
        sx += xs[i]
        sy += ys[i]
    mx, my = sx / n, sy / n
    num = dx = dy = 0.0
    for i in range(n):
        a, b = xs[i] - mx, ys[i] - my
        num += a * b
        dx += a * a
        dy += b * b
    if dx == 0 or dy == 0:
        return None
    return num / math.sqrt(dx * dy)


# ------------------------------------------------------------------ the report

def compress_timeline(entries):
    """Collapse a year-by-year holder list into runs, so repeats read as one tenure."""
    runs = []
    for year, names in entries:
        if runs and runs[-1][0] == names:
            runs[-1][2] = year
        else:
            runs.append([names, year, year])
    return runs


def build():
    by_year = OrderedDict()
    for year, filename in SNAPSHOTS:
        path = os.path.join(DATA, filename)
        if not os.path.exists(path):
            raise SystemExit("missing snapshot: %s — run ./fetch-data.sh" % path)
        by_year[year] = load(path)

    corpus = Corpus(by_year)
    rows = []

    def add(metric, key, value):
        rows.append((metric, key, value))

    # ------------------------------------------------------------------ inputs
    for year in YEARS:
        staff = by_year[year]
        keys = [employee_key(e) for e in staff]
        add("input.rows", str(year), fmt_int(len(staff)))
        add("input.digest", str(year), fmt_text(sha256("\n".join("\x1f".join(e.fields()) for e in staff))))
        add("input.keydigest", str(year), fmt_text(sha256("\n".join(sorted(set(keys))))))
        add("input.keys", str(year), fmt_int(len(set(keys))))
        add("input.teachers", str(year), fmt_int(sum(1 for e in staff if is_teacher(e))))
        add("input.principals", str(year), fmt_int(sum(1 for e in staff if is_school_leader(e))))

    # -------------------------------------------------------------- peer sets
    for level, tok in LEVELS:
        for year in YEARS:
            add("peer.orgs", "%d|%s" % (year, tok), fmt_int(len(corpus.peer_orgs(year, level, MIN_TEACHERS))))
        for frm, to in TRANSITIONS:
            for scope in SCOPES:
                rates = sorted(corpus.peer_rates(frm, to, level, MIN_TEACHERS, scope))
                key = "%d-%d|%s|%s" % (frm, to, tok, scope)
                add("peer.n", key, fmt_int(len(rates)))
                add("peer.median", key, fmt_real(median(rates)))
                add("peer.q1", key, fmt_real(quantile_at(rates, 0.25)))
                add("peer.q3", key, fmt_real(quantile_at(rates, 0.75)))

    # ----------------------------------------------------------- per school
    pooled_teacher = {}
    principal_changes = {}
    longest_tenure = {}

    for org, level in SCHOOLS:
        tok = dict(LEVELS)[level]
        add("school.level", org, fmt_text(tok))

        series = {}
        for scope in SCOPES:
            out = []
            for frm, to in TRANSITIONS:
                stayed, start, rate = retention_between(
                    corpus.staff_at(frm, org, scope), corpus.staff_at(to, org, scope))
                peers = corpus.peer_rates(frm, to, level, MIN_TEACHERS, scope)
                pct = None if rate is None else percentile_rank(peers, rate)
                out.append((frm, to, stayed, start, pct))
                key = "%s|%d-%d" % (org, frm, to)
                add("series.%s.stayed" % scope, key, fmt_int(stayed))
                add("series.%s.start" % scope, key, fmt_int(start))
                add("series.%s.percentile" % scope, key, fmt_real(pct))
            series[scope] = out

        pooled = [
            ("teacher", series["teacher"]),
            ("teacherWithAlt", series["teacherWithAlt"]),
            ("all", series["all"]),
            ("teacherSince2020", [s for s in series["teacher"] if s[0] >= 2020]),
        ]
        for name, rowset in pooled:
            stayed = sum(s[2] for s in rowset)
            start = sum(s[3] for s in rowset)
            add("pooled.%s.stayed" % name, org, fmt_int(stayed))
            add("pooled.%s.start" % name, org, fmt_int(start))
            if name == "teacher":
                pooled_teacher[org] = None if start == 0 else stayed / start

        for year in YEARS:
            pft = len(corpus.staff_at(year, org, "teacher"))
            alt = len(corpus.staff_at(year, org, "teacherWithAlt")) - pft
            add("alt.pft", "%s|%d" % (org, year), fmt_int(pft))
            add("alt.alt", "%s|%d" % (org, year), fmt_int(alt))

        entries = []
        for year in YEARS:
            names = " + ".join(sorted("%s %s" % (l.first, l.last) for l in corpus.leaders_at(year, org)))
            entries.append((year, names if names else "(none listed)"))
        runs = compress_timeline(entries)
        add("principal.runs", org, fmt_int(len(runs)))
        # Runs minus one: three runs over the window means two changes.
        add("principal.changes", org, fmt_int(max(0, len(runs) - 1)))
        for n, (names, frm, to) in enumerate(runs):
            add("principal.run", "%s|%02d" % (org, n), fmt_text("%d-%d|%s" % (frm, to, names)))
        principal_changes[org] = max(0, len(runs) - 1)
        longest_tenure[org] = max(to - frm + 1 for _, frm, to in runs)

        since = [y for y in YEARS if y >= 2020]
        base_keys = {employee_key(e) for e in corpus.staff_at(since[0], org, "teacher")}
        add("cohort2020.base", org, fmt_int(len(base_keys)))
        for year in since:
            now = {employee_key(e) for e in corpus.staff_at(year, org, "teacher")}
            add("cohort2020.remaining", "%s|%d" % (org, year), fmt_int(len(base_keys & now)))

        add("teacherCount", org, fmt_int(len(corpus.staff_at(YEARS[-1], org, "teacher"))))

        dest = {"stayed": 0, "movedSchool": 0, "changedRole": 0, "leftDistrict": 0}
        for frm, to in TRANSITIONS:
            for k in {employee_key(e) for e in corpus.staff_at(frm, org, "teacher")}:
                where = corpus.teacher_org[to].get(k)
                if where == org:
                    dest["stayed"] += 1
                elif where is not None:
                    dest["movedSchool"] += 1
                elif k in corpus.payroll[to]:
                    dest["changedRole"] += 1
                else:
                    dest["leftDistrict"] += 1
        for name in ("stayed", "movedSchool", "changedRole", "leftDistrict"):
            add("destination.%s" % name, org, fmt_int(dest[name]))

        # A transition belongs to whoever held the building in its starting
        # October — that is the year whose staff either stayed or left.
        eras = []
        for names, frm, to in runs:
            mine = [s for s in series["teacher"] if frm <= s[0] <= to]
            stayed = sum(s[2] for s in mine)
            start = sum(s[3] for s in mine)
            if start > 0:
                eras.append((names, frm, to, stayed, start))
        add("era.n", org, fmt_int(len(eras)))
        for n, (names, frm, to, stayed, start) in enumerate(eras):
            add("era.span", "%s|%02d" % (org, n), fmt_text("%d-%d|%s" % (frm, to, names)))
            add("era.stayed", "%s|%02d" % (org, n), fmt_int(stayed))
            add("era.start", "%s|%02d" % (org, n), fmt_int(start))

    # ------------------------------- does principal churn predict retention?
    rankable = [org for org in K8_ORGS if pooled_teacher.get(org) is not None]
    for org in rankable:
        add("correlation.churn", org, fmt_int(principal_changes[org]))
        add("correlation.tenure", org, fmt_int(longest_tenure[org]))
    add("correlation.n", "-", fmt_int(len(rankable)))
    add("correlation.churnVsRetention", "-",
        fmt_real(pearson([principal_changes[o] for o in rankable], [pooled_teacher[o] for o in rankable])))
    add("correlation.tenureVsRetention", "-",
        fmt_real(pearson([longest_tenure[o] for o in rankable], [pooled_teacher[o] for o in rankable])))

    return rows


def main():
    out_dir = os.path.join(ROOT, "calc", "out")
    os.makedirs(out_dir, exist_ok=True)
    tsv = serialize(build())
    with open(os.path.join(out_dir, "retention-python.tsv"), "w", encoding="utf-8", newline="\n") as fh:
        fh.write(tsv)
    sys.stdout.write("wrote calc/out/retention-python.tsv (%d facts)\n" % (tsv.count("\n") - 1))


if __name__ == "__main__":
    main()