"""(iv) Computed-value validation.

Derived metrics: re-compute from the validated raw inputs and compare to the
served DB value. Three families covered in v1:

1. OSS distribution (F.10): we derive oss_pct_any / multiple / chronic from
   the source's distribution columns (% zero / 1 / 2 / 3 / 4+).
2. PSES topic rollups (O.6): we average subtopic scores per (school, year,
   topic, respondent) and store the mean.
3. Demographic shares aggregated from grade × school rows.

For each, this script re-derives from the cached source file and compares
to philly_school_year_metrics. Reports per-metric agreement + worst diffs.

Output: docs/qa_reports/philly/__computed_value.md (+ .json)
Run: python pipeline_philly/verify/reconcile_computed.py
"""
from __future__ import annotations

import json
import pathlib
import statistics
import sys
import zipfile
import io

import pandas as pd

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2]))

from pipeline_philly.verify import common  # noqa: E402

TOL = 0.5  # ±pp; tighter than (iii) since we're checking math, not cross-publisher
EXACT_FLOOR = 0.01  # delta below which we treat as exact (rounding)


def _toN(v):
    try:
        if v is None or str(v).strip() in {"", "*", "S", "n.a.", "i.s."}:
            return None
        return float(str(v).strip())
    except (ValueError, TypeError):
        return None


def derive_oss_from_source() -> dict[tuple[str, str, str], dict[str, float | None]]:
    """Returns {(ulcs, year, subgroup_enum): {any, multiple, chronic}} re-derived."""
    SOURCE_ID = "sdp_suspensions_school"
    path = common.latest_source_path(SOURCE_ID)
    df = pd.read_csv(path, low_memory=False)
    # Same subgroup map as the loader
    SDP_GROUP_MAP = {
        "All Students": "ALL",
        "American Indian/Alaskan Native": "AMER_INDIAN_AK_NATIVE",
        "Asian": "ASIAN",
        "Native Hawaiian/Pacific Islander": "HAWAIIAN_PAC_ISL",
        "Black/African American": "BLACK",
        "Hispanic/Latino": "HISPANIC",
        "White": "WHITE",
        "Multi Racial/Other": "TWO_OR_MORE_RACES",
        "Economically Disadvantaged": "ECON_DISADV",
        "Not Economically Disadvantaged": "NOT_ECON_DISADV",
        "EL": "ELL",
        "Non-EL": "NOT_ELL",
        "Has IEP": "IEP",
        "Does Not Have IEP": "NOT_IEP",
        "Female": "FEMALE",
        "Male": "MALE",
        "Non-Binary": "NON_BINARY",
    }
    out: dict[tuple[str, str, str], dict[str, float | None]] = {}
    for _, r in df.iterrows():
        ulcs = str(r["ULCS Code"]).strip()
        year = str(r["School Year"]).strip()
        # "2020-2021" → "2020-21"
        if len(year) == 9 and year[4] == "-":
            year = f"{year[:4]}-{year[7:9]}"
        sg = SDP_GROUP_MAP.get(str(r.get("Group", "")).strip())
        if not sg:
            continue
        pZero = _toN(r.get("% with Zero OS Suspensions (Yearly)"))
        pOne = _toN(r.get("% with 1 OS Suspension (Yearly)"))
        pTwo = _toN(r.get("% with 2 OS Suspensions (Yearly)"))
        pThree = _toN(r.get("% with 3 OS Suspensions (Yearly)"))
        four_plus = None
        if all(v is not None for v in (pZero, pOne, pTwo, pThree)):
            four_plus = max(0.0, 100 - (pZero + pOne + pTwo + pThree))
        pAny = (100 - pZero) if pZero is not None else None
        pMultiple = (pTwo + pThree + four_plus) if (pTwo is not None and pThree is not None and four_plus is not None) else None
        pChronic = four_plus
        out[(ulcs, year, sg)] = {"oss_pct_any": pAny, "oss_pct_multiple": pMultiple, "oss_pct_chronic": pChronic}
    return out


def derive_pses_from_source() -> dict[tuple[str, str, str], float | None]:
    """Returns {(ulcs, year, metric_key): mean_score} re-derived from subtopic scores."""
    SOURCE_ID = "pses_topic_subtopic_scores_all_years"
    zip_path = common.latest_source_path(SOURCE_ID)
    TOPIC_MAP = {
        "School Climate": "school_climate",
        "Instructional Environment": "instructional_environment",
        "School Leadership": "school_leadership",
        "Professional Capacity": "professional_capacity",
        "Family Engagement": "family_engagement",
        "Diversity, Equity, and Inclusion": "dei",
    }
    RESP_MAP = {"Teacher": "teacher", "Student": "student"}  # O.6 first-class only
    with zipfile.ZipFile(zip_path) as z:
        with z.open("open_data_subtopic_scores.xlsx") as f:
            content = f.read()
    df = pd.read_excel(io.BytesIO(content), sheet_name="School")
    # Aggregate per (ulcs, year, topic, respondent) → mean of subtopic scores
    sums: dict[tuple[str, str, str, str], list[float]] = {}
    for _, r in df.iterrows():
        ulcs = str(r["ulcs_code"]).strip()
        year = str(r["year_academic"]).strip()
        if len(year) == 9 and year[4] == "-":
            year = f"{year[:4]}-{year[7:9]}"
        topic = TOPIC_MAP.get(str(r.get("topic", "")).strip())
        resp = RESP_MAP.get(str(r.get("survey", "")).strip())
        if not topic or not resp:
            continue
        v = _toN(r.get("score_display"))
        if v is None:
            continue
        sums.setdefault((ulcs, year, topic, resp), []).append(v)
    out: dict[tuple[str, str, str], float | None] = {}
    for (ulcs, year, topic, resp), vals in sums.items():
        if not vals:
            continue
        out[(ulcs, year, f"survey_{resp}_{topic}")] = sum(vals) / len(vals)
    return out


def derive_demographics_from_source() -> dict[tuple[str, str, str], float]:
    """Returns {(ulcs, year, metric_key): aggregated_share} re-derived from grade-level rows."""
    SOURCE_ID = "sdp_enrollment_demographics_school_2024-25"
    path = common.latest_source_path(SOURCE_ID)
    df = pd.read_csv(path, low_memory=False)
    COUNT_TO_METRIC = {
        "ell": "pct_english_learner",
        "iep": "pct_special_ed",
        "female": "pct_female",
        "male": "pct_male",
        "indian": "pct_amer_indian",
        "asian": "pct_asian",
        "hawaiian": "pct_native_hawaiian",
        "black": "pct_black",
        "hispanic": "pct_hispanic",
        "white": "pct_white",
        "mult": "pct_two_or_more_races",
    }
    # Aggregate
    agg: dict[str, dict] = {}
    for _, r in df.iterrows():
        ulcs = str(r.get("ulcscode", "")).strip()
        if not ulcs:
            continue
        all_n = _toN(r.get("allstudents"))
        if all_n is None or all_n <= 0:
            continue
        a = agg.setdefault(ulcs, {"all": 0, "counts": {}, "cep": None})
        a["all"] += all_n
        for col in COUNT_TO_METRIC:
            n = _toN(r.get(col))
            if n is not None:
                a["counts"][col] = a["counts"].get(col, 0) + n
        if a["cep"] is None:
            a["cep"] = _toN(r.get("ceppct"))
    out: dict[tuple[str, str, str], float] = {}
    for ulcs, a in agg.items():
        if a["all"] <= 0:
            continue
        out[(ulcs, "2024-25", "enrollment")] = a["all"]
        for col, mk in COUNT_TO_METRIC.items():
            if col in a["counts"]:
                out[(ulcs, "2024-25", mk)] = a["counts"][col] / a["all"] * 100
        if a["cep"] is not None:
            out[(ulcs, "2024-25", "pct_econ_disadv")] = a["cep"]
    return out


def compare(name: str, expected_by_key: dict[tuple, float | None], db_query: str, key_indexer) -> dict:
    """Compare per-cell. expected_by_key maps tuple → value. db_query returns (key, served)."""
    dsn = common.served_db_dsn()
    import psycopg
    served: dict[tuple, float] = {}
    with psycopg.connect(dsn) as conn:
        conn.read_only = True
        with conn.cursor() as cur:
            cur.execute(db_query)
            for row in cur:
                k = key_indexer(row)
                if k is not None and row[-1] is not None:
                    served[k] = float(row[-1])

    common_keys = set(expected_by_key.keys()) & set(served.keys())
    diffs: list[dict] = []
    matches = 0
    mismatches = 0
    expected_null = 0
    for k in common_keys:
        e = expected_by_key.get(k)
        if e is None:
            expected_null += 1
            continue
        s = served[k]
        d = s - e
        if abs(d) <= TOL:
            matches += 1
        else:
            mismatches += 1
            diffs.append({"key": list(k), "expected": round(e, 3), "served": round(s, 3), "delta": round(d, 3)})

    return {
        "name": name,
        "expected_total": len(expected_by_key),
        "served_total": len(served),
        "common_keys": len(common_keys),
        "expected_null_in_common": expected_null,
        "matches": matches,
        "mismatches": mismatches,
        "expected_only": len(set(expected_by_key.keys()) - set(served.keys())),
        "served_only": len(set(served.keys()) - set(expected_by_key.keys())),
        "agreement_rate": round(matches / (matches + mismatches), 4) if (matches + mismatches) else None,
        "worst_diffs_top10": sorted(diffs, key=lambda r: abs(r["delta"]), reverse=True)[:10],
    }


def main() -> None:
    common.ensure_qa_dir()
    print("[iv] re-deriving OSS distribution from raw...")
    oss = derive_oss_from_source()
    # OSS is structured as one expected dict per metric_key.
    oss_by_metric: dict[str, dict[tuple, float | None]] = {
        "oss_pct_any": {},
        "oss_pct_multiple": {},
        "oss_pct_chronic": {},
    }
    for (ulcs, year, sg), values in oss.items():
        for mk, v in values.items():
            oss_by_metric[mk][(ulcs, year, sg)] = v

    print("[iv] re-deriving PSES topic rollups from subtopic xlsx...")
    pses = derive_pses_from_source()

    print("[iv] re-deriving demographic shares from grade-level rows...")
    demo = derive_demographics_from_source()

    summaries: list[dict] = []

    for metric, expected in oss_by_metric.items():
        s = compare(
            f"OSS — {metric}",
            expected,
            f"""SELECT school_ulcs, year, subgroup::text, value::float8
                  FROM philly_school_year_metrics
                 WHERE metric_key = '{metric}' AND value IS NOT NULL""",
            lambda row: (str(row[0]), str(row[1]), str(row[2])),
        )
        summaries.append(s)
        print(f"  {s['name']:35s} matches={s['matches']:>5d}  mismatches={s['mismatches']:>4d}  agreement={s['agreement_rate']}")

    s = compare(
        "PSES topic rollups (teacher + student × 6 topics)",
        pses,
        """SELECT school_ulcs, year, metric_key, value::float8
             FROM philly_school_year_metrics
            WHERE metric_key LIKE 'survey_teacher_%' OR metric_key LIKE 'survey_student_%'
              AND value IS NOT NULL""",
        lambda row: (str(row[0]), str(row[1]), str(row[2])),
    )
    summaries.append(s)
    print(f"  {s['name']:35s} matches={s['matches']:>5d}  mismatches={s['mismatches']:>4d}  agreement={s['agreement_rate']}")

    s = compare(
        "Demographic shares (aggregated from grade rows)",
        demo,
        """SELECT school_ulcs, year, metric_key, value::float8
             FROM philly_school_year_metrics
            WHERE metric_key IN ('enrollment','pct_english_learner','pct_special_ed','pct_female','pct_male',
                                 'pct_amer_indian','pct_asian','pct_native_hawaiian','pct_black',
                                 'pct_hispanic','pct_white','pct_two_or_more_races','pct_econ_disadv')
              AND subgroup = 'ALL'
              AND value IS NOT NULL""",
        lambda row: (str(row[0]), str(row[1]), str(row[2])),
    )
    summaries.append(s)
    print(f"  {s['name']:35s} matches={s['matches']:>5d}  mismatches={s['mismatches']:>4d}  agreement={s['agreement_rate']}")

    out: list[str] = []
    out.append("# (iv) Computed-value validation — Philly\n")
    out.append("Re-derives the computed metrics from the validated raw source and compares\n")
    out.append(f"to the served DB. Tolerance: ±{TOL} pp.\n")
    for s in summaries:
        out.append(f"\n## {s['name']}\n")
        out.append(f"```json\n{json.dumps(s, indent=2, default=str)}\n```")
    (common.QA_REPORTS_DIR / "__computed_value.md").write_text("\n".join(out))
    # NaN is not valid JSON (and breaks the site's static import of this file);
    # emit null instead.
    def _nan_to_null(o):
        if isinstance(o, float) and o != o:
            return None
        if isinstance(o, dict):
            return {k: _nan_to_null(v) for k, v in o.items()}
        if isinstance(o, list):
            return [_nan_to_null(v) for v in o]
        return o
    (common.QA_REPORTS_DIR / "__computed_value.json").write_text(
        json.dumps(_nan_to_null(summaries), indent=2, default=str, allow_nan=False),
    )
    print(f"\n[iv] report -> docs/qa_reports/philly/__computed_value.md")


if __name__ == "__main__":
    main()