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

Pennsylvania and Alabama use different tests with different cut scores, so their
published proficiency rates are not comparable. The Stanford Education Data
Archive links every state's tests to one national scale, which is what makes a
West Philadelphia K-8 and a Tuscaloosa City school comparable at all. Both
measures here are deviations from the national average, in grade levels.

Third of three implementations; see calc/CONTRACT.md. Standard library only.

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

from __future__ import annotations

import csv
import json
import math
import os
import re
import sys

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

csv.field_size_limit(sys.maxsize)

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
SEDA = os.path.join(ROOT, "seda")
DEMOGRAPHICS = os.path.join(ROOT, "demographics")

# LEAs the current federal enrolment counts are read for.
CCD_LEAS = ["4218990", "0103360", "3600091", "3600092"]

# CCD codes: 99 is all students, 2 is Black students.
CCD_ALL, CCD_BLACK = 99, 2

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

PHILA_LEA = "4218990"
TUSC_CITY_LEA = "103360"
TUSC_COUNTY_LEA = "103390"

DISTRICT_NAMES = {
    PHILA_LEA: "Philadelphia City SD",
    TUSC_CITY_LEA: "Tuscaloosa City",
    TUSC_COUNTY_LEA: "Tuscaloosa County",
}

K8 = [
    ("421899006528", "Harrington"),
    ("421899000619", "Penn Alexander"),
    ("421899006530", "Lea"),
    ("421899003782", "Comegys"),
    ("421899003783", "Powel"),
    ("421899004732", "Alain Locke"),
    ("421899003649", "Meredith"),
    ("421899003650", "Nebinger"),
    ("421899005186", "Greenfield"),
    ("421899003805", "Marian Anderson"),
    ("421899003618", "Bache-Martin"),
]

# SEDA files all of New York City under one synthetic citywide LEA rather than
# the 32 geographic-district LEAs the federal collections use. The five-digit
# school number is the same in both, so an id crosswalks by swapping the prefix.
NYC_SEDA_LEA = "3620580"

NYC = [
    ("360009102074", "PS 11 Purvis J. Behan"),
    ("360009102069", "PS 9 Teunis G. Bergen"),
    ("360009202818", "PS 321 William Penn"),
    ("360009202403", "PS 107 John W. Kimball"),
    ("360009202175", "PS 39 Henry Bristow"),
    ("360009202226", "MS 51 William Alexander"),
    ("360009202557", "PS 154 Windsor Terrace"),
    ("360009202072", "PS 10 Magnet School"),
    ("360009206406", "PS 118 Maurice Sendak"),
    ("360009102490", "PS 133 William A. Butler"),
    ("360009102793", "PS 282 Park Slope"),
]


def nyc_seda_id(nces):
    """A federal NCES school id as SEDA writes it for New York City."""
    if not re.match(r"^\d{12}$", nces):
        raise ValueError("nyc_seda_id: not a 12-digit NCES id: %r" % (nces,))
    return NYC_SEDA_LEA + nces[-5:]


SOURCES = [
    ("achievement", "seda_school_pool_cs_4.1.csv",
     ["sedasch", "sedaschname", "cs_mn_avg_eb", "cs_mn_grd_eb"]),
    ("covariates", "seda_cov_school_pool_4.1.csv",
     ["sedasch", "sedalea", "schcity", "schnam", "level", "perblk", "totenrl", "perecd"]),
    ("districts", "seda2022_dist.csv",
     ["sedaadmin", "sedaadminname", "subject", "subgroup", "ys_mn19_eb", "ys_mn22_eb", "ys_chg_eb"]),
]


def current_enrollment(rows):
    """
    Fold the long-format CCD extract into one record per school.

    Only the sex = 99 rows are read: the per-sex rows are the same students
    counted again. Suppressed cells are written as negative numbers and must
    not become a negative share.
    """
    total, black = {}, {}
    for r in rows:
        if r.get("ncessch") is None or r.get("sex") != 99:
            continue
        e = r.get("enrollment")
        if not isinstance(e, (int, float)) or e < 0:
            continue
        if r.get("race") == CCD_ALL:
            total[r["ncessch"]] = e
        elif r.get("race") == CCD_BLACK:
            black[r["ncessch"]] = e
    out = {}
    for sid in set(total) | set(black):
        out[sid] = (total.get(sid), black.get(sid))
    return out


def read(rel):
    path = os.path.join(SEDA, rel)
    if not os.path.exists(path):
        raise SystemExit("missing %s — run ./fetch-seda.sh" % path)
    with open(path, encoding="latin-1", newline="") as fh:
        return list(csv.DictReader(fh))


# A JavaScript-compatible numeric literal. Python's float() accepts things
# JavaScript's Number() does not — "1_000", "nan", "infinity" — so the shape is
# checked before conversion rather than after.
_NUMERIC = re.compile(r"^[+-]?(\d+\.?\d*|\.\d+)([eE][+-]?\d+)?$")


def num(v):
    """Parse a SEDA decimal. Values may be blank, 'NA', or written as '.075'."""
    if v is None:
        return None
    t = v.strip()
    if t == "" or t.upper() == "NA" or not _NUMERIC.match(t):
        return None
    n = float(t)
    return n if math.isfinite(n) else None


def js_round(x):
    """
    JavaScript's Math.round: halves go up.

    Python's round() and R's round() both round halves to even, so a percentile
    of exactly 46.5 would come out 46 there and 47 in the original. This is the
    single most likely place for three implementations to disagree by one.
    """
    return math.floor(x + 0.5)


def percentile_of(sorted_values, v):
    """Where `v` falls in an ascending distribution, 0-100; None if unrankable."""
    if v is None or not sorted_values:
        return None
    lo, hi = 0, len(sorted_values)
    while lo < hi:
        mid = (lo + hi) // 2
        if sorted_values[mid] < v:
            lo = mid + 1
        else:
            hi = mid
    return js_round((lo / len(sorted_values)) * 100)


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 pearson(xs, ys):
    """Accumulated left to right, so all three implementations land on the same bits."""
    n = min(len(xs), len(ys))
    if n < 2:
        return None
    sx = 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)


def join_seda(achievement, covariates, keep, labels, levels, growths):
    """
    Join achievement to covariates, keeping schools that pass `keep`.

    Percentiles are against the whole national file, not the kept subset: the
    point is to place these schools among all US schools.
    """
    cov_by_id = {}
    for c in covariates:
        cid = c.get("sedasch")
        if cid is not None:
            cov_by_id[cid] = c

    out = []
    for a in achievement:
        sid = a.get("sedasch")
        if sid is None:
            continue
        level = num(a.get("cs_mn_avg_eb"))
        if level is None:
            continue
        cov = cov_by_id.get(sid)
        if cov is None or not keep(cov):
            continue
        growth = num(a.get("cs_mn_grd_eb"))
        # `??` in the original coalesces on absence, not on emptiness: a blank
        # school name stays blank rather than falling through to the next one.
        if sid in labels:
            name = labels[sid]
        elif cov.get("schnam") is not None:
            name = cov["schnam"].strip()
        elif a.get("sedaschname") is not None:
            name = a["sedaschname"].strip()
        else:
            name = sid
        out.append({
            "id": sid,
            "name": name,
            "district": (cov.get("sedalea") or "").strip(),
            "gradeLevel": (cov.get("level") or "").strip(),
            "level": level,
            "growth": growth,
            "levelPct": percentile_of(levels, level),
            "growthPct": percentile_of(growths, growth),
            "pctBlack": num(cov.get("perblk")),
            "enrollment": num(cov.get("totenrl")),
        })
    # Descending by level. Python's sort is stable, as the original's is, so
    # schools with the same level keep their file order.
    out.sort(key=lambda s: -s["level"])
    return out


def build():
    rows = []

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

    parsed = {}
    for name, rel, columns in SOURCES:
        parsed[name] = read(rel)
        add("input.rows", name, fmt_int(len(parsed[name])))
        dump = "\n".join("\x1f".join((r.get(c) or "") for c in columns) for r in parsed[name])
        add("input.digest", name, fmt_text(sha256(dump)))

    # Current demographics, so no decade-old racial composition is emitted alone.
    ccd_rows = []
    for lea in CCD_LEAS:
        path = os.path.join(DEMOGRAPHICS, "enrollment_%s.json" % lea)
        if not os.path.exists(path):
            raise SystemExit("missing %s — run ./fetch-demographics.sh" % path)
        with open(path, encoding="utf-8") as fh:
            body = fh.read()
        add("input.filedigest", "enrollment_%s" % lea, fmt_text(sha256(body)))
        ccd_rows.extend(json.loads(body).get("results", []))
    now_by_nces = current_enrollment(ccd_rows)
    add("input.rows", "ccd_enrollment", fmt_int(len(ccd_rows)))
    add("input.schools", "ccd_enrollment", fmt_int(len(now_by_nces)))

    achievement = parsed["achievement"]
    covariates = parsed["covariates"]
    dist_rows = parsed["districts"]

    levels = sorted(v for v in (num(a.get("cs_mn_avg_eb")) for a in achievement) if v is not None)
    growths = sorted(v for v in (num(a.get("cs_mn_grd_eb")) for a in achievement) if v is not None)

    labels = {nces: label for nces, label in K8}
    wanted = set(labels)
    nyc_labels = {nyc_seda_id(nces): label for nces, label in NYC}
    nyc_federal = {nyc_seda_id(nces): nces for nces, _ in NYC}
    nyc_wanted = set(nyc_labels)
    def is_brooklyn(c):
        return c.get("sedalea") == NYC_SEDA_LEA and (c.get("schcity") or "").upper() == "BROOKLYN"

    def now_of(sid):
        """
        The current federal counts for one SEDA school, or None.

        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.
        """
        fed = nyc_federal.get(sid, ("0" + sid) if len(sid) == 11 else sid)
        return now_by_nces.get(fed)

    def school_dump(schools):
        """
        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.
        """
        lines = []
        for i, s in enumerate(schools):
            now = now_of(s["id"])
            lines.append("\x1f".join([
                s["id"], s["name"], s["district"], s["gradeLevel"],
                fmt_real(s["level"]), fmt_real(s["growth"]),
                NA if s["levelPct"] is None else fmt_int(s["levelPct"]),
                NA if s["growthPct"] is None else fmt_int(s["growthPct"]),
                fmt_real(s["pctBlack"]), fmt_real(s["enrollment"]),
                NA if now is None or now[0] is None else fmt_int(int(now[0])),
                NA if now is None or now[1] is None else fmt_int(int(now[1])),
                fmt_int(i),
            ]))
        return "\n".join(lines)

    groups = [
        ("phila", join_seda(achievement, covariates, lambda c: (c.get("sedasch") or "") in wanted, labels, levels, growths), True),
        ("tuscaloosa", join_seda(achievement, covariates, lambda c: c.get("sedalea") == TUSC_CITY_LEA, {}, levels, growths), True),
        ("brooklyn", join_seda(achievement, covariates, lambda c: (c.get("sedasch") or "") in nyc_wanted, nyc_labels, levels, growths), True),
        ("allBrooklyn", join_seda(achievement, covariates, is_brooklyn, {}, levels, growths), False),
    ]

    for group, schools, per_school in groups:
        add("seda.n", group, fmt_int(len(schools)))
        # 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, fmt_real(median([s["level"] for s in schools])))
        add("seda.withGrowth", group, fmt_int(sum(1 for s in schools if s["growth"] is not None)))
        add("seda.digest", group, fmt_text(sha256(school_dump(schools))))
        if not per_school:
            continue
        for i, s in enumerate(schools):
            key = "%s|%s" % (group, s["id"])
            add("seda.level", key, fmt_real(s["level"]))
            add("seda.growth", key, fmt_real(s["growth"]))
            add("seda.levelPct", key, NA if s["levelPct"] is None else fmt_int(s["levelPct"]))
            add("seda.growthPct", key, NA if s["growthPct"] is None else fmt_int(s["growthPct"]))
            add("seda.pctBlack", key, fmt_real(s["pctBlack"]))
            add("seda.enrollment", key, fmt_real(s["enrollment"]))
            add("seda.name", key, fmt_text(s["name"]))
            add("seda.district", key, fmt_text(s["district"]))
            add("seda.gradeLevel", key, fmt_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, fmt_int(i))
            now = now_of(s["id"])
            add("seda.nowEnrollment", key, NA if now is None or now[0] is None else fmt_int(int(now[0])))
            add("seda.nowBlack", key, NA if now is None or now[1] is None else fmt_int(int(now[1])))

    add("national.levels", "-", fmt_int(len(levels)))
    add("national.growths", "-", fmt_int(len(growths)))

    # 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.
    for place, lea in (("Philadelphia", PHILA_LEA), ("Tuscaloosa_City", TUSC_CITY_LEA), ("New_York_City", NYC_SEDA_LEA)):
        ids = {c.get("sedasch") for c in covariates if c.get("sedalea") == lea}
        subset = [a for a in achievement if (a.get("sedasch") or "") in ids]
        add("coverage.schools", place, fmt_int(len(subset)))
        add("coverage.level", place, fmt_int(sum(1 for a in subset if num(a.get("cs_mn_avg_eb")) is not None)))
        add("coverage.growth", place, fmt_int(sum(1 for a in subset if num(a.get("cs_mn_grd_eb")) is not None)))

    # New York is absent from the 2022 district release entirely.
    add("district.newYorkRows", "-", fmt_int(sum(1 for r in dist_rows if r.get("stateabb") == "NY")))

    # ---------------------------------------------------------- the districts
    for r in dist_rows:
        rid = r.get("sedaadmin") or ""
        if r.get("subgroup") != "all" or rid not in DISTRICT_NAMES:
            continue
        subject = "math" if r.get("subject") == "mth" else "reading"
        add("district.level2019", "%s|%s" % (rid, subject), fmt_real(num(r.get("ys_mn19_eb"))))
        add("district.level2022", "%s|%s" % (rid, subject), fmt_real(num(r.get("ys_mn22_eb"))))
        add("district.change", "%s|%s" % (rid, subject), fmt_real(num(r.get("ys_chg_eb"))))
    for rid, name in DISTRICT_NAMES.items():
        add("district.name", rid, fmt_text(name))

    for subj in ("mth", "rla"):
        subject = "math" if subj == "mth" else "reading"
        for_subject = [r for r in dist_rows if r.get("subgroup") == "all" and r.get("subject") == subj]
        add("usMedian.level2019", subject,
            fmt_real(median([v for v in (num(r.get("ys_mn19_eb")) for r in for_subject) if v is not None])))
        add("usMedian.change", subject,
            fmt_real(median([v for v in (num(r.get("ys_chg_eb")) for r in for_subject) if v is not None])))
        # Every district row for the subject, including the ones with no
        # estimate: this is the denominator the page prints.
        add("usMedian.n", subject, fmt_int(len(for_subject)))

    # ------------------------------------------------- how much is just intake
    poverty = {}
    for c in covariates:
        cid = c.get("sedasch")
        p = num(c.get("perecd"))
        if cid is not None and p is not None and 0 <= p <= 1:
            poverty[cid] = p
    add("poverty.schools", "-", fmt_int(len(poverty)))

    # 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", "-", fmt_int(sum(1 for v in poverty.values() if v >= POVERTY_CEILING)))
    add("poverty.ceilingAt", "-", fmt_real(POVERTY_CEILING))

    for name, column in (("level", "cs_mn_avg_eb"), ("growth", "cs_mn_grd_eb")):
        xs, ys = [], []
        for a in achievement:
            sid = a.get("sedasch")
            v = num(a.get(column))
            pov = None if sid is None else poverty.get(sid)
            if v is None or pov is None:
                continue
            xs.append(pov)
            ys.append(v)
        add("poverty.%s.r" % name, "-", fmt_real(pearson(xs, ys)))
        add("poverty.%s.n" % name, "-", fmt_int(len(xs)))

    # ------------------------------------------- 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.
    level_of = {}
    for a in achievement:
        sid = a.get("sedasch")
        v = num(a.get("cs_mn_avg_eb"))
        if sid is not None and v is not None:
            level_of[sid] = v

    def set_stats(name, ids):
        levels = [level_of[i] for i in ids if i in level_of]
        povs = [poverty[i] for i in ids if i in poverty]
        add("set.n", name, fmt_int(len(levels)))
        add("set.level", name, fmt_real(median(levels)))
        add("set.poverty", name, fmt_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", name, fmt_int(sum(1 for v in povs if v >= POVERTY_CEILING)))
        add("set.ceilingOf", name, fmt_int(len(povs)))

    def cov_ids_where(keep):
        return [c["sedasch"] for c in covariates if c.get("sedasch") is not None and keep(c)]

    set_stats("philaEleven", [nces for nces, _ in K8])
    set_stats("philaAll", cov_ids_where(lambda c: c.get("sedalea") == PHILA_LEA))
    set_stats("brooklynTen", [nyc_seda_id(nces) for nces, _ in NYC])
    set_stats("brooklynAll", cov_ids_where(is_brooklyn))
    set_stats("tuscaloosaCity", cov_ids_where(lambda c: c.get("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.
    xs, ys = [], []
    for a in achievement:
        sid = a.get("sedasch")
        v = num(a.get("cs_mn_avg_eb"))
        pv = None if sid is None else poverty.get(sid)
        if v is None or pv is None:
            continue
        xs.append(pv)
        ys.append(v)
    n = len(xs)
    sx = sy = 0.0
    for i in range(n):
        sx += xs[i]
        sy += ys[i]
    mx, my = sx / n, sy / n
    sxy = sxx = 0.0
    for i in range(n):
        dx = xs[i] - mx
        sxy += dx * (ys[i] - my)
        sxx += dx * dx
    slope = sxy / sxx
    add("povertyFit.slope", "-", fmt_real(slope))
    add("povertyFit.intercept", "-", fmt_real(my - slope * mx))
    add("povertyFit.n", "-", fmt_int(n))

    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, "seda-python.tsv"), "w", encoding="utf-8", newline="\n") as fh:
        fh.write(tsv)
    print("wrote calc/out/seda-python.tsv (%d facts)" % (tsv.count("\n") - 1))


if __name__ == "__main__":
    main()