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

DATA = "data"; OUT = "out"
SRC = os.path.join(DATA, "sources")
os.makedirs(OUT, exist_ok=True)
Z = 1.96


def lines_of(name):
    with open(os.path.join(SRC, name + ".txt"), encoding="utf-8") as fh:
        return fh.read().split("\n")


def find(name, pattern, nth=0):
    """Return (line_no, line) of the nth line matching regex pattern; print it."""
    L = lines_of(name)
    hits = [(i + 1, l) for i, l in enumerate(L) if re.search(pattern, l)]
    if len(hits) <= nth:
        raise RuntimeError(f"pattern {pattern!r} not found in {name}")
    ln, l = hits[nth]
    print(f"[{name} L{ln}] {l.strip()[:300]}")
    return ln, l


def next_nonempty(name, ln):
    """Line following line number ln (1-based) that is non-empty; print it."""
    L = lines_of(name); j = ln
    while not L[j].strip():
        j += 1
    print(f"[{name} L{j+1}] {L[j].strip()[:300]}")
    return j + 1, L[j]


def wc_l(path):
    with open(path, "rb") as fh:
        return fh.read().count(b"\n")


studies = json.load(open(os.path.join(DATA, "studies.json"), encoding="utf-8"))
assert [s["id"] for s in studies] == ["S1", "S2", "S3", "S4", "S5"]
print("== Source lines relied on ==")
items = {}

# ---------------- S1 Paradis et al. 2024 ----------------
s = "arxiv-2410.12944v3"
ln_b, l_b = find(s, r"ExpCon \|")
ln_se, l_se = next_nonempty(s, ln_b)
coefs = re.findall(r"-?\d+\.\d+[.*]*", l_b.split("ExpCon")[0])
ses = re.findall(r"\((\d+\.\d+)\)", l_se)
b1_str, se1_str = coefs[1], ses[1]          # Model 2 = second column (collapsed table)
b1 = float(re.sub(r"[.*]+$", "", b1_str)); se1 = float(se1_str)
_, l_ci = find(s, r"best-fit Model 2")
ci = re.search(r"\[(-?\d+\.\d+),\s*(-?\d+\.\d+)\]", l_ci); ci_l, ci_u = float(ci.group(1)), float(ci.group(2))
p1 = re.search(r"p=(\d*\.\d+)", l_ci).group(1)
find(s, r"TABLE III: .*n = 96")
find(s, r"completed the task in 96 min")
assert (b1, se1, ci_l, ci_u) == (-0.24, 0.14, -0.51, 0.03)
se1_ci = (ci_u - ci_l) / (2 * Z)
items["S1"] = dict(
    estimate_as_printed=f"{b1_str} (beta_1 for ExpCon, Model 2, Table III; Section VI-C: beta_1=-0.24, i.e. 21% faster)",
    uncertainty_as_printed=f"SE ({se1_str}); 95% CI = [{ci.group(1)}, {ci.group(2)}], p={p1} (NS)",
    sample_as_printed="n = 96 (Table III caption); Section VI-B: N=47 AI, N=46 no-AI",
    definition="OLS coefficient on experimental condition (AI exposure) with natural-log time on task (LogToT) as dependent variable, developer-level covariates (Model 2), one enterprise coding task, randomised",
    locator=f"{s}, Section VI-C Table III, row ExpCon, column (2) 'Dev' (second value on collapsed line L{ln_b}; SE line L{ln_se}); CI and p in Section VI-C text",
    harmonisable=True, harmonised_estimate=b1, harmonised_se=se1, poolable=True,
    transformation=f"log-time coefficient: theta = b = {b1}; SE(theta) = SE(b) = {se1} (printed). Check: SE from CI = (0.03-(-0.51))/(2*1.96) = {se1_ci:.4f}",
    reason_if_not=None)

# ---------------- S2 Becker et al. 2025 (METR) ----------------
s = "arxiv-2507.09089v2"
_, l2 = find(s, r"allowing AI actually increases completion time by 19%")
find(s, r"implement issues in 19% more time on average")
find(s, r"Developers complete 136 issues with AI-allowed and 110")
find(s, r"16 developers with moderate AI experience complete 246 tasks")
c2 = float(re.search(r"by (\d+)%", l2).group(1)) / 100.0
th2 = math.log(1 + c2)
items["S2"] = dict(
    estimate_as_printed="+19% (allowing AI increases completion time by 19%; regression-based, Section 3 / Section D.1)",
    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 issues AI-allowed, 110 AI-disallowed",
    definition="S = E[time AI allowed]/E[time AI disallowed] - 1, estimated as exp(beta)-1 from OLS of log implementation time on AI indicator and log forecast time; issue-level randomisation",
    locator=f"{s}, Abstract and Section 3 'Results' (opening paragraph); estimator Sections 2.3, D.1-D.3",
    harmonisable=True, harmonised_estimate=th2, harmonised_se=None, poolable=False,
    transformation=f"percentage change in time: theta = ln(1 + c) = ln(1 + {c2}) = {th2:.6f}",
    reason_if_not="No SE, CI or p-value printed in the extract; the CI exists only in a figure that is not reproduced, so SE(theta) cannot be obtained (harmonisable, not poolable).")

# ---------------- S3 Bono & Xu 2024 ----------------
s = "arxiv-2411.01067v2"
_, l3 = find(s, r"^All Content \| 53\.84")
find(s, r"We had 181 subjects")
cells = [c.strip() for c in l3.split("|")]
ov3, tr3, co3, imp3 = cells[1], cells[2], cells[3], cells[4]
c3 = float(re.match(r"(-?\d+\.\d+)%", imp3).group(1)) / 100.0
th3 = math.log(1 + c3)
items["S3"] = dict(
    estimate_as_printed=f"Overall {ov3}, Treatment {tr3}, Control {co3} (minutes); Improvement {imp3}",
    uncertainty_as_printed="significance stars only: *** (paper's definition 0<p<=0.01); no SE or CI printed",
    sample_as_printed="181 subjects (randomised half treatment, half control)",
    definition="difference in mean task duration (minutes) between Copilot and control groups across all tasks, holding accuracy constant (control duration predicted at treatment accuracy, Appendix B), randomised",
    locator=f"{s}, Section 4.2 Table 3, 'Aggregate Results' block, row 'All Content', columns Treatment/Control/Improvement",
    harmonisable=True, harmonised_estimate=th3, harmonised_se=None, poolable=False,
    transformation=f"percentage change in time: theta = ln(1 + c) = ln(1 + ({c3})) = {th3:.6f} (cf. ln({tr3}/{co3}) = {math.log(float(tr3)/float(co3)):.6f})",
    reason_if_not="Only significance stars (p<=0.01) are printed; a p-value threshold does not identify SE(theta) (harmonisable, not poolable).")

# ---------------- S4 Bono 2025 ----------------
s = "arxiv-2511.13860v1"
_, l4a = find(s, r"^Aware Time per Email")
_, l4c = find(s, r"^Control Time per Email")
_, l4s = find(s, r"^% Time Savings")
_, l4p = find(s, r"^p-value \|")
find(s, r"Control \( N\{=\}52")
find(s, r"Aware \( N\{=\}56")
last = lambda l: l.split("|")[-1].strip()
a4, ct4, sv4, p4s = last(l4a), last(l4c), last(l4s), last(l4p)
c4 = -float(sv4) / 100.0        # 11.1% time saving -> change in time c = -0.111
p4 = float(p4s)
th4 = math.log(1 + c4)
z4 = norm.ppf(1 - p4 / 2)
se4 = abs(th4) / z4
items["S4"] = dict(
    estimate_as_printed=f"Aware {a4} s, Control {ct4} s per email; % Time Savings {sv4}",
    uncertainty_as_printed=f"p-value {p4s}",
    sample_as_printed="Aware N=56, Control N=52 (167 analysts across three arms incl. Blind N=59)",
    definition="percentage time saving in mean time per email over the full 25-email queue, Aware (agent verdicts/rationale + prioritised queue) vs Control (no AI), randomised arms",
    locator=f"{s}, Section 3.3 Table 3 'Comparison of Email Processing Times', column 'Full Corpus', rows Aware/Control time, % Time Savings, p-value",
    harmonisable=True, harmonised_estimate=th4, harmonised_se=se4, poolable=True,
    transformation=f"percentage change in time c = -{sv4}/100 = {c4}; theta = ln(1 + c) = {th4:.6f}; approximate SE from two-sided p: SE = |theta|/z(1-{p4s}/2) = {abs(th4):.6f}/{z4:.6f} = {se4:.6f} (approximate SE)",
    reason_if_not=None)

# ---------------- S5 Kreitmeir & Raschky 2023 ----------------
s = "arxiv-2304.09339v1"
ln5, l5 = find(s, r"TreatPost \| 0\.115")
_, l5se = next_nonempty(s, ln5)
_, l5n = find(s, r"^N \| 18332")
b5 = [c.strip() for c in l5.split("|")][3]
se5 = [c.strip() for c in l5se.split("|")][3]
n5 = [c.strip() for c in l5n.split("|")][3]
items["S5"] = dict(
    estimate_as_printed=f"{b5} (Treat x Post, Output, column 3, panel A)",
    uncertainty_as_printed=f"SE {se5} (clustered at user level); no significance star",
    sample_as_printed=f"N = {n5} (user-day observations, users with affiliation)",
    definition="difference-in-differences coefficient on aggregate GitHub Output measure (sum of Push, PR, PR comments, commit comments, Create, Issues events) for Italian vs Austrian/French users, pre 27/28.03 vs post 03/04.04.2023 around Italy's ChatGPT ban",
    locator=f"{s}, Section 4 Table 1, panel A, row Treat x Post, column (3) Output (Users with Affiliation)",
    harmonisable=False, harmonised_estimate=None, harmonised_se=None, poolable=False,
    transformation="none applied (not harmonisable)",
    reason_if_not="Effect on output counts (GitHub events) from a DiD around withdrawal of access, not a comparison of task-completion times with vs without AI on the same tasks, and not randomised.")

# ---------------- Pooling ----------------
order = ["S1", "S2", "S3", "S4", "S5"]
pool = [k for k in order if items[k]["poolable"]]
y = np.array([items[k]["harmonised_estimate"] for k in pool]); se = np.array([items[k]["harmonised_se"] for k in pool])


def meta(y, se):
    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) if C > 0 else 0.0
    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)))
    return dict(fe=fe, fe_se=fe_se, q=q, df=df, tau2=tau2, i2=i2, re=re_, re_se=re_se,
                w_fe=(w / w.sum()).tolist(), w_re=(wr / wr.sum()).tolist())


m = meta(y, se)
interp = "low" if m["i2"] < 25 else ("moderate" if m["i2"] < 50 else "substantial")
re_lo, re_hi = m["re"] - Z * m["re_se"], m["re"] + Z * m["re_se"]

# sensitivity: S1 SE from the printed CI instead of the printed SE
se_alt = se.copy(); se_alt[pool.index("S1")] = se1_ci
m_alt = meta(y, se_alt)
# baseline: unweighted means of theta
harm = [k for k in order if items[k]["harmonisable"]]
base_h = float(np.mean([items[k]["harmonised_estimate"] for k in harm]))
base_p = float(np.mean(y))

n_input_rows = {"studies.json": len(studies), "README.md": wc_l(os.path.join(DATA, "README.md"))}
for x in studies:
    n_input_rows[f"sources/{x['source']}.txt"] = wc_l(os.path.join(SRC, x["source"] + ".txt"))

data_issues = [
    {"issue": "S1: Section VI-B reports N=47 (AI) + N=46 (no-AI) = 93 developers, while Table III states n = 96.",
     "handling": "Not resolvable from the extract; reported. sample_as_printed lists both; does not affect theta or SE, which are taken from Table III/Section VI-C."},
    {"issue": "S1: Section VI-B gives Mean LogToT=4.6 for the AI group and 4.3 for no-AI while stating AI was faster (96 vs 114 min); the log means appear swapped/inconsistent.",
     "handling": "These log means are not used."},
    {"issue": "S1: Table III is rendered with the four model columns collapsed into one line; SE (0.14) is rounded, and the CI [-0.51, 0.03] implies SE 0.1378.",
     "handling": "Model 2 taken as the second value per the editorial note; printed SE 0.14 used; CI-based SE reported as sensitivity (sensitivity_s1_se_from_ci)."},
    {"issue": "S2: The uncertainty of the headline +19% estimate is shown only in Figure 1/Figure 15, which are not reproduced; no SE or p-value printed.",
     "handling": "Harmonised (theta = ln 1.19) but not poolable; harmonised_se = null."},
    {"issue": "S3: Treatment/Control means 44.29/63.10 imply -29.81%, printed improvement is -29.79% (slightly beyond rounding of the printed means); significance shown only by stars with an unusual convention (* = 0.05<p<=0.1).",
     "handling": "Printed -29.79% used as instructed (difference in theta ~0.0003); not poolable because stars do not identify an SE."},
    {"issue": "S4: Table 3 units ambiguous: Malicious/Non-Malicious times per email (1.8-2.8 's') cannot average to Full Corpus times (50.5/56.9 s) under any common unit; Full Corpus ratio 50.5/56.9 implies 11.25% saving vs printed 11.1% (within rounding).",
     "handling": "Only the Full Corpus column is used; theta depends only on the printed % saving (unit-free). Flagged as a source inconsistency."},
    {"issue": "S4: SE is approximate, back-calculated from a two-sided p-value (0.471) under a normal approximation on the log scale, although the paper's test is presumably on the time difference.",
     "handling": "Marked as approximate SE in transformation; pooled results carry this caveat."},
    {"issue": "S5: Table 1 note prints '** p<0.5' (presumably p<0.05); row label rendered as '\\times TreatPost'; panel B header lists 11.04 although the text says it was excluded; rendered date line 'August 24, 2026' is an artefact.",
     "handling": "No effect on the extracted panel A column 3 value; S5 is not harmonisable anyway."},
]

res = {
    "task_id": "T3-synthesis",
    "target_definition": "theta = ln(T_AI / T_noAI): natural log of the ratio of mean task-completion time with generative-AI assistance to mean completion time without it, for the same tasks, from a randomised comparison (negative = faster with AI).",
    "items": items,
    "k_extracted": sum(1 for k in order if items[k]["estimate_as_printed"]),
    "k_harmonised": sum(1 for k in order if items[k]["harmonised_estimate"] is not None),
    "k_pooled": len(pool),
    "pooled_studies": pool,
    "fe_estimate": m["fe"], "fe_se": m["fe_se"], "fe_ci_low": m["fe"] - Z * m["fe_se"], "fe_ci_high": m["fe"] + Z * m["fe_se"],
    "re_estimate": m["re"], "re_se": m["re_se"], "re_ci_low": re_lo, "re_ci_high": re_hi,
    "q_stat": m["q"], "q_df": m["df"], "q_pvalue": float(chi2.sf(m["q"], m["df"])),
    "i2": m["i2"], "tau2": m["tau2"],
    "heterogeneity_interpretation": interp,
    "pooled_significant_5pct": bool(re_lo > 0 or re_hi < 0),
    "pooled_pct_change_re": 100 * (math.exp(m["re"]) - 1),
    "pooled_pct_change_re_ci": [100 * (math.exp(re_lo) - 1), 100 * (math.exp(re_hi) - 1)],
    "weights_fe": dict(zip(pool, m["w_fe"])), "weights_re": dict(zip(pool, m["w_re"])),
    "baseline_unweighted_mean_theta_harmonised": base_h,
    "baseline_unweighted_mean_theta_poolable": base_p,
    "sensitivity_s1_se_from_ci": {"s1_se": se1_ci, "re_estimate": m_alt["re"], "re_se": m_alt["re_se"],
                                   "re_ci_low": m_alt["re"] - Z * m_alt["re_se"], "re_ci_high": m_alt["re"] + Z * m_alt["re_se"],
                                   "q_stat": m_alt["q"], "i2": m_alt["i2"], "tau2": m_alt["tau2"]},
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": data_issues,
    "n_input_rows": n_input_rows,
    "n_analysed": len(order),
    "random_seed": None,
}
with open(os.path.join(OUT, "results.json"), "w") as fh:
    json.dump(res, fh, indent=1)

cols = ["study", "estimate_as_printed", "uncertainty_as_printed", "sample_as_printed", "definition", "locator",
        "harmonised_estimate", "harmonised_se", "harmonisable", "poolable", "transformation", "reason_if_not"]
with open(os.path.join(OUT, "table1.csv"), "w", newline="") as fh:
    wr = csv.writer(fh); wr.writerow(cols)
    for k in order:
        it = items[k]
        wr.writerow([k] + [("" if it[c] is None else it[c]) for c in cols[1:]])

# ---------------- Forest plot ----------------
labels = {"S1": "S1 Paradis et al. 2024 (Google)", "S2": "S2 Becker et al. 2025 (METR)",
          "S3": "S3 Bono & Xu 2024 (Security Copilot)", "S4": "S4 Bono 2025 (phishing agent; approx. SE)"}
fig, ax = plt.subplots(figsize=(9, 4.2))
ypos = list(range(len(harm), 0, -1))
for yy, k in zip(ypos, harm):
    th = items[k]["harmonised_estimate"]; s_ = items[k]["harmonised_se"]
    if s_ is not None:
        ax.errorbar(th, yy, xerr=Z * s_, fmt="s", color="black", capsize=3)
        txt = f"{th:.3f} [{th - Z*s_:.3f}, {th + Z*s_:.3f}]"
    else:
        ax.plot(th, yy, "o", mfc="white", mec="gray")
        txt = f"{th:.3f} (no SE; not pooled)"
    ax.text(1.02, yy, txt, transform=ax.get_yaxis_transform(), va="center", fontsize=8)
d = np.array([[re_lo, 0], [m["re"], 0.2], [re_hi, 0], [m["re"], -0.2]])
ax.fill(d[:, 0], d[:, 1], color="tab:blue")
ax.text(1.02, 0, f"{m['re']:.3f} [{re_lo:.3f}, {re_hi:.3f}]", transform=ax.get_yaxis_transform(), va="center", fontsize=8)
ax.axvline(0, color="gray", lw=0.8, ls="--")
ax.set_yticks(ypos + [0]); ax.set_yticklabels([labels[k] for k in harm] + [f"RE pooled (DL), k={len(pool)}, I2={m['i2']:.0f}%"])
ax.set_xlabel("theta = ln(T_AI / T_noAI)  (negative = faster with AI)")
ax.set_title("Harmonised effects of generative-AI access on task time")
plt.tight_layout(rect=(0, 0, 0.8, 1))
fig.savefig(os.path.join(OUT, "figure1.png"), dpi=150)

print("== Results ==")
for k in order:
    print(k, items[k]["harmonised_estimate"], items[k]["harmonised_se"], "poolable" if items[k]["poolable"] else "not pooled")
print(f"FE {m['fe']:.4f} (SE {m['fe_se']:.4f}); RE {m['re']:.4f} [{re_lo:.4f}, {re_hi:.4f}]; Q {m['q']:.4f}; I2 {m['i2']:.2f}; tau2 {m['tau2']:.4f}; pct {res['pooled_pct_change_re']:.2f}%")