"""F4-conf-01: extract printed treatment effects / uncertainty from two source extracts.
Parses data/sources/*.txt, prints every matched source line, derives two quantities,
runs internal consistency checks, writes out/results.json. Deterministic; no randomness."""
import json, math, re, os
from scipy.stats import norm

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
SRC = os.path.join(ROOT, "data", "sources")
DCR_ID, RI_ID = "arxiv-2609.27732v1", "arxiv-2606.14621v1"
Z = 1.96  # normal critical value for 95% two-sided intervals

def read(name):
    with open(os.path.join(ROOT, "data", name), encoding="utf-8") as f:
        return f.read()

def nlines(txt):  # equals `wc -l` (newline count)
    return txt.count("\n")

dcr_txt = read(f"sources/{DCR_ID}.txt"); ri_txt = read(f"sources/{RI_ID}.txt"); readme = read("README.md")
dcr = dcr_txt.split("\n"); ri = ri_txt.split("\n")

def block(lines, start_pat, end_pat):
    s = next(i for i, l in enumerate(lines) if l.startswith(start_pat))
    e = next(i for i in range(s + 1, len(lines)) if lines[i].startswith(end_pat))
    return s, e

def cells(line):
    return [c.strip() for c in line.split("|")]

def num(cell):
    """Return (printed string without stars/parentheses, float, stars)."""
    m = re.match(r"^\(?(-?[\d,]*\.?\d+)\)?(\**)$", cell.strip())
    s = m.group(1).replace(",", "")
    return s, float(s), m.group(2)

def reg_table(start_pat, end_pat):
    s, e = block(dcr, start_pat, end_pat)
    rows = [(s + i, l) for i, l in enumerate(dcr[s:e]) if l.strip()]
    header = None; out = {}; last = None
    for ln, l in rows[1:]:
        c = cells(l)
        if header is None:
            header = c; header_ln = ln; continue
        label = c[0]
        if label == "" and last is not None:  # SE line belongs to previous coef row
            out[last]["se"] = (ln + 1, l, c[1:])
        else:
            out[label] = {"coef": (ln + 1, l, c[1:])}; last = label
    return header_ln + 1, header, out

def get(tab, row, col, kind="coef"):
    hln, header, rows = tab
    j = header.index(col) - 1
    ln, line, c = rows[row][kind]
    print(f"  [line {ln}] {line.strip()}   -> row '{row}', col '{col}' ({kind}): {c[j]}")
    return num(c[j])

print(f"== {DCR_ID} Table 6")
t6 = reg_table("Table 6:", "Note:")
att5 = get(t6, "ATT", "5 km"); att5se = get(t6, "ATT", "5 km", "se")
att3 = get(t6, "ATT", "3 km"); att3se = get(t6, "ATT", "3 km", "se")
obs5 = get(t6, "Observations", "5 km")
print(f"== {DCR_ID} Table 7")
t7 = reg_table("Table 7:", "Note:")
nar3 = get(t7, "Narcotics", "3 km"); nar3se = get(t7, "Narcotics", "3 km", "se")
fat1 = get(t7, "Drug-related (fatal)", "1 km"); fat1se = get(t7, "Drug-related (fatal)", "1 km", "se")
hiv5 = get(t7, "HIV", "5 km"); hiv5se = get(t7, "HIV", "5 km", "se")
print(f"== {DCR_ID} Table A19")
ta19 = reg_table("Table A19:", "Note:")
a19 = get(ta19, "ATT", "CS-DR"); a19se = get(ta19, "ATT", "CS-DR", "se")

# exact p-value for HIV 5 km: does the extract print one? search for p= near HIV
hiv_p_printed = re.search(r"HIV[^\n]*p\s*=\s*0\.\d+", dcr_txt) is not None
print(f"  exact p-value printed for HIV ATT anywhere in {DCR_ID}: {hiv_p_printed}")

# ---- RESEA Table 2
print(f"== {RI_ID} Table 2")
s = next(i for i, l in enumerate(ri) if l.startswith("Outcome | Treatment | Control"))
hdr = cells(ri[s]); t2 = {}
for i in range(s + 1, len(ri)):
    if not ri[i].strip(): break
    c = cells(ri[i]); t2[c[0]] = (i + 1, ri[i], c)

def ri_cell(row, col):
    ln, line, c = t2[row]; v = c[hdr.index(col)]
    print(f"  [line {ln}] {line.strip()}\n      -> row '{row}', col '{col}': {v}")
    return v

def parse_effect(v):
    # simpler robust split
    pt, p, ci = [x.strip() for x in v.split(";")]
    est = re.sub(r"[^\d.\-]", "", pt.replace("-$", "-"))
    lo, hi = re.match(r"95% CI:\s*\[(-?[\d,.]+),\s*(-?[\d,.]+)\]", ci).groups()
    return est, p, lo.replace(",", ""), hi.replace(",", "")

ATE = "Average treatment effect estimate"
w_est, w_p, w_lo, w_hi = parse_effect(ri_cell("Wage", ATE))
r_est, r_p, r_lo, r_hi = parse_effect(ri_cell("Reemployment", ATE))
u_est, u_p, u_lo, u_hi = parse_effect(ri_cell("Weeks on UI", ATE))
cm = ri_cell("Wage", "Control")
ctrl_wage = re.match(r"\$([\d,]+)", cm).group(1).replace(",", "")
w_se_printed = bool(re.search(r"\b(SE|s\.e\.|standard error)\b", t2["Wage"][1][len("Wage"):], re.I))
print(f"  standard error printed for wage ATE: {w_se_printed}")

def asnum(s):
    return int(s) if re.fullmatch(r"-?\d+", s) else float(s)

# ---- derived quantities
ci_lo = round(att5[1] - Z * att5se[1], 3); ci_hi = round(att5[1] + Z * att5se[1], 3)
w_se = round((float(w_hi) - float(w_lo)) / (2 * Z), 1)

# ---- consistency checks (reported as data issues where they fail)
checks = []
def star_ok(coef, se, stars):
    p = 2 * norm.sf(abs(coef / se))
    exp = "***" if p < .01 else "**" if p < .05 else "*" if p < .10 else ""
    return p, exp == stars
for name, (c, s_) in {"T6 5km": (att5, att5se), "T6 3km": (att3, att3se), "T7 Narc 3km": (nar3, nar3se),
                      "T7 fatal 1km": (fat1, fat1se), "T7 HIV 5km": (hiv5, hiv5se), "A19 CS-DR": (a19, a19se)}.items():
    p, ok = star_ok(c[1], s_[1], c[2]); checks.append((name, round(p, 4), ok))
    print(f"  star check {name}: z-based p={p:.4f}, stars '{c[2]}' consistent={ok}")
w_p_implied = 2 * norm.sf(float(w_est) / w_se)
w_mid = (float(w_lo) + float(w_hi)) / 2
print(f"  wage ATE: CI midpoint={w_mid}, implied p (normal)={w_p_implied:.4f} vs printed {w_p}")
hiv_p_implied = 2 * norm.sf(abs(hiv5[1] / hiv5se[1]))
print(f"  (not reported) normal-approx p for HIV 5km would be {hiv_p_implied:.3f}; NOT used as a result")

data_issues = [
    {"issue": f"Same estimate (CS-DR, 5 km ATT -0.431) is printed with SE (0.210) in Table 6 but SE ({a19se[0]}) in Appendix Table A19 of {DCR_ID}.",
     "handling": "Keys name Table 6, so dcr_att_5km_se = 0.210 and the derived CI uses 0.210; A19 value not used."},
    {"issue": f"Control-group size inconsistent in {DCR_ID}: Tables 6/7 and Figure 6 note say 608 never-treated regions; Table A19 note says 640.",
     "handling": "Recorded only; no requested key depends on it."},
    {"issue": f"{DCR_ID} abstract says fatal hospitalisations 'fall by 35%' and narcotics 'rise by a comparable amount'; Table 7 5 km fatal ATT -0.354 log points corresponds to exp(-0.354)-1 = -29.8% (Section 6.1 says ~30%), and narcotics +0.390 to +47.7%. The abstract appears to read log points as percent.",
     "handling": "Recorded; values extracted as printed in Table 7 (log points), no conversion applied."},
    {"issue": f"{DCR_ID} text says 'HIV estimates are insignificant and close to zero', while the 5 km HIV ATT is -0.340 (SE 0.294), not small in magnitude though insignificant.",
     "handling": "Recorded; significance flag follows the table's stars (none => not significant at 10%)."},
    {"issue": f"{DCR_ID} prints no exact p-values, only significance stars.",
     "handling": "dcr_hiv_att_5km_p_value = null; no p-value was back-computed into the results."},
    {"issue": f"{RI_ID} running text reports the weeks-on-UI CI as magnitudes '[1.75 weeks, 2.24 weeks]' for a 'reduction', whereas Table 2 prints the signed interval [-2.24, -1.75].",
     "handling": "Used Table 2 signed values (lower bound -2.24 first)."},
    {"issue": f"Rendering defect in {RI_ID} text: 'p ¡ 0.001' (should read p < 0.001); Tables 2-3 were re-set by the extract editor.",
     "handling": "No effect on extracted values, which come from the re-set Table 2 whose values are stated to be verbatim."},
    {"issue": f"{RI_ID} prints no standard errors for ATEs, only p-values and 95% CIs; Table 2 caption says CIs follow 'the methodology described in the main text', which is omitted from the extract.",
     "handling": "SE derived as CI width / (2*1.96) assuming a symmetric normal-approximation interval; check: implied p = %.3f vs printed %s." % (w_p_implied, w_p)},
]
bad = [c for c in checks if not c[2]]
if bad:
    data_issues.append({"issue": "Star/SE mismatch: " + str(bad), "handling": "Recorded; printed values kept."})

L = lambda src, t, r, c: f'{src} | Table {t} | row "{r}" | column "{c}"'
res = {
    "task_id": "F4-conf-01",
    "dcr_att_5km": asnum(att5[0]), "dcr_att_5km_se": asnum(att5se[0]),
    "dcr_att_3km": asnum(att3[0]), "dcr_att_3km_se": asnum(att3se[0]),
    "dcr_obs_5km": int(obs5[0]),
    "dcr_narcotics_att_3km": asnum(nar3[0]), "dcr_narcotics_att_3km_se": asnum(nar3se[0]),
    "dcr_fatal_att_1km": asnum(fat1[0]), "dcr_fatal_att_1km_se": asnum(fat1se[0]),
    "dcr_hiv_att_5km": asnum(hiv5[0]), "dcr_hiv_att_5km_se": asnum(hiv5se[0]),
    "dcr_hiv_att_5km_p_value": None,
    "dcr_hiv_att_5km_significant_10pct": "*" in hiv5[2],
    "dcr_att_5km_ci_low": ci_lo, "dcr_att_5km_ci_high": ci_hi,
    "ri_wage_ate": asnum(w_est), "ri_wage_ate_ci_low": asnum(w_lo), "ri_wage_ate_ci_high": asnum(w_hi),
    "ri_wage_ate_se_reported": w_se_printed, "ri_wage_ate_se_derived": w_se,
    "ri_control_mean_wage": asnum(ctrl_wage),
    "ri_reemployment_ate_pp": asnum(r_est),
    "ri_weeks_ui_ate": asnum(u_est), "ri_weeks_ui_ate_ci_low": asnum(u_lo), "ri_weeks_ui_ate_ci_high": asnum(u_hi),
    "locators": {
        "dcr_att_5km": L(DCR_ID, 6, "ATT", "5 km"),
        "dcr_att_5km_se": L(DCR_ID, 6, "ATT", "5 km") + " | standard error (parenthesised line below)",
        "dcr_att_3km": L(DCR_ID, 6, "ATT", "3 km"),
        "dcr_att_3km_se": L(DCR_ID, 6, "ATT", "3 km") + " | standard error (parenthesised line below)",
        "dcr_obs_5km": L(DCR_ID, 6, "Observations", "5 km"),
        "dcr_narcotics_att_3km": L(DCR_ID, 7, "Narcotics", "3 km"),
        "dcr_narcotics_att_3km_se": L(DCR_ID, 7, "Narcotics", "3 km") + " | standard error (parenthesised line below)",
        "dcr_fatal_att_1km": L(DCR_ID, 7, "Drug-related (fatal)", "1 km"),
        "dcr_fatal_att_1km_se": L(DCR_ID, 7, "Drug-related (fatal)", "1 km") + " | standard error (parenthesised line below)",
        "dcr_hiv_att_5km": L(DCR_ID, 7, "HIV", "5 km"),
        "dcr_hiv_att_5km_se": L(DCR_ID, 7, "HIV", "5 km") + " | standard error (parenthesised line below)",
        "dcr_hiv_att_5km_p_value": f"not reported in {DCR_ID} (Table 7 prints only significance stars; no exact p-values anywhere in the extract)",
        "dcr_att_5km_ci_low": L(DCR_ID, 6, "ATT", "5 km") + " | derived from printed coefficient and SE (lower bound)",
        "dcr_att_5km_ci_high": L(DCR_ID, 6, "ATT", "5 km") + " | derived from printed coefficient and SE (upper bound)",
        "ri_wage_ate": L(RI_ID, 2, "Wage", ATE) + " | point estimate",
        "ri_wage_ate_ci_low": L(RI_ID, 2, "Wage", ATE) + " | 95% CI lower bound",
        "ri_wage_ate_ci_high": L(RI_ID, 2, "Wage", ATE) + " | 95% CI upper bound",
        "ri_wage_ate_se_derived": L(RI_ID, 2, "Wage", ATE) + " | derived from printed 95% CI bounds (SE not reported in source)",
        "ri_control_mean_wage": L(RI_ID, 2, "Wage", "Control") + " | mean",
        "ri_reemployment_ate_pp": L(RI_ID, 2, "Reemployment", ATE) + " | point estimate",
        "ri_weeks_ui_ate": L(RI_ID, 2, "Weeks on UI", ATE) + " | point estimate",
        "ri_weeks_ui_ate_ci_low": L(RI_ID, 2, "Weeks on UI", ATE) + " | 95% CI lower bound",
        "ri_weeks_ui_ate_ci_high": L(RI_ID, 2, "Weeks on UI", ATE) + " | 95% CI upper bound",
    },
    "derivation_notes": {
        "dcr_att_5km_ci": "Normal approximation, critical value 1.96: ATT +/- 1.96*SE = -0.431 +/- 1.96*0.210 using Table 6 printed values; rounded to 3 decimals. Assumes asymptotic normality of the clustered-SE ATT estimator (no t / small-cluster correction).",
        "ri_wage_ate_se": "SE = (upper - lower)/(2*1.96) = (2141 - 166)/3.92 from Table 2 printed 95%% CI, assuming a symmetric normal-approximation (Wald) interval; rounded to 0.1 USD. With N = 23,549 a t critical value is indistinguishable from 1.96. Consistency check: implied two-sided p = %.3f vs printed p=0.022." % w_p_implied,
    },
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": data_issues,
    "n_input_rows": {"README.md": nlines(readme), f"{DCR_ID}.txt": nlines(dcr_txt), f"{RI_ID}.txt": nlines(ri_txt)},
    "n_analysed": 26,
    "random_seed": None,
}
os.makedirs(os.path.join(ROOT, "out"), exist_ok=True)
with open(os.path.join(ROOT, "out", "results.json"), "w") as f:
    json.dump(res, f, indent=2, ensure_ascii=False)
print(json.dumps({k: v for k, v in res.items() if k not in ("locators", "data_issues", "derivation_notes")}, indent=1))
# sensitivity figures quoted in report.md (not stored as results)
print("sensitivity: CI with z=2.0:", round(att5[1]-2*att5se[1],3), round(att5[1]+2*att5se[1],3),
      "| CI with A19 SE:", round(att5[1]-Z*a19se[1],3), round(att5[1]+Z*a19se[1],3),
      "| fatal 5km % =", round(100*(math.exp(-0.354)-1),1), "| narcotics 5km % =", round(100*(math.exp(0.390)-1),1))