"""T3-synthesis: extraction, harmonisation onto theta = ln(T_AI/T_noAI), IV fixed-effect and
DerSimonian-Laird pooling, forest plot. Reads only data/; writes out/results.json, out/table1.csv,
out/figure1.png. Deterministic; no randomness (random_seed = null).
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
"""
import csv
import json
import math
import os
import re

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from scipy.stats import norm

DATA = "data"
SRC = os.path.join(DATA, "sources")
OUT = "out"
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 get_line(src, lineno, must_contain):
    """Return 1-based line `lineno` of a source, asserting it contains the expected text; print it."""
    ln = lines_of(src)[lineno - 1]
    for s in must_contain:
        assert s in ln, f"{src}:{lineno} does not contain {s!r}: {ln!r}"
    print(f"[{src} L{lineno}] {ln.strip()}")
    return ln


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


print("=== Source lines relied on ===")
# ---------------- S1 Paradis et al. 2024 ----------------
s1_coef = get_line("arxiv-2410.12944v3", 76, ["ExpCon"])
s1_se = get_line("arxiv-2410.12944v3", 78, ["(0.14)"])
s1_txt = get_line("arxiv-2410.12944v3", 58, ["95\\%CI = [-0.51,0.03]", "p=0.086"])
s1_n = get_line("arxiv-2410.12944v3", 66, ["n = 96"])
s1_n2 = get_line("arxiv-2410.12944v3", 48, ["N=47", "N=46"])
coefs = s1_coef.split("ExpCon")[0].split()          # ['-0.25.', '-0.24.', '-0.30*', '-0.14']
ses = re.findall(r"\(([-\d.]+)\)", s1_se)
b1 = float(coefs[1].rstrip(".*"))                    # Model 2
se1 = float(ses[1])
m = re.search(r"95\\%CI = \[([-\d.]+),([-\d.]+)\]", s1_txt)
L1, U1 = float(m.group(1)), float(m.group(2))
p1 = float(re.search(r"p=([\d.]+)", s1_txt).group(1))
se1_from_ci = (U1 - L1) / (2 * Z)

# ---------------- S2 Becker et al. 2025 ----------------
s2_abs = get_line("arxiv-2507.09089v2", 34, ["increases completion time by 19%"])
s2_res = get_line("arxiv-2507.09089v2", 62, ["19% more time", "136 issues", "110 issues"])
s2_hdr = get_line("arxiv-2507.09089v2", 8, ["Figure 1, which displays the confidence interval"])
c2 = float(re.search(r"by (\d+)%", s2_abs).group(1)) / 100.0

# ---------------- S3 Bono & Xu 2024 ----------------
s3_row = get_line("arxiv-2411.01067v2", 55, ["All Content", "-29.79%***"])
s3_hdr = get_line("arxiv-2411.01067v2", 51, ["Treatment", "Control", "Improvement"])
s3_n = get_line("arxiv-2411.01067v2", 38, ["181 subjects"])
s3_note = get_line("arxiv-2411.01067v2", 10, ["no standard errors or confidence intervals"])
cells3 = [c.strip() for c in s3_row.split("|")]
overall3, t3, ctl3 = float(cells3[1]), float(cells3[2]), float(cells3[3])
c3 = float(cells3[4].replace("%", "").replace("*", "")) / 100.0

# ---------------- S4 Bono 2025 ----------------
s4_hdr = get_line("arxiv-2511.13860v1", 88, ["Full Corpus"])
s4_aw = get_line("arxiv-2511.13860v1", 90, ["Aware Time per Email (s)"])
s4_ct = get_line("arxiv-2511.13860v1", 92, ["Control Time per Email (s)"])
s4_pct = get_line("arxiv-2511.13860v1", 94, ["% Time Savings"])
s4_p = get_line("arxiv-2511.13860v1", 96, ["p-value"])
s4_nc = get_line("arxiv-2511.13860v1", 31, ["Control ( N{=}52 )"])
s4_na = get_line("arxiv-2511.13860v1", 35, ["Aware ( N{=}56 )"])
s4_f = get_line("arxiv-2511.13860v1", 44, ["We filter out 13 subject responses"])
last = lambda ln: ln.split("|")[-1].strip()
aw4, ct4, sav4, p4 = float(last(s4_aw)), float(last(s4_ct)), float(last(s4_pct)), float(last(s4_p))
c4 = -sav4 / 100.0          # "11.1 % time savings" -> change in time c = -0.111

# ---------------- S5 Kreitmeir & Raschky 2023 ----------------
s5_hd = get_line("arxiv-2304.09339v1", 74, ["Events", "Release", "Output", "PushPull"])
s5_c = get_line("arxiv-2304.09339v1", 80, ["TreatPost"])
s5_se = get_line("arxiv-2304.09339v1", 82, ["(0.295)"])
s5_n = get_line("arxiv-2304.09339v1", 88, ["N | 18332"])
s5_def = get_line("arxiv-2304.09339v1", 42, ["aggregate Output measure"])
b5 = [c.strip() for c in s5_c.split("|")][3]
se5 = [c.strip() for c in s5_se.split("|")][3]
n5 = [c.strip() for c in s5_n.split("|")][3]
assert b5 == "-0.278" and se5 == "(0.295)" and n5 == "18332"

# ---------------- Harmonisation ----------------
theta1 = b1
theta2 = math.log(1 + c2)
theta3 = math.log(1 + c3)
theta4 = math.log(1 + c4)
z4 = norm.ppf(1 - p4 / 2)
se4 = abs(theta4) / z4

items = {
    "S1": {
        "estimate_as_printed": f"ExpCon coefficient (Model 2) = {coefs[1]} (marker '.' = p<.1) on LogToT; text: beta_1=-0.24, effect of AI = 21% (faster)",
        "uncertainty_as_printed": f"SE ({ses[1]}); 95% CI = [{m.group(1)}, {m.group(2)}]; p=0.086 (NS)",
        "sample_as_printed": "n = 96 (Table III); time-on-task means N=47 (AI) and N=46 (no AI) in Sec. VI-B",
        "definition": "OLS coefficient of the AI experimental-condition indicator on the natural log of time on task, controlling for developer-level covariates (Model 2); randomised, one common enterprise coding task",
        "locator": "arxiv-2410.12944v3, Table III, row ExpCon (+ SE row below), column (2) 'Dev' [extract lines 76-78]; Section VI-C, 95% CI and p-value [line 58]",
        "harmonisable": True,
        "harmonised_estimate": theta1,
        "harmonised_se": se1,
        "poolable": True,
        "transformation": f"theta = b = {b1}; SE(theta) = SE(b) = {se1} (cross-check from printed CI: ({U1} - ({L1}))/(2*1.96) = {se1_from_ci:.4f})",
        "reason_if_not": None,
    },
    "S2": {
        "estimate_as_printed": "+19% (allowing AI increases completion time by 19%; 'implement issues in 19% more time on average')",
        "uncertainty_as_printed": "none printed (the 95% CI is shown only in Figure 1, which is not reproduced in the extract)",
        "sample_as_printed": "16 developers; 246 tasks (136 issues AI-allowed, 110 AI-disallowed)",
        "definition": "S = E[time with AI allowed]/E[time with AI disallowed] - 1, estimated as exp(beta)-1 from a log-linear OLS of total implementation time on the AI indicator and log forecast time (issue-level randomisation)",
        "locator": "arxiv-2507.09089v2, Abstract [line 34] and Section 3 Results opening [line 62]; estimator Sections 2.3 and D.1-D.3",
        "harmonisable": True,
        "harmonised_estimate": theta2,
        "harmonised_se": None,
        "poolable": False,
        "transformation": f"theta = ln(1 + c) = ln(1 + {c2}) = {theta2:.6f}",
        "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, which is not reproduced, so SE(theta) cannot be obtained.",
    },
    "S3": {
        "estimate_as_printed": f"All Content: Overall {cells3[1]} min, Treatment {cells3[2]} min, Control {cells3[3]} min, Improvement {cells3[4]}",
        "uncertainty_as_printed": "significance stars only: *** (defined as 0<p<=0.01); no SE or CI printed",
        "sample_as_printed": "181 subjects completed the task satisfactorily (randomised half to Copilot)",
        "definition": "Percentage difference in mean task duration (minutes), Copilot vs control, holding accuracy constant (control duration predicted at treatment accuracy; Appendix B), all tasks across three IT-admin scenarios",
        "locator": "arxiv-2411.01067v2, Section 4.2, Table 3, row 'Aggregate Results - All Content', columns Treatment/Control/Improvement [line 55]",
        "harmonisable": True,
        "harmonised_estimate": theta3,
        "harmonised_se": None,
        "poolable": False,
        "transformation": f"theta = ln(1 + c) = ln(1 + ({c3})) = {theta3:.6f} (printed ratio check: ln({t3}/{ctl3}) = {math.log(t3/ctl3):.6f})",
        "reason_if_not": "Not poolable: uncertainty is reported only as significance stars (*** = p<=0.01); bootstrap SEs/CIs are not printed, so SE(theta) cannot be obtained.",
    },
    "S4": {
        "estimate_as_printed": f"Full Corpus: Aware Time per Email (s) {last(s4_aw)}, Control Time per Email (s) {last(s4_ct)}, % Time Savings {last(s4_pct)}",
        "uncertainty_as_printed": f"p-value {last(s4_p)}",
        "sample_as_printed": "Control N=52, Aware N=56 (Blind N=59; 167 analysts in total; 13 responses scoring <40% filtered: 3 Aware, 6 Blind, 4 Control)",
        "definition": "Percentage time saving per processed email (mean time per email), Aware arm (agent queue ordering + verdicts/rationales) vs Control (no AI), full 25-email queue",
        "locator": "arxiv-2511.13860v1, Section 3.3, Table 3, column 'Full Corpus', rows Aware/Control time, % Time Savings, p-value [lines 88-96]",
        "harmonisable": True,
        "harmonised_estimate": theta4,
        "harmonised_se": se4,
        "poolable": True,
        "transformation": f"c = -{sav4}/100 = {c4}; theta = ln(1 + c) = ln({1 + c4:.3f}) = {theta4:.6f}; approximate SE = |theta| / z_(1-p/2) = {abs(theta4):.6f} / z_(1-{p4}/2) = {abs(theta4):.6f} / {z4:.6f} = {se4:.6f} (approximate SE from p-value only)",
        "reason_if_not": None,
    },
    "S5": {
        "estimate_as_printed": f"Treat x Post (Output, column 3, panel A) = {b5}",
        "uncertainty_as_printed": f"SE {se5} (clustered at user level); not significant",
        "sample_as_printed": f"N = {n5} (user-day observations, users with affiliation)",
        "definition": "Difference-in-differences coefficient on an aggregate GitHub output count measure (Push, PRs, comments, Create, Issues) for Italian vs Austrian/French users around Italy's ChatGPT ban (pre 27/28.03 vs post 03/04.04.2023)",
        "locator": "arxiv-2304.09339v1, Section 4, Table 1, panel A, row Treat x Post (+ SE row), column (3) Output; N row [lines 80-88]",
        "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 (activity volume) around the withdrawal of access, not a comparison of completion times with vs without AI on the same tasks, and it is not a randomised comparison.",
    },
}

# ---------------- Pooling ----------------
pool = [(k, v["harmonised_estimate"], v["harmonised_se"]) for k, v in items.items() if v["poolable"]]
y = [p[1] for p in pool]
w = [1 / p[2] ** 2 for p in pool]
sw = sum(w)
fe = sum(wi * yi for wi, yi in zip(w, y)) / sw
fe_se = math.sqrt(1 / sw)
Q = sum(wi * (yi - fe) ** 2 for wi, yi in zip(w, y))
df = len(pool) - 1
C = sw - sum(wi ** 2 for wi in w) / sw
tau2 = max(0.0, (Q - df) / C)
i2 = max(0.0, (Q - df) / Q) * 100 if Q > 0 else 0.0
wr = [1 / (p[2] ** 2 + tau2) for p in pool]
re_est = sum(wi * yi for wi, yi in zip(wr, y)) / sum(wr)
re_se = math.sqrt(1 / sum(wr))
re_lo, re_hi = re_est - Z * re_se, re_est + Z * re_se
het = "low" if i2 < 25 else ("moderate" if i2 < 50 else "substantial")

harm = {k: v["harmonised_estimate"] for k, v in items.items() if v["harmonised_estimate"] is not None}

n_input_rows = {
    "README.md": wc_l(os.path.join(DATA, "README.md")),
    "studies.json": len(json.load(open(os.path.join(DATA, "studies.json")))),
}
for f in sorted(os.listdir(SRC)):
    n_input_rows[f] = wc_l(os.path.join(SRC, f))

data_issues = [
    {"issue": "S1 (Paradis): Sec. VI-B prints Mean LogToT = 4.6 for the AI group and 4.3 for the no-AI group, which contradicts the reported faster AI group (96 vs 114 min).",
     "handling": "Not used; the extracted quantity is the Table III Model 2 coefficient, which is internally consistent with the Sec. VI-C CI and p-value."},
    {"issue": "S1: Table III states n = 96 while Sec. VI-B reports N = 47 + 46 = 93 for time on task.",
     "handling": "Both reported verbatim in sample_as_printed; does not affect theta or SE."},
    {"issue": "S1: printed SE 0.14 vs SE implied by the printed 95% CI [-0.51, 0.03] = 0.1378 (rounding).",
     "handling": "Printed SE(b)=0.14 used per convention; CI-implied SE recorded in transformation as cross-check."},
    {"issue": "S1: Table III rendering collapses four model columns into one line with the variable name at the end and garbles the 'Adj. R^2' / 'p value' row labels.",
     "handling": "Columns read left to right as Models 1-4 per the editorial note; Model 2 = -0.24 confirmed by text (Sec. VI-C)."},
    {"issue": "S2 (Becker): the confidence interval of the headline 19% estimate is only displayed in Figure 1, not printed in the extract.",
     "handling": "Harmonised (theta = ln 1.19) but not poolable; no remembered value substituted."},
    {"issue": "S2: a raw ratio estimate of 34% is also printed (footnote 12), differing from the 19% regression estimate.",
     "handling": "The regression-based 19% headline estimate requested in studies.json is used."},
    {"issue": "S3 (Bono & Xu): printed Improvement -29.79% vs ratio of printed means 44.29/63.10 - 1 = -29.81%.",
     "handling": "Printed percentage used (difference 0.0003 on theta scale, rounding)."},
    {"issue": "S3: only significance stars are printed, with a nonstandard definition (* 0.05<p<=0.1, ** 0.01<p<=0.05, *** p<=0.01).",
     "handling": "SE not recoverable; study marked not poolable."},
    {"issue": "S4 (Bono 2025): printed % Time Savings 11.1 vs 1 - 50.5/56.9 = 11.25% from printed means.",
     "handling": "Printed percentage and p-value used; sensitivity with the ratio of means reported in results (sensitivity_s4_ratio_of_means)."},
    {"issue": "S4: Table 3 Malicious/Non-Malicious 'time per email (s)' values (1.8-2.8) are implausibly small vs the Full Corpus (~50 s), and the Malicious % (-52.9) does not match 2.8/1.9; units ambiguous in those columns.",
     "handling": "Only the Full Corpus column is used; anomaly documented."},
    {"issue": "S4: arm sizes (52/56/59, summing to 167) appear to be before the filtering of 13 responses; analysed N per arm not printed.",
     "handling": "Reported verbatim in sample_as_printed; not needed for theta or the p-value-based SE."},
    {"issue": "S4: SE is approximate, recovered from a two-sided p-value (0.471) assuming normality.",
     "handling": "Marked as approximate; S4 enters pooling per the stated convention."},
    {"issue": "S5 (Kreitmeir & Raschky): table notes print '** p<0.5' (presumably p<0.05); row label rendered as '\\times TreatPost'; date line 'August 24, 2026' is a rendering artefact.",
     "handling": "No effect: S5 is not harmonisable."},
]

# sensitivity (not part of required keys): S4 theta from ratio of printed means, same z
theta4_alt = math.log(aw4 / ct4)
se4_alt = abs(theta4_alt) / z4
w_alt = [1 / se1 ** 2, 1 / se4_alt ** 2]
fe_alt = (w_alt[0] * theta1 + w_alt[1] * theta4_alt) / sum(w_alt)
# sensitivity: S1 SE from CI
w_ci = [1 / se1_from_ci ** 2, 1 / se4 ** 2]
fe_ci = (w_ci[0] * theta1 + w_ci[1] * theta4) / sum(w_ci)

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 the generative-AI assistant to mean completion time without it, on the same tasks, from a randomised comparison; negative = faster with AI.",
    "items": items,
    "k_extracted": sum(1 for v in items.values() if v["estimate_as_printed"]),
    "k_harmonised": len(harm),
    "k_pooled": len(pool),
    "pooled_studies": [p[0] for p in pool],
    "fe_estimate": fe, "fe_se": fe_se, "fe_ci_low": fe - Z * fe_se, "fe_ci_high": fe + Z * fe_se,
    "re_estimate": re_est, "re_se": re_se, "re_ci_low": re_lo, "re_ci_high": re_hi,
    "q_stat": Q, "q_df": df, "q_pvalue": float(1 - __import__("scipy.stats", fromlist=["chi2"]).chi2.cdf(Q, df)),
    "i2": i2, "tau2": tau2,
    "heterogeneity_interpretation": het,
    "pooled_significant_5pct": bool(re_lo > 0 or re_hi < 0),
    "pooled_pct_change_re": 100 * (math.exp(re_est) - 1),
    "baseline_unweighted_mean_poolable": sum(y) / len(y),
    "baseline_unweighted_mean_harmonised": sum(harm.values()) / len(harm),
    "sensitivity_s4_ratio_of_means": {"theta_S4": theta4_alt, "se_S4": se4_alt, "fe_estimate": fe_alt},
    "sensitivity_s1_se_from_ci": {"se_S1": se1_from_ci, "fe_estimate": fe_ci},
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": data_issues,
    "n_input_rows": n_input_rows,
    "n_analysed": len(items),
    "random_seed": None,
}

with open(os.path.join(OUT, "results.json"), "w") as fh:
    json.dump(results, 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, v in items.items():
        wr_.writerow([k] + [("" if v[c] is None else v[c]) for c in cols[1:]])

# ---------------- Forest plot ----------------
short = {"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)"}
rows = [k for k in items if items[k]["harmonisable"]]
fig, ax = plt.subplots(figsize=(9.5, 4.2))
ypos = list(range(len(rows) + 1, 1, -1))
for yp, k in zip(ypos, rows):
    t, s = items[k]["harmonised_estimate"], items[k]["harmonised_se"]
    if s is not None:
        ax.errorbar(t, yp, xerr=Z * s, fmt="s", color="black", capsize=3)
        lab = f"{t:+.3f} [{t - Z*s:+.3f}, {t + Z*s:+.3f}]" + (" (approx. SE)" if k == "S4" else "")
    else:
        ax.plot(t, yp, "o", mfc="white", color="grey")
        lab = f"{t:+.3f} (no SE; not pooled)"
    ax.text(1.02, yp, lab, transform=ax.get_yaxis_transform(), va="center", fontsize=8)
ax.fill([re_lo, re_est, re_hi, re_est], [1, 1.25, 1, 0.75], color="tab:blue")
ax.text(1.02, 1, f"{re_est:+.3f} [{re_lo:+.3f}, {re_hi:+.3f}]", transform=ax.get_yaxis_transform(), va="center", fontsize=8)
ax.axvline(0, color="grey", ls="--", lw=0.8)
ax.set_yticks(ypos + [1])
ax.set_yticklabels([short[k] for k in rows] + [f"RE pooled (DL), k={len(pool)}"])
ax.set_xlabel("theta = ln(T_AI / T_noAI)  (negative = faster with AI)")
ax.set_title("Forest plot: harmonised effects of AI access on completion time")
ax.set_ylim(0.3, len(rows) + 1.7)
fig.subplots_adjust(left=0.30, right=0.74)
fig.savefig(os.path.join(OUT, "figure1.png"), dpi=150)

print("=== Results ===")
for k, v in items.items():
    print(k, "theta =", v["harmonised_estimate"], "SE =", v["harmonised_se"], "poolable =", v["poolable"])
print(f"FE {fe:.4f} (SE {fe_se:.4f}); RE {re_est:.4f} [{re_lo:.4f}, {re_hi:.4f}]; Q {Q:.4f}; I2 {i2:.2f}; tau2 {tau2:.5f}; "
      f"pct change RE {results['pooled_pct_change_re']:.2f}%")