#!/usr/bin/env python3
"""Original descriptive reanalysis. No network, student-level output, or causal claims.

Usage: python analyze.py --data /path/to/main_regressions/final_data.csv --out ./results
See README.md for the pinned source, design, estimands, and limitations.
"""

import argparse
import hashlib
import json
from pathlib import Path

import numpy as np
import pandas as pd

SOURCE_COMMIT = "2f63dae1a01d51453826fe07ef5cf6678e339588"
SOURCE_SHA256 = "430055b4633eaaf87879bb38eb4b1c9425b91e1d7f529029dc1883e33d91df0a"
ARMS = (("control", "Control"), ("vanilla", "GPT Base"), ("augmented", "GPT Tutor"))
SEED = 20260908
REPLICATES = 10_000
PRACTICE_CUTOFF = 0.80
EXAM_CUTOFF = 0.50
SCORE_TOLERANCE = 1e-12


def validate_data(raw):
    required = ["Student ID", "Class", "Year", "Session", "Part2Tot", "Part3Tot",
                "GPTBase", "GPTTutor", "Treatment arm", "Honors"]
    if not set(required).issubset(raw.columns):
        raise ValueError("Missing required source fields")
    if raw[required].isna().any().any():
        raise ValueError("Missing values: inspect and explicitly revise the sample accounting")
    if raw.duplicated(["Student ID", "Session"]).any():
        raise ValueError("Duplicate student-session observations")
    if not raw[["Part2Tot", "Part3Tot"]].apply(lambda s: s.between(0, 1)).all().all():
        raise ValueError("Scores must already be fractions in [0, 1]")
    if not raw["Honors"].isin([0, 1]).all():
        raise ValueError("Unexpected honors coding")
    if not raw["Session"].isin([1, 2, 3, 4]).all():
        raise ValueError("Unexpected session coding")
    if not raw["Treatment arm"].isin([a for a, _ in ARMS]).all():
        raise ValueError("Unexpected treatment coding")
    for arm, _ in ARMS:
        subset = raw[raw["Treatment arm"] == arm]
        expected = {"GPTBase": int(arm == "vanilla"), "GPTTutor": int(arm == "augmented")}
        for field, value in expected.items():
            if not subset[field].eq(value).all():
                raise ValueError("Treatment labels and indicators disagree")
    if raw.groupby("Class")["Treatment arm"].nunique().max() != 1:
        raise ValueError("Treatment is not constant within class")
    if raw.groupby("Student ID")["Class"].nunique().max() != 1:
        raise ValueError("A student occurs in more than one class")


def score_masks(practice, exam, practice_cutoff=PRACTICE_CUTOFF, exam_cutoff=EXAM_CUTOFF):
    # The source includes 0.4999999999999999 for an effectively 50% score.
    # Numerical noise must not turn equality into an exclusive-threshold event.
    p, e = np.asarray(practice), np.asarray(exam)
    high = (p >= practice_cutoff) | np.isclose(p, practice_cutoff, rtol=0, atol=SCORE_TOLERANCE)
    low = (e < exam_cutoff) & ~np.isclose(e, exam_cutoff, rtol=0, atol=SCORE_TOLERANCE)
    return high, high & low


def summarize(frame, practice_cutoff=PRACTICE_CUTOFF, exam_cutoff=EXAM_CUTOFF, weights=None):
    p = frame["Part2Tot"].to_numpy()
    e = frame["Part3Tot"].to_numpy()
    w = np.ones(len(frame)) if weights is None else np.asarray(weights, dtype=float)
    high, joint = score_masks(p, e, practice_cutoff, exam_cutoff)
    denominator = w[high].sum()
    return {
        "pairs": int(len(frame)),
        "students": int(frame["Student ID"].nunique()),
        "classes": int(frame["Class"].nunique()),
        "practice_mean": float(np.average(p, weights=w)),
        "exam_mean": float(np.average(e, weights=w)),
        "high_practice_count": int(high.sum()),
        "joint_count": int(joint.sum()),
        "high_practice_share": float(w[high].sum() / w.sum()),
        "joint_share": float(w[joint].sum() / w.sum()),
        "conditional_low_exam": float(w[joint].sum() / denominator) if denominator else None,
        "high_practice_students": int(frame.loc[high, "Student ID"].nunique()),
        "high_practice_classes": int(frame.loc[high, "Class"].nunique()),
    }


def class_bootstrap(frame, rng, replicates=REPLICATES):
    # Class sufficient statistics preserve every student's repeated observations.
    # Sampling rows independently would falsely inflate the effective sample size.
    rows = []
    for _, g in frame.groupby("Class", sort=True):
        high, joint = score_masks(g.Part2Tot, g.Part3Tot)
        rows.append([len(g), g.Part2Tot.sum(), g.Part3Tot.sum(), high.sum(),
                     joint.sum()])
    stats = np.asarray(rows, dtype=float)
    draws = rng.integers(0, len(stats), size=(replicates, len(stats)))
    totals = stats[draws].sum(axis=1)
    conditional = np.divide(totals[:, 4], totals[:, 3],
                            out=np.full(replicates, np.nan), where=totals[:, 3] != 0)
    return {
        "practice_mean": totals[:, 1] / totals[:, 0],
        "exam_mean": totals[:, 2] / totals[:, 0],
        "high_practice_share": totals[:, 3] / totals[:, 0],
        "joint_share": totals[:, 4] / totals[:, 0],
        "conditional_low_exam": conditional,
    }


def interval(values):
    values = values[np.isfinite(values)]
    if not len(values):
        return None
    return [float(x) for x in np.quantile(values, [0.025, 0.975], method="linear")]


def run(data_path, out):
    source_bytes = data_path.read_bytes()
    if hashlib.sha256(source_bytes).hexdigest() != SOURCE_SHA256:
        raise ValueError("Source checksum differs from the pinned replication data")
    raw = pd.read_csv(data_path)
    validate_data(raw)
    frame = raw[raw.Honors == 0].copy()
    rng = np.random.default_rng(SEED)
    bootstraps = {}
    arms = []
    for arm, label in ARMS:
        g = frame[frame["Treatment arm"] == arm]
        result = {"arm": arm, "label": label, **summarize(g)}
        samples = class_bootstrap(g, rng)
        result["ci95"] = {key: interval(value) for key, value in samples.items()}
        result["valid_conditional_replicates"] = int(np.isfinite(samples["conditional_low_exam"]).sum())
        bootstraps[arm] = samples
        arms.append(result)

    contrasts = []
    for arm, label in ARMS[1:]:
        for metric in ["conditional_low_exam", "joint_share", "high_practice_share"]:
            contrast = bootstraps[arm][metric] - bootstraps["control"][metric]
            point = next(a[metric] for a in arms if a["arm"] == arm) - arms[0][metric]
            contrasts.append({"comparison": f"{label} minus Control", "metric": metric,
                              "estimate": point, "ci95": interval(contrast)})

    grid = []
    for practice in [0.60, 0.70, 0.80, 0.90]:
        for exam in [0.40, 0.50, 0.60]:
            for arm, label in ARMS:
                grid.append({"arm": arm, "label": label, "practice_cutoff": practice,
                             "exam_cutoff": exam,
                             **summarize(frame[frame["Treatment arm"] == arm], practice, exam)})

    subgroups = []
    for field in ["Year", "Session"]:
        for level, group in frame.groupby(field, sort=True):
            for arm, label in ARMS:
                subgroups.append({"split": field, "level": int(level), "arm": arm,
                                  "label": label,
                                  **summarize(group[group["Treatment arm"] == arm])})

    variants = []
    for arm, label in ARMS:
        regular = frame[frame["Treatment arm"] == arm]
        all_students = raw[raw["Treatment arm"] == arm]
        weights = 1 / regular.groupby("Student ID")["Session"].transform("size")
        complete = regular[regular.groupby("Student ID")["Session"].transform("size") == 4]
        for name, data, w in [("Include honors", all_students, None),
                              ("Equal student weight", regular, weights),
                              ("Four-session attendees", complete, None)]:
            variants.append({"variant": name, "arm": arm, "label": label,
                             **summarize(data, weights=w)})
    leave_one_out = []
    for arm, label in ARMS:
        g = frame[frame["Treatment arm"] == arm]
        values = [summarize(g[g.Class != c])["conditional_low_exam"] for c in sorted(g.Class.unique())]
        leave_one_out.append({"arm": arm, "label": label,
                              "minimum": min(values), "maximum": max(values)})

    result = {
        "provenance": {
            "source_repository": "https://github.com/obastani/GenAICanHarmLearning",
            "source_commit": SOURCE_COMMIT,
            "source_file": "main_regressions/final_data.csv",
            "source_sha256": hashlib.sha256(source_bytes).hexdigest(),
            "analysis_date": "2026-09-08",
            "analysis_type": "Exploratory descriptive secondary analysis; not preregistered",
            "numpy": np.__version__, "pandas": pd.__version__,
        },
        "method": {"seed": SEED, "bootstrap_replicates": REPLICATES,
                   "resampling_unit": "Whole classroom, independently within assigned arm",
                   "interval": "2.5th and 97.5th percentiles; NumPy linear quantiles",
                   "practice_cutoff_inclusive": PRACTICE_CUTOFF,
                   "exam_cutoff_exclusive": EXAM_CUTOFF,
                   "score_equality_tolerance": SCORE_TOLERANCE},
        "sample": {"source_pairs": len(raw), "source_students": raw["Student ID"].nunique(),
                   "source_classes": raw.Class.nunique(),
                   "excluded_honors_pairs": int((raw.Honors == 1).sum()),
                   "excluded_honors_students": raw[raw.Honors == 1]["Student ID"].nunique(),
                   "analyzed_pairs": len(frame), "analyzed_students": frame["Student ID"].nunique(),
                   "analyzed_classes": frame.Class.nunique(), "missing_score_pairs": 0,
                   "duplicate_pairs": 0},
        "arms": arms, "descriptive_contrasts": contrasts,
        "threshold_grid": grid, "subgroups": subgroups,
        "sample_sensitivity": variants, "leave_one_class_out": leave_one_out,
    }
    out.mkdir(parents=True, exist_ok=True)
    (out / "results.json").write_text(json.dumps(result, indent=2, allow_nan=False) + "\n")
    summary_rows = []
    for a in arms:
        row = {k: v for k, v in a.items() if k != "ci95"}
        for metric, bounds in a["ci95"].items():
            row[f"{metric}_ci_lower"], row[f"{metric}_ci_upper"] = bounds
        summary_rows.append(row)
    pd.DataFrame(summary_rows).to_csv(out / "summary.csv", index=False)
    pd.DataFrame(grid).to_csv(out / "thresholds.csv", index=False)
    pd.DataFrame(subgroups).to_csv(out / "subgroups.csv", index=False)
    pd.DataFrame(variants).to_csv(out / "sample-sensitivity.csv", index=False)
    print(json.dumps({"sample": result["sample"], "arms": arms,
                      "descriptive_contrasts": contrasts,
                      "leave_one_class_out": leave_one_out}, indent=2))
    return result


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--data", required=True, type=Path)
    parser.add_argument("--out", required=True, type=Path)
    options = parser.parse_args()
    run(options.data, options.out)
