"""compute_benchmarks: state-school benchmarks computed from our dataset
(spec §5/§8.6). The disadvantaged average must be weighted by cohort size,
medians must ignore NaN, and only the latest year counts."""

import numpy as np
import pandas as pd

from backend.data_loader import compute_benchmarks

LATEST = 202425


def _df():
    rows = [
        # Six primary schools, latest year. Disadvantaged RWM chosen so the
        # weighted average differs clearly from the unweighted mean:
        # weighted = (40*100 + 60*300) / 400 = 55.0 ; unweighted mean = 50.0
        dict(year=LATEST, attainment_8_score=np.nan, eligible_pupils=100,
             rwm_expected_disadvantaged_pct=40.0, eal_pct=10.0,
             sen_support_pct=10.0, disadvantaged_pct=20.0, fsm_pct=15.0, total_pupils=200),
        dict(year=LATEST, attainment_8_score=np.nan, eligible_pupils=300,
             rwm_expected_disadvantaged_pct=60.0, eal_pct=20.0,
             sen_support_pct=14.0, disadvantaged_pct=24.0, fsm_pct=17.0, total_pupils=280),
        dict(year=LATEST, attainment_8_score=np.nan, eligible_pupils=np.nan,
             rwm_expected_disadvantaged_pct=99.0, eal_pct=30.0,
             sen_support_pct=18.0, disadvantaged_pct=30.0, fsm_pct=19.0, total_pupils=300),
        dict(year=LATEST, attainment_8_score=np.nan, eligible_pupils=50,
             rwm_expected_disadvantaged_pct=np.nan, eal_pct=np.nan,
             sen_support_pct=np.nan, disadvantaged_pct=np.nan, fsm_pct=np.nan, total_pupils=np.nan),
        dict(year=LATEST, attainment_8_score=np.nan, eligible_pupils=40,
             rwm_expected_disadvantaged_pct=np.nan, eal_pct=40.0,
             sen_support_pct=20.0, disadvantaged_pct=40.0, fsm_pct=21.0, total_pupils=350),
        dict(year=LATEST, attainment_8_score=np.nan, eligible_pupils=60,
             rwm_expected_disadvantaged_pct=np.nan, eal_pct=50.0,
             sen_support_pct=22.0, disadvantaged_pct=44.0, fsm_pct=23.0, total_pupils=400),
        # Two secondary schools (attainment_8 non-null)
        dict(year=LATEST, attainment_8_score=45.0, eligible_pupils=180,
             rwm_expected_disadvantaged_pct=np.nan, eal_pct=15.0,
             sen_support_pct=12.0, disadvantaged_pct=22.0, fsm_pct=12.0, total_pupils=1000),
        dict(year=LATEST, attainment_8_score=50.0, eligible_pupils=200,
             rwm_expected_disadvantaged_pct=np.nan, eal_pct=25.0,
             sen_support_pct=16.0, disadvantaged_pct=26.0, fsm_pct=14.0, total_pupils=1200),
        # An older-year primary row that must NOT influence anything
        dict(year=202324, attainment_8_score=np.nan, eligible_pupils=500,
             rwm_expected_disadvantaged_pct=1.0, eal_pct=99.0,
             sen_support_pct=99.0, disadvantaged_pct=99.0, fsm_pct=99.0, total_pupils=9999),
    ]
    return pd.DataFrame(rows)


def test_weighted_disadvantaged_average():
    b = compute_benchmarks(_df())
    # Row 3 has NaN eligible_pupils and must be excluded from the weighting.
    assert b["primary"]["disadvantaged_rwm_expected_pct"] == 55.0


def test_medians_ignore_nan_and_older_years():
    b = compute_benchmarks(_df())
    assert b["year"] == LATEST
    # sen medians over [10,14,18,20,22] = 18 — the only context measure still
    # sourced from the performance df (the rest come from the census mart).
    assert b["primary"]["sen_support_pct"] == 18.0
    # disadvantaged_pct medians over [20,24,30,40,44] = 30
    assert b["primary"]["disadvantaged_pct"] == 30.0


def test_benchmarks_use_census_mart_for_context():
    census = {
        "primary": {"year": LATEST, "fsm_pct": 25.3, "eal_pct": 21.8, "median_pupils": 240},
        "secondary": {"year": LATEST, "fsm_pct": 24.1, "eal_pct": 18.9, "median_pupils": 980},
    }
    b = compute_benchmarks(_df(), census_benchmarks=census)
    assert b["primary"]["fsm_pct"] == 25.3
    assert b["primary"]["eal_pct"] == 21.8
    assert b["secondary"]["eal_pct"] == 18.9
    assert b["secondary"]["median_pupils"] == 980


def test_benchmarks_context_none_when_mart_missing():
    # The performance df has no fsm_pct and its eal/disadvantaged columns are
    # KS2-only — never silently fall back to medianing them for context.
    b = compute_benchmarks(_df(), census_benchmarks=None)
    assert b["primary"]["fsm_pct"] is None
    assert b["primary"]["eal_pct"] is None
    assert b["primary"]["median_pupils"] is None


def test_secondary_block_has_no_disadvantaged_rwm():
    b = compute_benchmarks(_df())
    assert "disadvantaged_rwm_expected_pct" not in b["secondary"]
    # KS2-only columns must not produce a fake secondary disadvantaged anchor
    # (the old median over all-through schools' KS2 rows produced 50%).
    assert b["secondary"]["disadvantaged_pct"] is None


def test_provenance_string():
    b = compute_benchmarks(_df())
    assert b["source"] == "state-school average (computed from our dataset)"


def test_empty_df():
    assert compute_benchmarks(pd.DataFrame()) == {}