"""F4-conf-01: extract printed treatment effects/uncertainty from two source extracts.
Reads only data/; writes out/results.json. No randomness."""
import json, re, os, hashlib
from decimal import Decimal

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

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

def wc_l(p):
    with open(os.path.join(D, p), "rb") as f:
        return f.read().count(b"\n")

dcr = read(f"sources/{DCR_ID}.txt").split("\n")
ri = read(f"sources/{RI_ID}.txt").split("\n")

def table_block(lines, title_prefix):
    i = next(k for k, l in enumerate(lines) if l.startswith(title_prefix))
    j = next(k for k in range(i + 1, len(lines)) if lines[k].startswith("Note:"))
    return [(k + 1, lines[k]) for k in range(i, j) if lines[k].strip()]

def num(s):
    """Parse printed number keeping printed decimals; returns (float, printed str)."""
    s = s.strip().replace("$", "").replace(",", "").replace("*", "").strip("()")
    s = s.replace("%", "")
    return (int(s) if re.fullmatch(r"-?\d+", s) else float(s)), s

def dcr_cell(block, row, col):
    header = [c.strip() for c in block[1][1].split("|")]
    ci = header.index(col)
    for idx, (ln, l) in enumerate(block):
        cells = [c.strip() for c in l.split("|")]
        if cells[0] == row:
            raw = cells[ci]
            print(f"[{DCR_ID} line {ln}] {l}")
            out = {"value": num(raw)[0], "stars": raw.count("*"), "raw": raw}
            nxt = block[idx + 1] if idx + 1 < len(block) else None
            if nxt and nxt[1].lstrip().startswith("| ("):
                print(f"[{DCR_ID} line {nxt[0]}] {nxt[1]}")
                se_raw = [c.strip() for c in nxt[1].split("|")][ci]
                out["se"] = num(se_raw)[0]; out["se_raw"] = se_raw
            return out
    raise KeyError(row)

t6 = table_block(dcr, "Table 6:")
t7 = table_block(dcr, "Table 7:")
for ln, l in t6[:2] + t7[:2]:
    print(f"[{DCR_ID} line {ln}] {l}")

a5 = dcr_cell(t6, "ATT", "5 km"); a3 = dcr_cell(t6, "ATT", "3 km")
obs5 = dcr_cell(t6, "Observations", "5 km")
narc3 = dcr_cell(t7, "Narcotics", "3 km")
fat1 = dcr_cell(t7, "Drug-related (fatal)", "1 km")
hiv5 = dcr_cell(t7, "HIV", "5 km")

# exact p-value search: any "p=" / "p-value" in DCR extract tied to HIV?
dcr_text = "\n".join(dcr)
p_hits = [l for l in dcr if re.search(r"p\s*=\s*0?\.\d", l)]
hiv_p = None  # stars only (none on HIV 5 km); no exact p-value printed
print(f"DCR lines with an exact 'p=' value: {len(p_hits)}")

# derived CI (Table 6, 5 km); round to 4 dp (one more than the printed 3 dp)
ci_lo = round(a5["value"] - Z * a5["se"], 4)
ci_hi = round(a5["value"] + Z * a5["se"], 4)

# RESEA Table 2
t2_hdr = next(k for k, l in enumerate(ri) if l.startswith("Outcome | Treatment | Control"))
header = [c.strip() for c in ri[t2_hdr].split("|")]
print(f"[{RI_ID} line {t2_hdr+1}] {ri[t2_hdr]}")
def ri_row(label):
    for k in range(t2_hdr + 1, len(ri)):
        cells = [c.strip() for c in ri[k].split("|")]
        if cells[0] == label:
            print(f"[{RI_ID} line {k+1}] {ri[k]}")
            return dict(zip(header, cells))
    raise KeyError(label)
ATE = "Average treatment effect estimate"
def parse_effect(cell):
    est = cell.split(";")[0]
    est = re.sub(r"\s*(pp|weeks|wks)$", "", est.strip())
    m = re.search(r"95% CI: \[(.+?),\s*(-?[\d.]+(?:,\d{3})*(?:\.\d+)?)\]", cell)
    lo, hi = m.group(1), m.group(2)
    se_printed = bool(re.search(r"\bSE\b|standard error", cell, re.I))
    return num(est)[0], num(lo)[0], num(hi)[0], se_printed
w = ri_row("Wage"); r = ri_row("Reemployment"); u = ri_row("Weeks on UI")
w_ate, w_lo, w_hi, w_se_printed = parse_effect(w[ATE])
r_ate, _, _, _ = parse_effect(r[ATE])
u_ate, u_lo, u_hi, _ = parse_effect(u[ATE])
ctrl_mean = num(w["Control"].split("(")[0])[0]
w_se = round((w_hi - w_lo) / (2 * Z), 2)

S6 = f"{DCR_ID} | Table 6"; S7 = f"{DCR_ID} | Table 7"; S2 = f"{RI_ID} | Table 2"
loc = {
 "dcr_att_5km": f'{S6} | row "ATT" | column "5 km"',
 "dcr_att_5km_se": f'{S6} | row "ATT" (standard error line in parentheses) | column "5 km"',
 "dcr_att_3km": f'{S6} | row "ATT" | column "3 km"',
 "dcr_att_3km_se": f'{S6} | row "ATT" (standard error line in parentheses) | column "3 km"',
 "dcr_obs_5km": f'{S6} | row "Observations" | column "5 km"',
 "dcr_narcotics_att_3km": f'{S7} | row "Narcotics" | column "3 km"',
 "dcr_narcotics_att_3km_se": f'{S7} | row "Narcotics" (standard error line in parentheses) | column "3 km"',
 "dcr_fatal_att_1km": f'{S7} | row "Drug-related (fatal)" | column "1 km"',
 "dcr_fatal_att_1km_se": f'{S7} | row "Drug-related (fatal)" (standard error line in parentheses) | column "1 km"',
 "dcr_hiv_att_5km": f'{S7} | row "HIV" | column "5 km"',
 "dcr_hiv_att_5km_se": f'{S7} | row "HIV" (standard error line in parentheses) | column "5 km"',
 "dcr_hiv_att_5km_p_value": f"not reported in {DCR_ID} (Table 7 gives only significance stars; the HIV 5 km cell has none)",
 "dcr_att_5km_ci_low": f'{S6} | row "ATT" | column "5 km" | derived from printed ATT and SE (normal, 1.96)',
 "dcr_att_5km_ci_high": f'{S6} | row "ATT" | column "5 km" | derived from printed ATT and SE (normal, 1.96)',
 "ri_wage_ate": f'{S2} | row "Wage" | column "Average treatment effect estimate"',
 "ri_wage_ate_ci_low": f'{S2} | row "Wage" | column "Average treatment effect estimate" (95% CI lower bound)',
 "ri_wage_ate_ci_high": f'{S2} | row "Wage" | column "Average treatment effect estimate" (95% CI upper bound)',
 "ri_wage_ate_se_derived": f'{S2} | row "Wage" | column "Average treatment effect estimate" | derived from printed 95% CI (width / (2 x 1.96)); SE not reported in {RI_ID}',
 "ri_control_mean_wage": f'{S2} | row "Wage" | column "Control"',
 "ri_reemployment_ate_pp": f'{S2} | row "Reemployment" | column "Average treatment effect estimate"',
 "ri_weeks_ui_ate": f'{S2} | row "Weeks on UI" | column "Average treatment effect estimate"',
 "ri_weeks_ui_ate_ci_low": f'{S2} | row "Weeks on UI" | column "Average treatment effect estimate" (95% CI lower bound)',
 "ri_weeks_ui_ate_ci_high": f'{S2} | row "Weeks on UI" | column "Average treatment effect estimate" (95% CI upper bound)',
}

issues = [
 {"issue": f"{DCR_ID}: Table 6 prints SE (0.210) for the 5 km ATT -0.431**, while Appendix Table A19 (CS-DR column) prints the same ATT with SE (0.218).",
  "handling": "Used Table 6 value (0.210) as the key definition names Table 6; discrepancy reported."},
 {"issue": f"{DCR_ID}: Tables 6/7 and Figure 6 state 608 never-treated control regions, the A19 note states 640; the Table 6 observation counts (16,175/16,875/17,250) equal 25 x (640 + 7/35/50 treated regions), i.e. consistent with 640 controls.",
  "handling": "Observation count reported as printed (17250); inconsistency not resolvable from the extract, reported only."},
 {"issue": f"{DCR_ID}: abstract says fatal hospitalizations 'fall by 35%', whereas Table 7 fatal 5 km ATT -0.354 log points implies about -29.8% (Section 6.1 says ~30%); prose also calls HIV estimates 'close to zero' though 3/5 km ATTs are -0.239/-0.340.",
  "handling": "Prose not used; table values extracted as printed."},
 {"issue": f"{RI_ID}: prose reports the weeks-on-UI CI as magnitudes [1.75 weeks, 2.24 weeks]; Table 2 prints [-2.24, -1.75]. Prose also has a rendering defect 'p ¡ 0.001'.",
  "handling": "Table 2 values used (lower bound -2.24 first)."},
]

res = {
 "task_id": "F4-conf-01",
 "dcr_att_5km": a5["value"], "dcr_att_5km_se": a5["se"],
 "dcr_att_3km": a3["value"], "dcr_att_3km_se": a3["se"],
 "dcr_obs_5km": int(obs5["value"]),
 "dcr_narcotics_att_3km": narc3["value"], "dcr_narcotics_att_3km_se": narc3["se"],
 "dcr_fatal_att_1km": fat1["value"], "dcr_fatal_att_1km_se": fat1["se"],
 "dcr_hiv_att_5km": hiv5["value"], "dcr_hiv_att_5km_se": hiv5["se"],
 "dcr_hiv_att_5km_p_value": hiv_p,
 "dcr_hiv_att_5km_significant_10pct": hiv5["stars"] >= 1,
 "dcr_att_5km_ci_low": ci_lo, "dcr_att_5km_ci_high": ci_hi,
 "ri_wage_ate": w_ate, "ri_wage_ate_ci_low": w_lo, "ri_wage_ate_ci_high": w_hi,
 "ri_wage_ate_se_reported": w_se_printed,
 "ri_wage_ate_se_derived": w_se,
 "ri_control_mean_wage": ctrl_mean,
 "ri_reemployment_ate_pp": r_ate,
 "ri_weeks_ui_ate": u_ate, "ri_weeks_ui_ate_ci_low": u_lo, "ri_weeks_ui_ate_ci_high": u_hi,
 "locators": loc,
 "derivation_notes": {
  "dcr_att_5km_ci": "Normal (Wald) 95% CI: ATT +/- 1.96 x SE using Table 6 printed ATT -0.431 and SE 0.210 (clustered SE, large-sample normal approximation); rounded to 4 decimals.",
  "ri_wage_ate_se": "SE = (upper - lower) / (2 x 1.96) using Table 2 printed 95% CI [166, 2,141]; assumes a symmetric normal-based CI (printed CI is symmetric around 1,153 to within rounding: 987 vs 988); rounded to 2 decimals. Using 2.0 instead would give 493.75."},
 "feasibility": "complete",
 "feasibility_reason": None,
 "data_issues": issues,
 "n_input_rows": {"README.md": wc_l("README.md"),
                  f"sources/{DCR_ID}.txt": wc_l(f"sources/{DCR_ID}.txt"),
                  f"sources/{RI_ID}.txt": wc_l(f"sources/{RI_ID}.txt")},
 "n_analysed": 25,
 "random_seed": None,
}
# baseline: cruder "2-sigma" convention for the two derived quantities (sensitivity reference)
res["baseline_dcr_att_5km_ci_2se"] = [round(a5["value"] - 2.0 * a5["se"], 4), round(a5["value"] + 2.0 * a5["se"], 4)]
res["baseline_ri_wage_ate_se_2se"] = round((w_hi - w_lo) / 4.0, 2)
# printed strings (JSON numbers drop trailing zeros, e.g. 0.210 -> 0.21)
res["printed_values"] = {
 "dcr_att_5km": a5["raw"], "dcr_att_5km_se": a5["se_raw"],
 "dcr_att_3km": a3["raw"], "dcr_att_3km_se": a3["se_raw"],
 "dcr_obs_5km": obs5["raw"],
 "dcr_narcotics_att_3km": narc3["raw"], "dcr_narcotics_att_3km_se": narc3["se_raw"],
 "dcr_fatal_att_1km": fat1["raw"], "dcr_fatal_att_1km_se": fat1["se_raw"],
 "dcr_hiv_att_5km": hiv5["raw"], "dcr_hiv_att_5km_se": hiv5["se_raw"],
 "ri_wage_cell": w[ATE], "ri_control_wage_cell": w["Control"],
 "ri_reemployment_cell": r[ATE], "ri_weeks_ui_cell": u[ATE],
}
# consistency diagnostics (printed only)
import math
zw = w_ate / w_se
print(f"Check: wage ATE z={zw:.3f}, implied two-sided normal p={math.erfc(abs(zw)/math.sqrt(2)):.4f} (printed p=0.022)")
print(f"Check: DCR 5km |t|={abs(a5['value'])/a5['se']:.3f} (printed **, i.e. p<0.05)")
print(f"Check: HIV 5km |t|={abs(hiv5['value'])/hiv5['se']:.3f} (no star)")
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", "derivation_notes", "data_issues")}, indent=1))