#!/usr/bin/env python3
"""Reproduce BA+ slopes within school-poverty bands for Evidence Lab."""

from __future__ import annotations

from pathlib import Path

import numpy as np

from report_income_poverty_reassessment import DATASETS, load_school_rows, weighted_corr, weighted_mean


ROOT = Path(__file__).resolve().parents[1]
OUTPUT = ROOT / "docs" / "ba_signal_in_high_poverty_summary_2026-02-21.txt"


def weighted_fit(rows: list[dict[str, float]], features: list[str]) -> tuple[np.ndarray, float]:
    x = np.array([[row[name] for name in features] 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)
    design = np.column_stack([np.ones(len(x)), x])
    root_weight = np.sqrt(w)
    beta, *_ = np.linalg.lstsq(design * root_weight[:, None], y * root_weight, rcond=None)
    prediction = design @ beta
    center = weighted_mean(y, w)
    total = np.sum(w * (y - center) ** 2)
    residual = np.sum(w * (y - prediction) ** 2)
    return beta, float(1.0 - residual / total)


def band(rows: list[dict[str, float]], lower: float | None, upper: float | None) -> list[dict[str, float]]:
    return [
        row
        for row in rows
        if (lower is None or row["poverty"] >= lower)
        and (upper is None or row["poverty"] < upper)
    ]


def bivariate(rows: list[dict[str, float]], predictor: str) -> tuple[float, float]:
    x = np.array([row[predictor] 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)
    beta, _ = weighted_fit(rows, [predictor])
    return weighted_corr(x, y, w), float(beta[1] * 10.0)


def adjusted_ba(rows: list[dict[str, float]]) -> tuple[float, float]:
    _, base_r2 = weighted_fit(rows, ["poverty", "attendance"])
    beta, full_r2 = weighted_fit(rows, ["poverty", "attendance", "ba"])
    return float(beta[3] * 10.0), float(full_r2 - base_r2)


def main() -> None:
    analyses = {subject: load_school_rows(path) for subject, path in DATASETS.items()}
    band_specs = [
        ("Overall", None, None),
        ("Low poverty (<20%)", None, 20.0),
        ("Middle poverty (20-40%)", 20.0, 40.0),
        ("High poverty (>=40%)", 40.0, None),
        ("Very high poverty (>=50%)", 50.0, None),
    ]

    lines = [
        "BA+ Signal Across School-Poverty Bands",
        "",
        "Question",
        "Does the association between community adult BA+ and school proficiency remain equally strong across the school-poverty range?",
        "",
        "Scope",
        "- Oregon 2024-25 school-level Total Population results.",
        "- Charter, virtual, and curated special-enrollment schools excluded.",
        "- Percent Proficient and model weights use students with reported Level 1-4 results.",
        "- BA+ is the ACS tract share of adults with a bachelor's degree or higher.",
        "- Poverty is ODE Students Experiencing Poverty for the enrolled school population.",
        "",
        "Main findings",
        "- BA+ is positively associated with proficiency statewide and within every broad poverty band.",
        "- The BA+ slope is substantially steeper in low-poverty schools than in high-poverty schools.",
        "- Inside high-poverty schools, BA+ retains a modest positive association after poverty level and attendance are included.",
        "- These patterns describe association, not a causal BA+ effect.",
        "",
        "Statewide school-poverty gradient",
    ]
    for subject, rows in analyses.items():
        correlation, slope = bivariate(rows, "poverty")
        lines.append(f"- {subject}: {slope:+.2f} proficiency points per +10 poverty points (r={correlation:+.3f}).")

    lines.extend(["", "BA+ association by poverty band"])
    for label, lower, upper in band_specs:
        lines.append(f"- {label}:")
        for subject, rows in analyses.items():
            selected = band(rows, lower, upper)
            correlation, slope = bivariate(selected, "ba")
            students = sum(row["weight"] for row in selected)
            lines.append(
                f"  - {subject}: {slope:+.2f} proficiency points per +10 BA+ points "
                f"(r={correlation:+.3f}; {len(selected):,} schools; {students:,.0f} scored students)."
            )

    lines.extend(["", "Adjusted BA+ association inside high-poverty bands"])
    for label, lower in ((">=40% poverty", 40.0), (">=50% poverty", 50.0)):
        lines.append(f"- {label}:")
        for subject, rows in analyses.items():
            selected = band(rows, lower, None)
            slope, delta_r2 = adjusted_ba(selected)
            lines.append(f"  - {subject}: adjusted BA+ slope {slope:+.2f} per +10 BA+ points; added R^2={delta_r2:.3f}.")

    lines.extend(
        [
            "",
            "Interpretation",
            "- Community educational context remains informative after schools are separated by enrolled-student poverty.",
            "- The association is not uniform: concentrated poverty coincides with a much smaller measured BA+ advantage.",
            "- BA+ is attached to school-site geography rather than individual families or exact attendance areas, so band-specific estimates retain geographic measurement error.",
        ]
    )
    OUTPUT.write_text("\n".join(lines) + "\n", encoding="utf-8")
    print(f"Wrote {OUTPUT}")


if __name__ == "__main__":
    main()