"""Build the frozen analysis dataset from the state store (exports + data-quality assertions + negative controls)."""
import os, sys, json, hashlib, datetime
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")); sys.path.insert(0, os.path.dirname(__file__))
from pgar.state import StateStore
from pgar.exports import export_runs
from common import REPO, DATA, load_runs
import pandas as pd

def main():
    store = StateStore(os.path.join(REPO, "runs", "state.sqlite"))
    export_runs(store, REPO)
    df = pd.read_csv(os.path.join(DATA, "runs.csv"))
    # data-quality assertions
    problems = []
    for split in df["split"].unique():
        d = df[df["split"] == split]
        counts = d.groupby(["task_id", "treatment"]).size().unstack(fill_value=0)
        if split in ("pilot", "confirmatory") and (counts.nunique(axis=1) > 1).any():
            problems.append(f"{split}: unbalanced arms per task\n{counts[counts.nunique(axis=1) > 1]}")
        dup = d.duplicated(subset=["task_id", "treatment", "repeat_index", "notes"]).sum()   # different assignments (seeds/faults) of the same cell are distinct experiments
        if dup: problems.append(f"{split}: {dup} duplicate (task, arm, repeat) cells")
    ev = df[df["status"] == "evaluated"]
    if (ev["scientific_success"].isna()).any(): problems.append("evaluated runs without scientific_success")
    if (ev["spent_usd"].fillna(0) < 0).any(): problems.append("negative cost")
    # negative control: Y must be 0 whenever n_violations>0 or any criterion failed
    bad = ev[(ev["scientific_success"] == 1) & ((ev["n_violations"] > 0) | (ev["n_criteria_passed"] < ev["n_criteria"]))]
    if len(bad): problems.append(f"{len(bad)} runs with success=1 despite violations/failed criteria: {bad['run_id'].tolist()}")
    summary = {"built_utc": datetime.datetime.now(datetime.timezone.utc).isoformat(), "n_runs": int(len(df)), "by_split_status": df.groupby(["split", "status"]).size().to_dict(),
               "problems": problems, "runs_csv_sha256": hashlib.sha256(open(os.path.join(DATA, "runs.csv"), "rb").read()).hexdigest()}
    summary["by_split_status"] = {f"{k[0]}/{k[1]}": int(v) for k, v in summary["by_split_status"].items()}
    json.dump(summary, open(os.path.join(DATA, "build_summary.json"), "w"), indent=1)
    print(json.dumps(summary, indent=1))
    if problems:
        print("DATA-QUALITY PROBLEMS:", *problems, sep="\n", file=sys.stderr)

if __name__ == "__main__":
    main()