#!/usr/bin/env python3
"""Reproduce the Evidence Lab income-versus-poverty model comparison.

The report compares three ordinary-school model families using the 2024-25
processed assessment files. Outcomes and model weights are based on students
with reported Level 1-4 results. Five-fold predictions are produced without
using held-out outcomes when fitting each fold.
"""

from __future__ import annotations

import csv
import math
from pathlib import Path

import numpy as np
from sklearn.model_selection import KFold


ROOT = Path(__file__).resolve().parents[1]
DATA_DIR = ROOT / "data" / "processed"
OUTPUT = ROOT / "docs" / "income_poverty_reassessment_note_2026-02-20.txt"
DATASETS = {
    "ELA": DATA_DIR / "SchoolDataWithAddressesAndCensusSES.csv",
    "Math": DATA_DIR / "SchoolDataMathWithAddressesAndCensusSES.csv",
    "Science": DATA_DIR / "SchoolDataScienceWithAddressesAndCensusSES.csv",
}

FEATURES = {
    "BA+ + Attendance + Income": [
        "ba",
        "attendance",
        "household_income",
    ],
    "BA+ + Attendance + Poverty": [
        "ba",
        "attendance",
        "poverty",
    ],
    "BA+ + Attendance + Income + Poverty": [
        "ba",
        "attendance",
        "household_income",
        "poverty",
    ],
}


def number(value: object) -> float | None:
    try:
        text = str(value).strip().replace(",", "")
        return float(text) if text else None
    except (TypeError, ValueError):
        return None


def yes(value: object) -> bool:
    return str(value).strip().lower() in {"1", "true", "t", "yes", "y"}


def rate(value: float) -> float:
    return value * 100.0 if abs(value) <= 1.5 else value


def grade(value: object) -> str:
    text = str(value or "").strip().lower()
    if text in {"all", "all grade", "all grades"}:
        return "all"
    return text.replace("grade", "").strip()


def ordinary(row: dict[str, str]) -> bool:
    charter = yes(row.get("Is Charter School")) or str(row.get("ODE School Type", "")).strip().lower() == "charter"
    virtual = yes(row.get("Is Virtual School")) or str(row.get("ODE Virtual Status", "")).strip().lower() in {"focus virtual", "full virtual"}
    special = yes(row.get("Is Special Enrollment School"))
    return not (charter or virtual or special)


def load_school_rows(path: Path) -> list[dict[str, float]]:
    staged: list[tuple[tuple[str, str], str, float, float, dict[str, float]]] = []
    with path.open(newline="", encoding="utf-8-sig") as handle:
        for row in csv.DictReader(handle):
            if row.get("Student Group") != "Total Population (All Students)" or not ordinary(row):
                continue
            scored = number(row.get("Scored Performance Denominator"))
            proficient = number(row.get("Number Proficient"))
            values = {
                "ba": number(row.get("ACS_Ed_BA_or_Higher_Rate")),
                "attendance": number(row.get("Attendance_Percent_Regular_Attenders")),
                "household_income": number(row.get("ACS_Median_Household_Income")),
                "per_capita_income": number(row.get("ACS_Per_Capita_Income")),
                "poverty": number(row.get("Students Experiencing Poverty")),
            }
            if scored is None or scored <= 0 or proficient is None or any(v is None for v in values.values()):
                continue
            values["ba"] = rate(values["ba"])
            values["poverty"] = rate(values["poverty"])
            key = (str(row.get("District ID", "")), str(row.get("School ID", "")))
            staged.append((key, grade(row.get("Grade Level")), scored, proficient, values))

    has_grades: dict[tuple[str, str], bool] = {}
    for key, grade_code, *_ in staged:
        has_grades[key] = has_grades.get(key, False) or grade_code != "all"

    schools: dict[tuple[str, str], dict[str, float]] = {}
    for key, grade_code, scored, proficient, values in staged:
        if grade_code == "all" and has_grades.get(key):
            continue
        current = schools.setdefault(key, {"weight": 0.0, "proficient": 0.0, **{f"{name}_wx": 0.0 for name in values}})
        current["weight"] += scored
        current["proficient"] += proficient
        for name, value in values.items():
            current[f"{name}_wx"] += value * scored

    result = []
    for current in schools.values():
        weight = current["weight"]
        if weight <= 0:
            continue
        row = {"weight": weight, "outcome": 100.0 * current["proficient"] / weight}
        for name in ("ba", "attendance", "household_income", "per_capita_income", "poverty"):
            row[name] = current[f"{name}_wx"] / weight
        result.append(row)
    return result


def weighted_mean(values: np.ndarray, weights: np.ndarray) -> float:
    return float(np.sum(values * weights) / np.sum(weights))


def weighted_r2(actual: np.ndarray, predicted: np.ndarray, weights: np.ndarray) -> float:
    center = weighted_mean(actual, weights)
    residual = float(np.sum(weights * (actual - predicted) ** 2))
    total = float(np.sum(weights * (actual - center) ** 2))
    return math.nan if total <= 0 else 1.0 - residual / total


def weighted_corr(left: np.ndarray, right: np.ndarray, weights: np.ndarray) -> float:
    left_center = weighted_mean(left, weights)
    right_center = weighted_mean(right, weights)
    covariance = np.sum(weights * (left - left_center) * (right - right_center)) / np.sum(weights)
    left_var = np.sum(weights * (left - left_center) ** 2) / np.sum(weights)
    right_var = np.sum(weights * (right - right_center) ** 2) / np.sum(weights)
    return float(covariance / math.sqrt(left_var * right_var))


def fit_predict(train_x: np.ndarray, train_y: np.ndarray, train_w: np.ndarray, test_x: np.ndarray) -> np.ndarray:
    design = np.column_stack([np.ones(len(train_x)), train_x])
    test_design = np.column_stack([np.ones(len(test_x)), test_x])
    root_weight = np.sqrt(train_w)
    coefficients, *_ = np.linalg.lstsq(design * root_weight[:, None], train_y * root_weight, rcond=None)
    return test_design @ coefficients


def cross_validated_r2(rows: list[dict[str, float]], feature_names: list[str]) -> float:
    x = np.array([[row[name] for name in feature_names] for row in rows], dtype=float)
    y = np.array([row["outcome"] for row in rows], dtype=float)
    w = np.array([row["weight"] for row in rows], dtype=float)
    prediction = np.full(len(rows), np.nan)
    splitter = KFold(n_splits=5, shuffle=True, random_state=42)
    for train, test in splitter.split(x):
        prediction[test] = fit_predict(x[train], y[train], w[train], x[test])
    return weighted_r2(y, prediction, w)


def main() -> None:
    results: dict[str, dict[str, float]] = {}
    correlations: dict[str, tuple[float, float]] = {}
    counts: dict[str, tuple[int, float]] = {}
    for subject, path in DATASETS.items():
        rows = load_school_rows(path)
        results[subject] = {name: cross_validated_r2(rows, fields) for name, fields in FEATURES.items()}
        weights = np.array([row["weight"] for row in rows], dtype=float)
        poverty = np.array([row["poverty"] for row in rows], dtype=float)
        household = np.array([row["household_income"] for row in rows], dtype=float)
        per_capita = np.array([row["per_capita_income"] for row in rows], dtype=float)
        correlations[subject] = (
            weighted_corr(poverty, household, weights),
            weighted_corr(poverty, per_capita, weights),
        )
        counts[subject] = (len(rows), float(np.sum(weights)))

    lines = [
        "Income vs Poverty Reassessment Note",
        "",
        "Question",
        "How do school poverty and tract income compare when adult BA+, attendance, and ordinary-school scope are held consistent?",
        "",
        "Scope",
        "- School-level Total Population models for Oregon 2024-25.",
        "- Charter, virtual, and curated special-enrollment schools excluded.",
        "- Outcomes and weights use students with reported Level 1-4 results.",
        "- Five-fold out-of-sample predictions use a fixed random seed for reproducibility.",
        "- Percent Proficient is calculated from reported Level 3 and Level 4 counts.",
        "",
        "Main findings",
        "- Income and poverty overlap substantially but are not interchangeable.",
        "- School poverty adds substantial predictive information in ELA and Science and a smaller increment in Math.",
        "- Median household income contributes some additional information when school poverty is present, but much less than poverty contributes to the income model.",
        "",
        "Weighted correlations: Students Experiencing Poverty vs tract income",
    ]
    for subject in DATASETS:
        household, per_capita = correlations[subject]
        schools, students = counts[subject]
        lines.append(
            f"- {subject}: median household income r={household:.3f}; per-capita income r={per_capita:.3f}; "
            f"{schools:,} schools; {students:,.0f} students with reported Level 1-4 results."
        )
    lines.extend(["", "Five-fold cross-validated R^2"])
    for subject in DATASETS:
        lines.append(f"- {subject}:")
        for name in FEATURES:
            lines.append(f"  - {name}: {results[subject][name]:.4f}")
    lines.extend(
        [
            "",
            "Interpretation",
            "- ODE Students Experiencing Poverty is the more direct enrolled-student hardship measure.",
            "- Tract income remains useful as broader community context and adds a modest increment in some specifications.",
            "- Adult BA+ and attendance remain in every model so the income-versus-poverty comparison does not mistake their shared signal for a poverty or income effect.",
            "- These are descriptive school-level associations, not causal estimates.",
        ]
    )
    OUTPUT.write_text("\n".join(lines) + "\n", encoding="utf-8")
    print(f"Wrote {OUTPUT}")


if __name__ == "__main__":
    main()