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

Philadelphia and New York both publish school attendance, and they publish
mirrored forms of the same fact: Philadelphia the share of students attending at
least 90% of days, New York the share chronically absent. Everything here is
normalized to the second, which is the federal standard metric. The federal
Civil Rights Data Collection then supplies the one year in which all three
cities — Philadelphia, Brooklyn and Tuscaloosa City — are measured the same way.

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

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

from __future__ import annotations

import csv
import io
import json
import math
import os

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

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
A = os.path.join(ROOT, "attendance")

NYC_YEARS = ["2016-17", "2017-18", "2018-19", "2019-20", "2020-21"]
CRDC_LEAS = ["4218990", "0103360", "3600091", "3600092"]
# 2017-18 is the last clean pre-pandemic collection; 2022-23 the most recent.
CRDC_YEARS = ["2017", "2022"]
TUSCALOOSA_SEDA_LEA = "103360"

K8 = [
    ("1300", "Harrington", "421899006528"),
    ("1280", "Penn Alexander", "421899000619"),
    ("1340", "Lea", "421899006530"),
    ("1260", "Comegys", "421899003782"),
    ("1390", "Powel", "421899003783"),
    ("1470", "Alain Locke", "421899004732"),
    ("2380", "Meredith", "421899003649"),
    ("2590", "Nebinger", "421899003650"),
    ("2470", "Greenfield", "421899005186"),
    ("2480", "Marian Anderson", "421899003805"),
    ("2210", "Bache-Martin", "421899003618"),
]

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

# name -> (path, columns every implementation reads out of it)
CSV_SOURCES = [
    ("ada_archive", "ex/Student_ADA_Yearly.csv",
     ["School Year", "ULCS Code", "Average Daily Attendance"]),
    ("ada_current", "Student_ADA_School.csv",
     ["School Year", "ULCS Code", "Average Daily Attendance (YTD)"]),
    ("attending90", "Student_Attendance_90_School.csv",
     ["School Year", "ULCS Code", "Category", "% with 90%+ Attendance (Yearly)"]),
    ("attending95_archive", "ex/SDP_95_Attendance_School_S.csv",
     ["School Year", "ULCS Code", "Category", "% with 95%+ Attendance (Yearly)"]),
    ("attending95_current", "Student_Attendance_95_School.csv",
     ["School Year", "ULCS Code", "Category", "% with 95%+ Attendance (Yearly)"]),
    ("teacher_ada", "ex/Archive_Teacher_95_ADA_2012-13_2022-23/OD_Teacher_Attendance_Schools_Suppressed_20240328.csv",
     ["school_year", "ulcscode", "pct_att_ada"]),
    ("seda_covariates", "../seda/seda_cov_school_pool_4.1.csv",
     ["sedasch", "sedalea", "schnam"]),
]


# ----------------------------------------------------------------- reading

def read_csv(rel):
    """
    Read one attendance CSV.

    The files are latin-1, and one of them carries a byte-order mark. Left in
    place the mark becomes part of the first column's name, so every lookup on
    that column silently returns nothing — a failure that produces an empty
    chart rather than an error.
    """
    path = os.path.join(A, rel)
    if not os.path.exists(path):
        raise SystemExit("missing %s — run ./fetch-attendance.sh" % path)
    with open(path, encoding="latin-1", newline="") as fh:
        text = fh.read()
    if text.startswith(""):
        text = text[1:]
    elif text.startswith(""):
        text = text[3:]
    return list(csv.DictReader(io.StringIO(text)))


def read_json(rel):
    path = os.path.join(A, rel)
    if not os.path.exists(path):
        raise SystemExit("missing %s — run ./fetch-attendance.sh and ./fetch-crdc.sh" % path)
    with open(path, encoding="utf-8") as fh:
        return fh.read()


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

def normalize_ulcs(code):
    """SDP school codes appear both zero-padded and unpadded across files."""
    t = str(code).strip().lstrip("0")
    return "0" if t == "" else t


def parse_value(v):
    """A published cell as a number. Suppressed cells are written as 's'."""
    if v is None:
        return None
    t = v.strip()
    if t == "" or t.lower() == "s" or t.lower() == "n/a":
        return None
    try:
        n = float(t)
    except ValueError:
        return None
    return n if math.isfinite(n) else None


def series_from_rows(rows, year_key, code_key, value_key, keep=lambda r: True):
    """Collapse long-format rows into school -> year -> value."""
    out = {}
    for r in rows:
        if not keep(r):
            continue
        v = parse_value(r.get(value_key))
        if v is None:
            continue
        code = normalize_ulcs(r.get(code_key) or "")
        year = (r.get(year_key) or "").strip()
        if code == "" or year == "":
            continue
        out.setdefault(code, {})[year] = v
    return out


def all_students(r):
    return r.get("Category") == "All Students"


def chronic_from_attending90(pct):
    """Philadelphia publishes the attending share; the standard metric is its complement."""
    return 100 - pct


def federal_rate(absent, enrolled):
    """
    A CRDC absentee count over its CCD enrolment, as a percentage.

    CRDC encodes suppressed or not-applicable cells as negative numbers
    (-3, -5, -9). Those must not be divided; nor may a rate above 100 be
    reported, which would mean the count and the denominator came from
    mismatched records.
    """
    if absent is None or absent < 0:
        return None
    if enrolled is None or enrolled <= 0:
        return None
    rate = (absent / enrolled) * 100
    return None if rate > 100 else rate


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)


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

def merge(old, new):
    """Two releases of the same series; the newer file wins year by year."""
    out = {k: dict(v) for k, v in old.items()}
    for code, years in new.items():
        out.setdefault(code, {}).update(years)
    return out


def build():
    rows = []

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

    # ------------------------------------------------------------------ inputs
    parsed = {}
    for name, rel, columns in CSV_SOURCES:
        parsed[name] = read_csv(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)))

    # The JSON sources are hashed as raw bytes rather than as parsed columns:
    # there is no cross-language question about how to read a JSON string, and
    # what each implementation makes of the numbers is checked by the facts below.
    json_files = [("nyc_%s" % y, "nyc_attendance_%s.json" % y) for y in NYC_YEARS]
    for year in CRDC_YEARS:
        for lea in CRDC_LEAS:
            json_files.append(("crdc_absent_%s_%s" % (year, lea), "crdc_absent_%s_%s.json" % (year, lea)))
            json_files.append(("crdc_enrolled_%s_%s" % (year, lea), "crdc_enrolled_%s_%s.json" % (year, lea)))
    raw_json = {}
    for name, rel in json_files:
        raw_json[name] = read_json(rel)
        add("input.filedigest", name, fmt_text(sha256(raw_json[name])))

    # ------------------------------------------------------------ Philadelphia
    phl_ada = merge(
        series_from_rows(parsed["ada_archive"], "School Year", "ULCS Code", "Average Daily Attendance"),
        series_from_rows(parsed["ada_current"], "School Year", "ULCS Code", "Average Daily Attendance (YTD)"),
    )
    attending90 = series_from_rows(
        parsed["attending90"], "School Year", "ULCS Code", "% with 90%+ Attendance (Yearly)", all_students)
    phl_chronic = {code: {y: chronic_from_attending90(v) for y, v in years.items()}
                   for code, years in attending90.items()}
    phl_95 = merge(
        series_from_rows(parsed["attending95_archive"], "School Year", "ULCS Code",
                         "% with 95%+ Attendance (Yearly)", all_students),
        series_from_rows(parsed["attending95_current"], "School Year", "ULCS Code",
                         "% with 95%+ Attendance (Yearly)", all_students),
    )
    teacher_ada = series_from_rows(parsed["teacher_ada"], "school_year", "ulcscode", "pct_att_ada")

    series = [("ada", phl_ada), ("chronic", phl_chronic), ("a95", phl_95), ("teacherAda", teacher_ada)]

    for org, _label, _nces in K8:
        code = normalize_ulcs(org)
        for name, src in series:
            for year, value in src.get(code, {}).items():
                add("phl.%s" % name, "%s|%s" % (org, year), fmt_real(value))

    for name, src in series:
        if name == "teacherAda":
            continue
        by_year = {}
        for years in src.values():
            for y, v in years.items():
                by_year.setdefault(y, []).append(v)
        for year, values in by_year.items():
            add("phlMedian.%s" % name, year, fmt_real(median(values)))
            add("phlCount.%s" % name, year, fmt_int(len(values)))

    # ---------------------------------------------------------------- New York
    nyc_by_year = {}
    for y in NYC_YEARS:
        by_dbn = {}
        for r in json.loads(raw_json["nyc_%s" % y]):
            ada = parse_value(r.get("attendance"))
            chronic = parse_value(r.get("chronically_absent_1"))
            if r.get("dbn") is None or ada is None or chronic is None:
                continue
            by_dbn[r["dbn"]] = (ada, chronic)
        nyc_by_year[y] = by_dbn

    for dbn, _nces, _label in NYC:
        for y in NYC_YEARS:
            hit = nyc_by_year[y].get(dbn)
            if hit is None:
                continue
            add("nyc.ada", "%s|%s" % (dbn, y), fmt_real(hit[0]))
            add("nyc.chronic", "%s|%s" % (dbn, y), fmt_real(hit[1]))

    for y in NYC_YEARS:
        all_rows = list(nyc_by_year[y].values())
        add("nycMedian.ada", y, fmt_real(median([r[0] for r in all_rows])))
        add("nycMedian.chronic", y, fmt_real(median([r[1] for r in all_rows])))
        add("nycMedian.n", y, fmt_int(len(all_rows)))

    # ----------------------------------------------- federal, all three cities
    targets = {}
    for org, label, nces in K8:
        targets[nces] = ("Philadelphia", label)
    for dbn, nces, label in NYC:
        targets[nces] = ("Brooklyn", label)
    # SEDA drops Alabama's leading zero from the NCES id; CRDC keeps all twelve.
    for r in parsed["seda_covariates"]:
        if r.get("sedalea") != TUSCALOOSA_SEDA_LEA:
            continue
        sid = r.get("sedasch")
        if not sid:
            continue
        targets["0" + sid if len(sid) == 11 else sid] = ("Tuscaloosa City", (r.get("schnam") or sid).strip())

    for nces, (city, label) in targets.items():
        add("federal.city", nces, fmt_text(city.replace(" ", "_")))
        add("federal.label", nces, fmt_text(label))

    for year in CRDC_YEARS:
        crdc_absent, crdc_enrolled = {}, {}
        for lea in CRDC_LEAS:
            for r in json.loads(raw_json["crdc_absent_%s_%s" % (year, lea)]).get("results", []):
                # CRDC returns two rows per school even with every other
                # subgroup pinned to 99: one for all students and one for the
                # homeless subgroup. Without this filter the homeless count
                # silently becomes the numerator.
                if r.get("homeless") is not None and r.get("homeless") != 99:
                    continue
                if r.get("ncessch") is not None and isinstance(r.get("students_chronically_absent"), (int, float)):
                    crdc_absent[r["ncessch"]] = r["students_chronically_absent"]
            for r in json.loads(raw_json["crdc_enrolled_%s_%s" % (year, lea)]).get("results", []):
                if r.get("ncessch") is not None and isinstance(r.get("enrollment"), (int, float)):
                    crdc_enrolled[r["ncessch"]] = r["enrollment"]

        by_city = {}
        for nces, (city, label) in targets.items():
            rate = federal_rate(crdc_absent.get(nces), crdc_enrolled.get(nces))
            if rate is None:
                continue
            add("federal.rate", "%s|%s" % (year, nces), fmt_real(rate))
            by_city.setdefault(city, []).append(rate)

        for city, rates in by_city.items():
            key = "%s|%s" % (year, city.replace(" ", "_"))
            add("federalMedian.median", key, fmt_real(median(rates)))
            add("federalMedian.min", key, fmt_real(min(rates)))
            add("federalMedian.max", key, fmt_real(max(rates)))
            add("federalMedian.n", key, fmt_int(len(rates)))

    # ------------------------------------------------------------ correlation
    chronics, teachers = [], []
    for org, _label, _nces in K8:
        code = normalize_ulcs(org)
        chronic = phl_chronic.get(code, {}).get("2024-2025")
        # Year keys are compared as strings, exactly as the original does; they
        # are all ASCII digits, so every language orders them the same way.
        teacher = median([v for y, v in teacher_ada.get(code, {}).items() if y >= "2016"])
        if chronic is None or teacher is None:
            continue
        chronics.append(chronic)
        teachers.append(teacher)
        add("correlation.chronic", org, fmt_real(chronic))
        add("correlation.teacher", org, fmt_real(teacher))
    add("correlation.n", "-", fmt_int(len(chronics)))
    add("correlation.chronicVsTeacherAttendance", "-", fmt_real(pearson(chronics, teachers)))

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


if __name__ == "__main__":
    main()