"""T3-synthesis: extraction, harmonisation to 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
import numpy as np
import pandas as pd
from scipy.stats import norm
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DATA = os.path.join(ROOT, "data")
OUT = os.path.join(ROOT, "out")
Z = 1.96

studies = json.load(open(os.path.join(DATA, "studies.json")))
src = {s["id"]: open(os.path.join(DATA, "sources", s["source"] + ".txt"), encoding="utf-8").read().splitlines()
       for s in studies}
SRCNAME = {s["id"]: s["source"] for s in studies}


def find(sid, pattern, after=None):
    """Return (line_no, line) of first line matching regex `pattern` (optionally after a line matching `after`)."""
    lines = src[sid]
    start = 0
    if after is not None:
        start = next(i for i, l in enumerate(lines) if re.search(after, l)) + 1
    for i in range(start, len(lines)):
        if re.search(pattern, lines[i]):
            print(f"[{sid} {SRCNAME[sid]} L{i+1}] {lines[i].strip()}")
            return i, lines[i]
    raise ValueError(f"{sid}: pattern not found: {pattern}")


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


print("=== Source lines relied upon ===")
# ---------------- S1 ----------------
i, l = find("S1", r"ExpCon \|")
s1_coefs = re.findall(r"-?\d+\.\d+", l)            # Models 1-4
s1_b_str = l.split()[1]                             # Model 2 token as printed, e.g. "-0.24."
s1_b = float(s1_coefs[1])
_, l_se = i + 2, src["S1"][i + 2]
print(f"[S1 {SRCNAME['S1']} L{i+3}] {l_se.strip()}")
s1_se_str = re.findall(r"\(\d+\.\d+\)", l_se)[1]
s1_se = float(s1_se_str.strip("()"))
_, lci = find("S1", r"best-fit Model 2")
m = re.search(r"\[(-?\d+\.\d+),\s*(-?\d+\.\d+)\]", lci)
s1_L, s1_U = float(m.group(1)), float(m.group(2))
s1_p = re.search(r"p=(0?\.\d+)", lci).group(1)
_, lN = find("S1", r"TABLE III")
s1_n = re.search(r"n = (\d+)", lN).group(1)
find("S1", r"completed the task in 96 min")
assert s1_b == -0.24 and s1_se == 0.14

# ---------------- S2 ----------------
_, l2 = find("S2", r"increases completion time by 19%")
_, l2n = find("S2", r"Developers complete 136 issues")
find("S2", r"16 developers with moderate AI experience")
find("S2", r"Figure 1, which displays the confidence interval")
s2_c = float(re.search(r"by (\d+)%", l2).group(1)) / 100.0
assert s2_c == 0.19

# ---------------- S3 ----------------
_, l3 = find("S3", r"^All Content \|", after=r"Aggregate Results")
cells = [c.strip() for c in l3.split("|")]
s3_overall, s3_treat, s3_ctrl, s3_impr = cells[1], cells[2], cells[3], cells[4]
find("S3", r"Table 3: Comparison of task duration")
find("S3", r"We had 181 subjects")
find("S3", r"Editorial note: the paper marks significance")
s3_c = float(re.search(r"(-?\d+\.\d+)%", s3_impr).group(1)) / 100.0
assert s3_c == -0.2979

# ---------------- S4 ----------------
_, la = find("S4", r"^Aware Time per Email")
_, lc = find("S4", r"^Control Time per Email")
_, ls = find("S4", r"^% Time Savings")
_, lp = find("S4", r"^p-value \|")
find("S4", r"Control \( N\{=\}52")
find("S4", r"Aware \( N\{=\}56")
last = lambda s: s.split("|")[-1].strip()
s4_aware, s4_ctrl, s4_sav, s4_p = last(la), last(lc), last(ls), last(lp)
s4_c = -float(s4_sav) / 100.0          # "time savings" of 11.1% = change in time of -11.1%
s4_pv = float(s4_p)
assert s4_c == -0.111 and s4_pv == 0.471

# ---------------- S5 ----------------
i5, l5 = find("S5", r"TreatPost \|", after=r"A: Pre 27/28.03")
c5 = [c.strip() for c in l5.split("|")]
j5 = next(j for j in range(i5 + 1, len(src["S5"])) if src["S5"][j].strip())
l5se = src["S5"][j5]
print(f"[S5 {SRCNAME['S5']} L{j5+1}] {l5se.strip()}")
_, l5n = find("S5", r"^N \|", after=r"A: Pre 27/28.03")
s5_b, s5_se, s5_n = c5[3], [c.strip() for c in l5se.split("|")][3], [c.strip() for c in l5n.split("|")][3]
find("S5", r"aggregate Output measure")
assert s5_b == "-0.278" and s5_se == "(0.295)" and s5_n == "18332"
print("=== end of source lines ===\n")

# ---------------- Harmonisation ----------------
items = {}
items["S1"] = dict(
    estimate_as_printed=f"{s1_b_str} (ExpCon coefficient, Model 2, LogToT; marker '.' = p<.1); text: beta_1=-0.24, 'about 21%' faster",
    uncertainty_as_printed=f"SE {s1_se_str}; 95% CI = [{s1_L}, {s1_U}]; p={s1_p}",
    sample_as_printed=f"n = {s1_n} (Table III caption); Section VI-B: N=47 (AI), N=46 (no AI)",
    definition="OLS coefficient of the experimental-condition (AI) indicator on natural-log time on task, with developer-level covariates (Model 2); RCT, one complex enterprise coding task",
    locator=f"{SRCNAME['S1']}: Table III, row ExpCon, column (2) 'Dev' (+ SE line below); Section VI-C (CI and p for Model 2)",
    harmonisable=True, harmonised_estimate=s1_b, harmonised_se=s1_se, poolable=True,
    transformation=f"log-time regression coefficient: theta = b = {s1_b}; SE(theta) = SE(b) = {s1_se} (printed CI gives (0.03-(-0.51))/(2*1.96) = {(s1_U-s1_L)/(2*Z):.4f}, consistent)",
    reason_if_not=None)
th2 = math.log(1 + s2_c)
items["S2"] = dict(
    estimate_as_printed="+19% ('allowing AI actually increases completion time by 19%'; 'implement issues in 19% more time on average')",
    uncertainty_as_printed="none printed in the extract (95% CI shown only in Figure 1, not reproduced)",
    sample_as_printed="16 developers, 246 tasks; 136 issues AI-allowed, 110 issues 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 (Sections 2.3, D.1-D.3); issue-level RCT",
    locator=f"{SRCNAME['S2']}: 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: theta = ln(1 + 0.19) = {th2:.6f}",
    reason_if_not="Not poolable: no SE, CI or p-value for the headline estimate is printed in the extract; its CI appears only in Figure 1 (and Figure 15), which is not reproduced, so SE(theta) cannot be obtained.")
th3 = math.log(1 + s3_c)
items["S3"] = dict(
    estimate_as_printed=f"Treatment {s3_treat} min, Control {s3_ctrl} min, Overall {s3_overall} min, Improvement {s3_impr}",
    uncertainty_as_printed="significance stars only: *** (defined as 0<p<=0.01); no SE or CI printed",
    sample_as_printed="181 subjects (randomised half to treatment, half to control)",
    definition="percentage difference in mean task duration (minutes) for Copilot vs control, holding accuracy constant (control duration predicted at treatment accuracy, Appendix B); RCT",
    locator=f"{SRCNAME['S3']}: Table 3, block 'Aggregate Results', row 'All Content', columns Treatment/Control/Improvement",
    harmonisable=True, harmonised_estimate=th3, harmonised_se=None, poolable=False,
    transformation=f"percentage change: theta = ln(1 + (-0.2979)) = {th3:.6f} (ratio of printed durations ln(44.29/63.10) = {math.log(float(s3_treat)/float(s3_ctrl)):.6f})",
    reason_if_not="Not poolable: the only uncertainty information is significance stars (***, p<=0.01); SE(theta) cannot be obtained.")
th4 = math.log(1 + s4_c)
z4 = norm.ppf(1 - s4_pv / 2)
se4 = abs(th4) / z4
items["S4"] = dict(
    estimate_as_printed=f"Aware {s4_aware} s, Control {s4_ctrl} s per email; % Time Savings {s4_sav}",
    uncertainty_as_printed=f"p-value {s4_p}",
    sample_as_printed="Aware N=56, Control N=52 (Blind N=59 not used); 167 analysts in total",
    definition="percentage time saving per email (full 25-email queue) of the Aware arm (agent verdicts, rationale, agent-prioritised queue) vs Control (no AI); RCT",
    locator=f"{SRCNAME['S4']}: Section 3.3, Table 3, 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 c = -0.111 (11.1% saving): theta = ln(1 - 0.111) = {th4:.6f}; APPROXIMATE SE from two-sided p: SE = |theta|/z(1-0.471/2) = {abs(th4):.6f}/{z4:.6f} = {se4:.6f} (ratio of printed times ln(50.5/56.9) = {math.log(float(s4_aware)/float(s4_ctrl)):.6f})",
    reason_if_not=None)
items["S5"] = dict(
    estimate_as_printed=f"{s5_b} (Treat x Post, Output, Panel A column (3))",
    uncertainty_as_printed=f"SE {s5_se} (clustered at user level); not marked significant",
    sample_as_printed=f"N {s5_n}",
    definition="difference-in-differences coefficient of Italy x post-ban on GitHub Output (count of push, pull requests, PR comments, commit comments, create, issues events; users with affiliation), Pre 27/28.03 vs Post 03/04.04 2023",
    locator=f"{SRCNAME['S5']}: Table 1, panel A, row Treat x Post (+ SE line), column (3) Output; row N",
    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) around the withdrawal of ChatGPT access in a non-randomised difference-in-differences design, not a comparison of task-completion times with vs without AI on the same tasks.")

# ---------------- Baseline (P2): unweighted means ----------------
harm = {k: v for k, v in items.items() if v["harmonisable"]}
pool = {k: v for k, v in items.items() if v["poolable"]}
baseline_mean_harmonised = float(np.mean([v["harmonised_estimate"] for v in harm.values()]))
baseline_mean_poolable = float(np.mean([v["harmonised_estimate"] for v in pool.values()]))

# ---------------- Pooling ----------------
y = np.array([v["harmonised_estimate"] for v in pool.values()])
se = np.array([v["harmonised_se"] for v in pool.values()])
w = 1 / se**2
fe = float(np.sum(w * y) / np.sum(w)); fe_se = float(math.sqrt(1 / np.sum(w)))
k = len(y); df = k - 1
Q = float(np.sum(w * (y - fe) ** 2))
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_ - Z * re_se, re_ + Z * re_se

# leave-one-out / sensitivity (reported in report, not required keys)
sens = {}
# S1 with CI-derived SE
se_alt = se.copy(); se_alt[0] = (s1_U - s1_L) / (2 * Z)
w_a = 1 / se_alt**2; sens["fe_with_S1_ci_se"] = float(np.sum(w_a * y) / np.sum(w_a))
# S4 from printed times instead of printed %
y_alt = y.copy(); y_alt[1] = math.log(float(s4_aware) / float(s4_ctrl)); se_y = abs(y_alt[1]) / z4
w_b = 1 / np.array([se[0], se_y])**2; sens["fe_with_S4_from_times"] = float(np.sum(w_b * y_alt) / np.sum(w_b))

def r(x):
    return None if x is None else float(round(x, 10))

for v in items.values():
    v["harmonised_estimate"] = r(v["harmonised_estimate"]); v["harmonised_se"] = r(v["harmonised_se"])

n_input_rows = {"studies.json": len(studies)}
for fn in ["README.md"] + [s["source"] + ".txt" for s in studies]:
    p = os.path.join(DATA, fn) if fn == "README.md" else os.path.join(DATA, "sources", fn)
    with open(p, "rb") as fh:
        n_input_rows[fn] = fh.read().count(b"\n")   # equals `wc -l`

data_issues = [
    {"issue": "S1: Section VI-B prints Mean LogToT=4.6 for the AI group and 4.3 for the no-AI group while stating the AI group was faster (96 vs 114 min); the two values appear swapped.",
     "handling": "Not used; the harmonised value is the Table III Model 2 coefficient, which agrees with the text of VI-C (CI, p, 21%)."},
    {"issue": "S1: Table III caption says n = 96 whereas Section VI-B reports N=47 + N=46 = 93.",
     "handling": "Both reported as printed in sample_as_printed; no effect on theta or SE."},
    {"issue": "S1: printed SE (0.14) and SE implied by the printed 95% CI [-0.51, 0.03] (0.1378) differ by rounding.",
     "handling": "Printed SE(b)=0.14 used (first convention); CI-based SE used as sensitivity analysis."},
    {"issue": "S2: the confidence interval of the headline +19% estimate is displayed only in Figure 1/15, which is not reproduced.",
     "handling": "harmonised_se = null, poolable = false; no remembered value substituted."},
    {"issue": "S3: printed improvement -29.79% differs slightly from the ratio of printed durations (44.29/63.10 - 1 = -29.81%).",
     "handling": "Printed percentage used (theta=-0.35368); duration-ratio value (-0.35396) reported in the transformation string."},
    {"issue": "S3: uncertainty only as significance stars with a non-standard key (*** 0<p<=0.01, * 0.05<p<=0.1).",
     "handling": "Harmonisable but not poolable per convention."},
    {"issue": "S4: printed % Time Savings 11.1 differs from the ratio of printed full-corpus times (50.5/56.9 -> 11.25%).",
     "handling": "Printed percentage and its p-value used; times-based theta used as sensitivity analysis."},
    {"issue": "S4: Table 3 Malicious/Non-Malicious columns (2.8, 1.8 vs 1.9, 2.4, labelled seconds) cannot be reconciled with Full Corpus times of 50.5/56.9 s; units of those columns ambiguous.",
     "handling": "Those columns are not used; only the Full Corpus column (a unit-free ratio) enters."},
    {"issue": "S4: SE recovered from a two-sided p-value (0.471) - approximate, assumes a Wald-type test on the log-ratio scale; the Aware arm also bundles agent queue prioritisation.",
     "handling": "Marked poolable with approximate SE, flagged in transformation and report."},
    {"issue": "S5: Table 1 notes print '** p<0.5' (presumably p<0.05); row label rendered as '\\times TreatPost'; header date 'August 24, 2026' is a rendering artefact; whether Output is logged is not stated unambiguously.",
     "handling": "Value reported as printed; S5 is not harmonisable for substantive reasons regardless."},
]

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, 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": sum(1 for v in items.values() if v["harmonised_estimate"] is not None),
    "k_pooled": k,
    "fe_estimate": r(fe), "fe_se": r(fe_se), "fe_ci_low": r(fe - Z * fe_se), "fe_ci_high": r(fe + Z * fe_se),
    "re_estimate": r(re_), "re_se": r(re_se), "re_ci_low": r(re_lo), "re_ci_high": r(re_hi),
    "q_stat": r(Q), "i2": r(I2), "tau2": r(tau2),
    "heterogeneity_interpretation": het,
    "pooled_significant_5pct": bool(re_lo > 0 or re_hi < 0),
    "pooled_pct_change_re": r(100 * (math.exp(re_) - 1)),
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": data_issues,
    "n_input_rows": n_input_rows,
    "n_analysed": len(items),
    "random_seed": None,
    "baseline_unweighted_mean_theta_harmonised": r(baseline_mean_harmonised),
    "baseline_unweighted_mean_theta_poolable": r(baseline_mean_poolable),
    "pooled_studies": list(pool.keys()),
    "sensitivity_fe_with_S1_ci_se": r(sens["fe_with_S1_ci_se"]),
    "sensitivity_fe_with_S4_from_times": r(sens["fe_with_S4_from_times"]),
    "fe_weights_pct": {kk: r(100 * ww / np.sum(w)) for kk, ww in zip(pool.keys(), w)},
}
with open(os.path.join(OUT, "results.json"), "w") as fh:
    json.dump(results, fh, indent=2, ensure_ascii=False)

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

# ---------------- Forest plot ----------------
short = {s["id"]: s["short"].split(",")[0] for s in studies}
rows = list(harm.keys())
fig, ax = plt.subplots(figsize=(8, 4))
ypos = list(range(len(rows) + 1, 1, -1))
for yy, kk in zip(ypos, rows):
    v = items[kk]
    if v["harmonised_se"] is not None:
        lo, hi = v["harmonised_estimate"] - Z * v["harmonised_se"], v["harmonised_estimate"] + Z * v["harmonised_se"]
        ax.errorbar(v["harmonised_estimate"], yy, xerr=[[v["harmonised_estimate"] - lo], [hi - v["harmonised_estimate"]]],
                    fmt="s", color="k", capsize=3)
        lab = f"{kk} {short[kk]}  {v['harmonised_estimate']:.3f} [{lo:.3f}, {hi:.3f}]"
    else:
        ax.plot(v["harmonised_estimate"], yy, "o", mfc="white", mec="grey")
        lab = f"{kk} {short[kk]}  {v['harmonised_estimate']:.3f} (no SE; not pooled)"
    ax.text(1.02, yy, lab, transform=ax.get_yaxis_transform(), va="center", fontsize=8)
ax.fill([re_lo, re_, re_hi, re_], [1, 1.25, 1, 0.75], color="tab:blue")
ax.text(1.02, 1, f"RE pooled (S1, S4)  {re_:.3f} [{re_lo:.3f}, {re_hi:.3f}]", transform=ax.get_yaxis_transform(),
        va="center", fontsize=8, color="tab:blue")
ax.axvline(0, color="grey", ls="--", lw=0.8)
ax.set_yticks([]); ax.set_ylim(0.3, len(rows) + 1.7)
ax.set_xlabel("theta = ln(T_AI / T_noAI)   (negative = faster with AI)")
ax.set_title("Forest plot: harmonised effects of AI access on task time (S5 not harmonisable)", fontsize=9)
fig.subplots_adjust(right=0.55)
fig.savefig(os.path.join(OUT, "figure1.png"), dpi=150)

print(f"k_harmonised={results['k_harmonised']} k_pooled={k} ({', '.join(pool)})")
print(f"FE theta={fe:.4f} SE={fe_se:.4f} 95%CI [{fe-Z*fe_se:.4f}, {fe+Z*fe_se:.4f}]")
print(f"RE theta={re_:.4f} SE={re_se:.4f} 95%CI [{re_lo:.4f}, {re_hi:.4f}]  pct change {100*(math.exp(re_)-1):.2f}%")
print(f"Q={Q:.4f} df={df} I2={I2:.2f}% tau2={tau2:.4f} ({het})")