"""T3-synthesis: extraction, harmonisation onto theta = ln(T_AI/T_noAI), FE/DL pooling, forest plot.
Run from the workspace root:  ~/wp29/venv/bin/python out/analysis.py
Reads only data/; writes out/results.json, out/table1.csv, out/figure1.png, out/exclusions.csv.
Deterministic; no randomness (random_seed = null)."""
import hashlib, json, math, os, re
import numpy as np
import pandas as pd
from scipy.stats import norm
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

DATA = "data"; OUT = "out"
os.makedirs(OUT, exist_ok=True)
EXPECTED_SHA = {
    "arxiv-2410.12944v3": "d5270b4435cfbdd516b7268d29c03670509eac7c93604a4466bdb192cbfb803b",
    "arxiv-2507.09089v2": "b6af8a2bc41012ce64c75e7764fa3a999546df95986adb8860f3442ecd0aed2d",
    "arxiv-2411.01067v2": "46a1ee4e200abb52e3d5148ca0e7ee5c1b8cded80363c0cdceb66c249913053a",
    "arxiv-2511.13860v1": "8fb2520f4e9343e47eb39d551c9da49c0d56a7c58ce455fa9073642e1621dfa5",
    "arxiv-2304.09339v1": "1db48b1c1908b00cab3cc1bfe0ec8ef28bc7d85387aea2477ca709c40a0ec06e",
}

studies = json.load(open(os.path.join(DATA, "studies.json")))
src_lines = {}
n_input_rows = {"studies.json": len(studies)}
with open(os.path.join(DATA, "README.md"), "rb") as f:
    n_input_rows["README.md"] = f.read().count(b"\n")
for s in studies:
    p = os.path.join(DATA, "sources", s["source"] + ".txt")
    raw = open(p, "rb").read()
    sha = hashlib.sha256(raw).hexdigest()
    assert sha == EXPECTED_SHA[s["source"]], f"hash mismatch {p}"
    n_input_rows[s["source"] + ".txt"] = raw.count(b"\n")  # equals `wc -l`
    src_lines[s["id"]] = raw.decode("utf-8").split("\n")


def grab(sid, pattern, nth=0):
    """Return (line_no, line) of the nth source line matching regex `pattern`; print it."""
    hits = [(i + 1, l) for i, l in enumerate(src_lines[sid]) if re.search(pattern, l)]
    assert len(hits) > nth, f"{sid}: pattern not found: {pattern}"
    ln, line = hits[nth]
    print(f"[{sid} line {ln}] {line.strip()[:300]}")
    return ln, line


def next_nonblank(sid, ln):
    """Return (line_no, line) of the first non-blank line after 1-based line ln; print it."""
    L = src_lines[sid]
    i = ln
    while not L[i].strip():
        i += 1
    print(f"[{sid} line {i+1}] {L[i].strip()}")
    return i + 1, L[i]


def nums(s):
    return re.findall(r"-?\d+(?:\.\d+)?", s)

print("=== Source lines relied upon ===")
# ---------------- S1 Paradis et al. Table III Model 2 + Sec VI-C ----------------
ln, l = grab("S1", r"ExpCon \| ")
coefs = l.split("ExpCon")[0].split()
s1_b_txt = coefs[1]                     # Model 2 (second of four), printed with '.' marker
_, se_line = next_nonblank("S1", ln)    # following line holds SEs
s1_se_txt = se_line.split()[1]          # "(0.14)"
_, l3 = grab("S1", r"best-fit Model 2")
assert "[-0.51,0.03]" in l3.replace(" ", "") and "p=0.086" in l3.replace(" ", "")
grab("S1", r"TABLE III")
s1_b = float(re.sub(r"[^\d\.\-]", "", s1_b_txt).rstrip("."))
s1_se = float(s1_se_txt.strip("()"))
assert (s1_b, s1_se) == (-0.24, 0.14)
s1_ci = (-0.51, 0.03)
s1_se_from_ci = (s1_ci[1] - s1_ci[0]) / (2 * 1.96)
grab("S1", r"N=47")

# ---------------- S2 Becker et al. headline ----------------
grab("S2", r"increases completion time by 19%")
_, l = grab("S2", r"implement issues in 19% more time")
grab("S2", r"Developers complete 136 issues")
grab("S2", r"S=\\frac")
grab("S2", r"=\\exp\(\\beta\)-1")
s2_c = 0.19

# ---------------- S3 Bono & Xu Table 3 All Content (first = Aggregate) ----------------
_, l = grab("S3", r"^All Content \|", 0)
cells = [c.strip() for c in l.split("|")]
s3_overall, s3_t, s3_c = float(cells[1]), float(cells[2]), float(cells[3])
s3_impr_txt = cells[4]
assert s3_impr_txt == "-29.79%***"
s3_c_pct = float(re.sub(r"[%\*]", "", s3_impr_txt)) / 100.0
grab("S3", r"We had 181 subjects")
grab("S3", r"Editorial note")
grab("S3", r"Table 3: Comparison of task duration")

# ---------------- S4 Bono Table 3 Full Corpus ----------------
_, la = grab("S4", r"^Aware Time per Email")
_, lc = grab("S4", r"^Control Time per Email")
_, ls = grab("S4", r"^% Time Savings")
_, lp = grab("S4", r"^p-value \|")
s4_aware = float(la.split("|")[3]); s4_ctrl = float(lc.split("|")[3])
s4_sav_txt = ls.split("|")[3].strip(); s4_p_txt = lp.split("|")[3].strip()
s4_sav = float(s4_sav_txt) / 100.0; s4_p = float(s4_p_txt)
grab("S4", r"Control \( N\{=\}52")
grab("S4", r"Aware \( N\{=\}56")

# ---------------- S5 Kreitmeir & Raschky Table 1 panel A col 3 ----------------
ln5, l = grab("S5", r"TreatPost \|", 0)
s5_b_txt = l.split("|")[3].strip()
_, se5_line = next_nonblank("S5", ln5)
s5_se_txt = se5_line.split("|")[3].strip()
_, ln_n = grab("S5", r"^N \| 18332")
s5_n_txt = ln_n.split("|")[3].strip()
grab("S5", r"aggregate Output measure")
assert (s5_b_txt, s5_se_txt, s5_n_txt) == ("-0.278", "(0.295)", "18332")
print("=== end source lines ===\n")

# ---------------- Harmonisation ----------------
items = {}
s1_theta = s1_b
items["S1"] = dict(
    estimate_as_printed="ExpCon coefficient -0.24. (Model 2, logged time on task); text: 'effect of AI = 21%' faster, 'about 21%' in abstract",
    uncertainty_as_printed="SE (0.14); 95% CI = [-0.51, 0.03]; p=0.086 (NS); marker '.' p<.1",
    sample_as_printed="n = 96 (Table III); N=47 AI, N=46 no-AI (Sec VI-B)",
    definition="OLS coefficient of the AI-exposure indicator (ExpCon) on log time on task (LogToT), Model 2 with developer-level covariates; one complex enterprise coding task, randomised AI vs no-AI",
    locator="arxiv-2410.12944v3, Table III, row ExpCon (+ SE row below), column Model (2) 'Dev'; Sec VI-C sentence 'estimate of experimental condition on best-fit Model 2'",
    harmonisable=True, harmonised_estimate=s1_theta, harmonised_se=s1_se, poolable=True,
    transformation=f"theta = b = -0.24 (log-time coefficient); SE(theta) = SE(b) = 0.14 (printed). Cross-check from CI: (0.03-(-0.51))/(2*1.96) = {s1_se_from_ci:.4f}",
    reason_if_not=None)
s2_theta = math.log(1 + s2_c)
items["S2"] = dict(
    estimate_as_printed="+19% (allowing AI increases completion time by 19%; 'implement issues in 19% more time on average')",
    uncertainty_as_printed="none printed (95% CI shown only in Figure 1 / Figure 15, not reproduced in the extract)",
    sample_as_printed="16 developers, 246 tasks (136 AI-allowed, 110 AI-disallowed issues)",
    definition="S = E[time | AI allowed]/E[time | AI disallowed] - 1 = exp(beta) - 1, beta from OLS of log implementation time on AI indicator and log forecast time; issue-level randomisation",
    locator="arxiv-2507.09089v2, Abstract and Section 3 opening paragraph ('19% more time'); estimator Sec 2.3 and App. D.1-D.3",
    harmonisable=True, harmonised_estimate=s2_theta, harmonised_se=None, poolable=False,
    transformation=f"theta = ln(1 + c) = ln(1 + 0.19) = {s2_theta:.6f} (equivalently theta = beta, since S = exp(beta) - 1)",
    reason_if_not="Not poolable: no SE, CI or p-value for the 19% estimate is printed in the extract; the confidence interval appears only in Figure 1 (not reproduced), so SE(theta) cannot be obtained.")
s3_theta = math.log(1 + s3_c_pct)
items["S3"] = dict(
    estimate_as_printed="All Content: Overall 53.84, Treatment 44.29, Control 63.10 (minutes), Improvement -29.79%***",
    uncertainty_as_printed="significance stars only (*** = 0<p<=0.01 per the paper's definition); no SE or CI printed",
    sample_as_printed="181 subjects (treatment/control split ~half each; per-arm N not printed in extract)",
    definition="difference in task duration (minutes) between Copilot and control groups, holding accuracy constant (control duration predicted at treatment accuracy, App. B), all scenarios and tasks",
    locator="arxiv-2411.01067v2, Section 4.2, Table 3, block 'Aggregate Results', row 'All Content', columns Treatment / Control / Improvement",
    harmonisable=True, harmonised_estimate=s3_theta, harmonised_se=None, poolable=False,
    transformation=f"theta = ln(1 + c) = ln(1 - 0.2979) = {s3_theta:.6f} (ratio of printed times: ln(44.29/63.10) = {math.log(s3_t/s3_c):.6f})",
    reason_if_not="Not poolable: only significance stars (p<=0.01) are printed; no SE, CI or exact p-value, so SE(theta) cannot be obtained.")
s4_theta = math.log(1 - s4_sav)
s4_z = norm.ppf(1 - s4_p / 2)
s4_se = abs(s4_theta) / s4_z
items["S4"] = dict(
    estimate_as_printed="Full Corpus: Aware Time per Email (s) 50.5, Control Time per Email (s) 56.9, % Time Savings 11.1",
    uncertainty_as_printed="p-value 0.471",
    sample_as_printed="Aware N=56, Control N=52 (167 analysts across three arms after filtering)",
    definition="percentage time saving (Control vs Aware arm) in email-processing time over the full 25-email queue, without the resolve-benign protocol; Aware arm sees the agent's verdicts/rationales and an agent-ordered queue",
    locator="arxiv-2511.13860v1, Section 3.3, Table 3, column 'Full Corpus', rows Aware Time / Control Time / % Time Savings / p-value",
    harmonisable=True, harmonised_estimate=s4_theta, harmonised_se=s4_se, poolable=True,
    transformation=f"time saving 11.1% => c = -0.111; theta = ln(1 - 0.111) = {s4_theta:.6f}; approximate SE = |theta| / z_(1-0.471/2) = {abs(s4_theta):.6f}/{s4_z:.6f} = {s4_se:.6f} (p assumed two-sided). APPROXIMATE SE.",
    reason_if_not=None)
items["S5"] = dict(
    estimate_as_printed="Treat x Post -0.278 (Output, column 3, panel A)",
    uncertainty_as_printed="SE (0.295); no stars",
    sample_as_printed="N 18332 (user-day observations, users with affiliation)",
    definition="difference-in-differences coefficient of Italy x post-ban on the aggregate Output measure (count-based GitHub activity: pushes, PRs, comments, creates, issues), pre 27/28.03 vs post 03/04.04 2023, user FE",
    locator="arxiv-2304.09339v1, Section 4, Table 1, panel A, row 'Treat x Post' (+ SE row), column (3) Output, N row",
    harmonisable=False, harmonised_estimate=None, harmonised_se=None, poolable=False,
    transformation="none (not harmonisable)",
    reason_if_not="Not harmonisable: the effect is on output counts (GitHub events), not on task-completion time for the same tasks, and the contrast is a non-randomised difference-in-differences around the withdrawal of ChatGPT access (ban), not a randomised AI vs no-AI comparison.")

# ---------------- Exclusion ledger ----------------
n_raw = len(items)
harm = [k for k, v in items.items() if v["harmonisable"]]
pool = [k for k, v in items.items() if v["poolable"]]
excl = pd.DataFrame([
    dict(step=1, rule="printed value extracted from cited extract", n_before=n_raw, n_after=n_raw, n_removed=0,
         rationale="all 5 items have a printed value at the stated locator"),
    dict(step=2, rule="harmonisable to theta=ln(T_AI/T_noAI)", n_before=n_raw, n_after=len(harm), n_removed=n_raw - len(harm),
         rationale="S5 removed: output-count DiD around a ban, not a time comparison of AI vs no-AI on the same tasks"),
    dict(step=3, rule="SE(theta) obtainable from printed quantities (poolable)", n_before=len(harm), n_after=len(pool), n_removed=len(harm) - len(pool),
         rationale="S2 (CI only in a figure) and S3 (stars only) removed: no SE obtainable"),
])
excl.to_csv(os.path.join(OUT, "exclusions.csv"), index=False)
assert n_raw - excl.n_removed.sum() == len(pool)

# ---------------- Baseline (P2): unweighted means ----------------
th_h = np.array([items[k]["harmonised_estimate"] for k in harm])
th_p = np.array([items[k]["harmonised_estimate"] for k in pool])
baseline_h = float(th_h.mean()); baseline_p = float(th_p.mean())

# ---------------- Pooling ----------------
y = th_p
se = np.array([items[k]["harmonised_se"] for k in pool])
w = 1 / se**2
fe = float(np.sum(w * y) / np.sum(w)); fe_se = float(math.sqrt(1 / np.sum(w)))
Q = float(np.sum(w * (y - fe) ** 2)); df = len(y) - 1
C = float(np.sum(w) - np.sum(w**2) / np.sum(w))
tau2 = max(0.0, (Q - df) / C)
i2 = max(0.0, (Q - df) / Q) * 100 if Q > 0 else 0.0
wr = 1 / (se**2 + tau2)
re_ = float(np.sum(wr * y) / np.sum(wr)); re_se = float(math.sqrt(1 / np.sum(wr)))
het = "low" if i2 < 25 else ("moderate" if i2 < 50 else "substantial")
re_lo, re_hi = re_ - 1.96 * re_se, re_ + 1.96 * re_se

def fe_pool(yv, sev):
    wv = 1 / np.asarray(sev) ** 2; m = float(np.sum(wv * np.asarray(yv)) / np.sum(wv)); sd = float(math.sqrt(1 / np.sum(wv)))
    return m, m - 1.96 * sd, m + 1.96 * sd
# Sensitivity (not the primary result): S1 SE from printed CI; S4 theta from ratio of printed times
sens_s1ci = fe_pool([s1_theta, s4_theta], [s1_se_from_ci, s4_se])
_t4r = math.log(s4_aware / s4_ctrl)
sens_s4ratio = fe_pool([s1_theta, _t4r], [s1_se, abs(_t4r) / s4_z])

r6 = lambda x: None if x is None else round(float(x), 6)
for v in items.values():
    v["harmonised_estimate"] = r6(v["harmonised_estimate"]); v["harmonised_se"] = r6(v["harmonised_se"])

data_issues = [
    {"issue": "S1: Sec VI-B prints Mean LogToT = 4.6 for the AI group and 4.3 for the no-AI group, contradicting the minute means (96 vs 114) and the stated direction (AI faster).",
     "handling": "Not used; theta taken from the Table III Model 2 coefficient as instructed."},
    {"issue": "S1: Table III caption says n = 96 while Sec VI-B reports N=47 + N=46 = 93.",
     "handling": "Both reported verbatim in sample_as_printed; no effect on theta/SE."},
    {"issue": "S1: SE implied by the printed 95% CI [-0.51, 0.03] is 0.1378, vs printed SE 0.14 (rounding).",
     "handling": "Printed SE 0.14 used (first convention); CI-based SE reported in transformation as a cross-check."},
    {"issue": "S1: Table III model columns collapsed into one line per row (rendering); stray 'p<'/'p=' lines.",
     "handling": "Model 2 read as second value left to right per editorial note; consistent with -0.24 quoted in text. Stray lines ignored (README)."},
    {"issue": "S2: the 95% CI for the headline +19% is shown only in Figure 1/15, not reproduced; the raw ratio estimator gives 34% (footnote 12), a different estimator.",
     "handling": "Regression-based 19% used; SE set to null and study not pooled; 34% not used."},
    {"issue": "S3: printed improvement -29.79% differs slightly from the ratio of printed times (44.29/63.10 - 1 = -29.81%).",
     "handling": "Printed -29.79% used per the percentage-change convention; difference in theta < 0.0003."},
    {"issue": "S3: only significance stars printed, with unusual definitions (*** 0<p<=0.01, ** 0.01<p<=0.05, * 0.05<p<=0.1); 25% of control subjects used other AI tools.",
     "handling": "Not poolable; contamination noted as a limitation (estimate is incremental over other AI use)."},
    {"issue": "S4: Table 3 unit label '(s)' is implausible: per-email malicious/non-malicious values (2.8/1.8, 1.9/2.4) cannot average to the Full Corpus values (50.5/56.9); values look like per-email minutes and a full-queue total in minutes.",
     "handling": "theta is a unit-free log ratio, so unaffected; numbers reported as printed with the printed unit label."},
    {"issue": "S4: printed % time savings 11.1 vs 1 - 50.5/56.9 = 11.25%; p-value sidedness not stated.",
     "handling": "Printed 11.1% used; p treated as two-sided for the approximate SE."},
    {"issue": "S5: footnote star legend reads '** p<0.5' (apparent typo for 0.05); render-time date 'August 24, 2026'; text says April 11th excluded while panel B header lists 11.04.",
     "handling": "Irrelevant to the harmonisation (S5 not harmonisable); recorded only."},
]

results = {
    "task_id": "T3-synthesis",
    "target_definition": "theta = ln(T_AI / T_noAI): natural log of the ratio of mean task-completion time with access to a generative-AI assistant to mean completion time without it, for the same tasks, from a randomised comparison; negative values mean faster with AI.",
    "items": items,
    "k_extracted": sum(1 for v in items.values() if v["estimate_as_printed"]),
    "k_harmonised": sum(1 for v in items.values() if v["harmonised_estimate"] is not None),
    "k_pooled": len(pool),
    "fe_estimate": r6(fe), "fe_se": r6(fe_se), "fe_ci_low": r6(fe - 1.96 * fe_se), "fe_ci_high": r6(fe + 1.96 * fe_se),
    "re_estimate": r6(re_), "re_se": r6(re_se), "re_ci_low": r6(re_lo), "re_ci_high": r6(re_hi),
    "q_stat": r6(Q), "i2": r6(i2), "tau2": r6(tau2),
    "heterogeneity_interpretation": het,
    "pooled_significant_5pct": bool(re_lo > 0 or re_hi < 0),
    "pooled_pct_change_re": r6(100 * (math.exp(re_) - 1)),
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": data_issues,
    "n_input_rows": n_input_rows,
    "n_analysed": n_raw,
    "random_seed": None,
    "baseline_unweighted_mean_theta_harmonised": r6(baseline_h),
    "baseline_unweighted_mean_theta_poolable": r6(baseline_p),
    "pooled_studies": pool,
    "q_df": df,
    "q_pvalue": r6(1 - __import__("scipy.stats").stats.chi2.cdf(Q, df)),
    "s1_se_from_ci": r6(s1_se_from_ci),
    "sensitivity_fe_s1_se_from_ci": [r6(x) for x in sens_s1ci],
    "sensitivity_fe_s4_theta_from_time_ratio": [r6(x) for x in sens_s4ratio],
    "pooled_pct_change_re_ci": [r6(100 * (math.exp(re_lo) - 1)), r6(100 * (math.exp(re_hi) - 1))],
}
with open(os.path.join(OUT, "results.json"), "w") as f:
    json.dump(results, f, indent=2)

cols = ["estimate_as_printed", "uncertainty_as_printed", "sample_as_printed", "definition", "locator",
        "harmonised_estimate", "harmonised_se", "harmonisable", "poolable", "transformation", "reason_if_not"]
pd.DataFrame([{"study": k, **{c: items[k][c] for c in cols}} for k in items]).to_csv(os.path.join(OUT, "table1.csv"), index=False)

# ---------------- Forest plot ----------------
labels = {"S1": "S1 Paradis et al. 2024 (Google devs)", "S2": "S2 Becker et al. 2025 (METR OSS devs)",
          "S3": "S3 Bono & Xu 2024 (IT admins)", "S4": "S4 Bono 2025 (security analysts) [approx. SE]"}
fig, ax = plt.subplots(figsize=(10, 4.5))
ys = list(range(len(harm) + 1, 1, -1))
for yy, k in zip(ys, harm):
    t = items[k]["harmonised_estimate"]; s = items[k]["harmonised_se"]
    if s is not None:
        ax.errorbar(t, yy, xerr=1.96 * s, fmt="s", color="black", capsize=3)
    else:
        ax.plot(t, yy, "o", mfc="white", color="grey")
        ax.annotate("no SE (not pooled)", (t, yy), textcoords="offset points", xytext=(-30, 8), fontsize=7, color="grey")
ax.fill([re_lo, re_, re_hi, re_], [1, 1.25, 1, 0.75], color="tab:blue")
ax.axvline(0, color="k", lw=0.8, ls="--")
ax.set_yticks(ys + [1]); ax.set_yticklabels([labels[k] for k in harm] + [f"DL random effects (k={len(pool)}: {', '.join(pool)})"])
ax.set_xlabel("theta = ln(T_AI / T_noAI)   (<0: faster with AI)")
ax.set_title("Generative-AI access and task-completion time (harmonised)", fontsize=10)
ax.set_xlim(-0.6, 0.3)
plt.tight_layout(); plt.savefig(os.path.join(OUT, "figure1.png"), dpi=150)

print(f"k_extracted={results['k_extracted']} k_harmonised={results['k_harmonised']} k_pooled={results['k_pooled']}")
print(f"FE {fe:.4f} (SE {fe_se:.4f}); RE {re_:.4f} (SE {re_se:.4f}) CI [{re_lo:.4f}, {re_hi:.4f}]; Q={Q:.4f} I2={i2:.2f} tau2={tau2:.6f} ({het})")
print(f"RE % change in time: {100*(math.exp(re_)-1):.2f}%")