"""F4-conf-01: extract printed values from two source extracts, derive two quantities.
Run from workspace root: ~/wp29/venv/bin/python out/analysis.py
Values are parsed from the text; every matched source line is printed."""
import json, re, os

D = "data/sources"
DCR_ID, RI_ID = "arxiv-2609.27732v1", "arxiv-2606.14621v1"
dcr = open(f"{D}/{DCR_ID}.txt", encoding="utf-8").read().splitlines()
ri = open(f"{D}/{RI_ID}.txt", encoding="utf-8").read().splitlines()

def section(lines, start, end):
    i = next(k for k, l in enumerate(lines) if l.startswith(start))
    j = next(k for k in range(i + 1, len(lines)) if lines[k].startswith(end))
    return lines[i:j]

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

def row_with_se(tab, label, table_name):
    """Return (coef cells, se cells) for a coefficient row and the following SE line."""
    idx = [k for k, l in enumerate(tab) if cells(l)[0] == label]
    assert len(idx) == 1, (label, idx)
    k = idx[0]
    se_k = next(m for m in range(k + 1, len(tab)) if tab[m].strip())
    print(f"[{table_name}] matched: {tab[k]!r}\n[{table_name}] SE line : {tab[se_k]!r}")
    return cells(tab[k])[1:], cells(tab[se_k])[1:]

def num(s):
    m = re.match(r"^\(?(-?\d+(?:\.\d+)?)\)?(\**)$", s.strip())
    assert m, s
    return float(m.group(1)), len(m.group(2))

# ---------- DCR Table 6 ----------
t6 = section(dcr, "Table 6:", "Note:")
hdr6 = cells(next(l for l in t6 if l.startswith("| 1 km")))[1:]
col6 = {h: i for i, h in enumerate(hdr6)}
att, se = row_with_se(t6, "ATT", "Table 6")
obs_line = next(l for l in t6 if cells(l)[0] == "Observations")
print(f"[Table 6] matched: {obs_line!r}")
obs = cells(obs_line)[1:]

dcr_att_5km, stars_5 = num(att[col6["5 km"]])
dcr_att_5km_se, _ = num(se[col6["5 km"]])
dcr_att_3km, _ = num(att[col6["3 km"]])
dcr_att_3km_se, _ = num(se[col6["3 km"]])
dcr_obs_5km = int(obs[col6["5 km"]].replace(",", ""))

# ---------- DCR Table 7 ----------
t7 = section(dcr, "Table 7:", "Note:")
hdr7 = cells(next(l for l in t7 if l.startswith("Outcome |")))[1:]
col7 = {h: i for i, h in enumerate(hdr7)}
c, s = row_with_se(t7, "Narcotics", "Table 7")
narc3, _ = num(c[col7["3 km"]]); narc3_se, _ = num(s[col7["3 km"]])
c, s = row_with_se(t7, "Drug-related (fatal)", "Table 7")
fat1, _ = num(c[col7["1 km"]]); fat1_se, _ = num(s[col7["1 km"]])
c, s = row_with_se(t7, "HIV", "Table 7")
hiv5, hiv5_stars = num(c[col7["5 km"]]); hiv5_se, _ = num(s[col7["5 km"]])
hiv5_sig10 = hiv5_stars >= 1  # any star means p<0.10 per table note
# exact p-value: search the DCR extract for any printed p-value (p=...)
pvals = [l for l in dcr if re.search(r"\bp\s*=\s*0?\.\d", l)]
print(f"[DCR] lines with exact p-values printed: {len(pvals)}")
hiv_p = None

# ---------- derived 95% CI for Table 6 5 km ATT ----------
Z = 1.96
ci_lo = round(dcr_att_5km - Z * dcr_att_5km_se, 4)
ci_hi = round(dcr_att_5km + Z * dcr_att_5km_se, 4)

# ---------- RI Table 2 ----------
t2 = section(ri, "Outcome | Treatment | Control", "Table 2:")
hdr2 = cells(t2[0])
col2 = {h: i for i, h in enumerate(hdr2)}
rows2 = {cells(l)[0]: cells(l) for l in t2[1:] if l.strip()}

def ate_parse(cell):
    m = re.match(r"^(-?)\$?(-?[\d,]+(?:\.\d+)?)\s*(pp|weeks|wks)?;\s*p\s*([=<])\s*([\d.]+);\s*95% CI:\s*\[(-?[\d,.]+),\s*(-?[\d,.]+)\]$", cell)
    assert m, cell
    sign = -1 if m.group(1) == "-" else 1
    f = lambda x: float(x.replace(",", ""))
    v = sign * f(m.group(2))
    return v, f(m.group(6)), f(m.group(7))

ATE = "Average treatment effect estimate"
for r in ["Wage", "Reemployment", "Weeks on UI"]:
    print(f"[Table 2] matched: {' | '.join(rows2[r])!r}")
w, wlo, whi = ate_parse(rows2["Wage"][col2[ATE]])
re_ate, _, _ = ate_parse(rows2["Reemployment"][col2[ATE]])
wk, wklo, wkhi = ate_parse(rows2["Weeks on UI"][col2[ATE]])
cm = re.match(r"^\$([\d,]+)", rows2["Wage"][col2["Control"]])
ri_control_mean_wage = int(cm.group(1).replace(",", ""))
toint = lambda x: int(x) if float(x).is_integer() else x
se_reported = bool(re.search(r"\b(SE|s\.e\.|standard error)\b", rows2["Wage"][col2[ATE]], re.I))
se_derived = round((whi - wlo) / (2 * Z), 2)

loc = lambda src, t, r, c: f'{src} | Table {t} | row "{r}" | column "{c}"'
res = {
    "task_id": "F4-conf-01",
    "dcr_att_5km": dcr_att_5km, "dcr_att_5km_se": dcr_att_5km_se,
    "dcr_att_3km": dcr_att_3km, "dcr_att_3km_se": dcr_att_3km_se,
    "dcr_obs_5km": dcr_obs_5km,
    "dcr_narcotics_att_3km": narc3, "dcr_narcotics_att_3km_se": narc3_se,
    "dcr_fatal_att_1km": fat1, "dcr_fatal_att_1km_se": fat1_se,
    "dcr_hiv_att_5km": hiv5, "dcr_hiv_att_5km_se": hiv5_se,
    "dcr_hiv_att_5km_p_value": hiv_p,
    "dcr_hiv_att_5km_significant_10pct": hiv5_sig10,
    "dcr_att_5km_ci_low": ci_lo, "dcr_att_5km_ci_high": ci_hi,
    "ri_wage_ate": toint(w), "ri_wage_ate_ci_low": toint(wlo), "ri_wage_ate_ci_high": toint(whi),
    "ri_wage_ate_se_reported": se_reported,
    "ri_wage_ate_se_derived": se_derived,
    "ri_control_mean_wage": ri_control_mean_wage,
    "ri_reemployment_ate_pp": re_ate,
    "ri_weeks_ui_ate": wk, "ri_weeks_ui_ate_ci_low": wklo, "ri_weeks_ui_ate_ci_high": wkhi,
    "locators": {
        "dcr_att_5km": loc(DCR_ID, 6, "ATT", "5 km"),
        "dcr_att_5km_se": loc(DCR_ID, 6, "ATT (standard error line, in parentheses)", "5 km"),
        "dcr_att_3km": loc(DCR_ID, 6, "ATT", "3 km"),
        "dcr_att_3km_se": loc(DCR_ID, 6, "ATT (standard error line, in parentheses)", "3 km"),
        "dcr_obs_5km": loc(DCR_ID, 6, "Observations", "5 km"),
        "dcr_narcotics_att_3km": loc(DCR_ID, 7, "Narcotics", "3 km"),
        "dcr_narcotics_att_3km_se": loc(DCR_ID, 7, "Narcotics (standard error line, in parentheses)", "3 km"),
        "dcr_fatal_att_1km": loc(DCR_ID, 7, "Drug-related (fatal)", "1 km"),
        "dcr_fatal_att_1km_se": loc(DCR_ID, 7, "Drug-related (fatal) (standard error line, in parentheses)", "1 km"),
        "dcr_hiv_att_5km": loc(DCR_ID, 7, "HIV", "5 km"),
        "dcr_hiv_att_5km_se": loc(DCR_ID, 7, "HIV (standard error line, in parentheses)", "5 km"),
        "dcr_hiv_att_5km_p_value": f"not reported in {DCR_ID} (Table 7 prints only significance stars; no exact p-value; deriving one from estimate/SE would require a distributional assumption)",
        "dcr_att_5km_ci_low": loc(DCR_ID, 6, "ATT", "5 km") + " | derived (estimate - 1.96*SE)",
        "dcr_att_5km_ci_high": loc(DCR_ID, 6, "ATT", "5 km") + " | derived (estimate + 1.96*SE)",
        "ri_wage_ate": loc(RI_ID, 2, "Wage", ATE),
        "ri_wage_ate_ci_low": loc(RI_ID, 2, "Wage", ATE) + " (95% CI lower bound)",
        "ri_wage_ate_ci_high": loc(RI_ID, 2, "Wage", ATE) + " (95% CI upper bound)",
        "ri_wage_ate_se_derived": loc(RI_ID, 2, "Wage", ATE) + " | derived ((upper - lower)/(2*1.96)); SE not reported in " + RI_ID,
        "ri_control_mean_wage": loc(RI_ID, 2, "Wage", "Control"),
        "ri_reemployment_ate_pp": loc(RI_ID, 2, "Reemployment", ATE),
        "ri_weeks_ui_ate": loc(RI_ID, 2, "Weeks on UI", ATE),
        "ri_weeks_ui_ate_ci_low": loc(RI_ID, 2, "Weeks on UI", ATE) + " (95% CI lower bound)",
        "ri_weeks_ui_ate_ci_high": loc(RI_ID, 2, "Weeks on UI", ATE) + " (95% CI upper bound)",
    },
    "derivation_notes": {
        "dcr_att_5km_ci": "Normal approximation with critical value 1.96: CI = ATT +/- 1.96*SE using the Table 6 5 km ATT (-0.431) and its printed clustered SE (0.210); rounded to 4 decimals. Assumes asymptotic normality of the estimator; the paper's own CIs (figures) may use a different (e.g. multiplier-bootstrap) critical value.",
        "ri_wage_ate_se": "SE = (upper - lower)/(2*1.96) from the printed Table 2 95% CI [166, 2,141], assuming a symmetric normal-based (Wald) interval with critical value 1.96; rounded to 2 decimals. The printed CI is nearly symmetric around 1,153 (987 below, 988 above), consistent with this assumption; with N=23,549 a t critical value is indistinguishable from 1.96.",
    },
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": [
        {"issue": "DCR: the 5 km CS-DR ATT -0.431 has SE (0.210) in Table 6 but SE (0.218) in Appendix Table A19 (same estimator, buffer and N=17,250).",
         "handling": "Used Table 6 as the key definitions require; derived CI uses 0.210. Flagged as an internal inconsistency."},
        {"issue": "DCR: control group is 608 never-treated regions in Tables 6/7 and A19 rows, but the A19 note says 640 never-treated comparison regions.",
         "handling": "Not used for any key; flagged."},
        {"issue": "DCR: abstract says fatal outcomes 'fall by 35%', but Table 7 5 km fatal ATT is -0.354 log points (~-29.8%); Section 6.1 itself says ~30% at 5 km. The abstract appears to read log points as percent.",
         "handling": "Not used for any key; values taken from Table 7; flagged."},
        {"issue": "DCR: text describes HIV estimates as 'close to zero', while the Table 7 HIV 5 km ATT is -0.340 (SE 0.294), similar in size to other effects; it is insignificant (no star) but not close to zero.",
         "handling": "Reported printed values; significance taken from stars (none)."},
        {"issue": "RESEA: the text reports the weeks-on-UI CI as magnitudes '[1.75 weeks, 2.24 weeks]' whereas Table 2 prints [-2.24, -1.75].",
         "handling": "Used Table 2 values with printed signs, lower bound first."},
        {"issue": "RESEA: rendering defect 'p ¡ 0.001' in the short-term weeks-on-UI paragraph (should be 'p < 0.001').",
         "handling": "Not used for any key; noted."},
        {"issue": "RESEA: Tables 2 and 3 were re-set by the extract editor (one row per outcome).",
         "handling": "Accepted as verbatim per extract header; Table 2 values cross-checked equal to Table 3 'Difference in means' column."},
    ],
    "n_input_rows": {},
    "n_analysed": None,
    "random_seed": None,
}
# cross-check Table 3 first column equals Table 2 ATE column
t3 = section(ri, "Outcome | Difference in means", "Table 3:")
r3 = {cells(l)[0]: cells(l) for l in t3[1:] if l.strip()}
for r in ["Wage", "Reemployment", "Weeks on UI"]:
    a = ate_parse(rows2[r][col2[ATE]]); b = ate_parse(r3[r][1].replace("wks", "weeks"))
    print(f"[check] Table2 vs Table3 {r}: {a} vs {b} -> {'equal' if a == b else 'DIFFERENT'}")
# Table 6 vs A19 SE check
a19 = section(dcr, "Table A19:", "Note:")
a_att, a_se = row_with_se(a19, "ATT", "Table A19")
print(f"[check] Table 6 5 km SE {dcr_att_5km_se} vs A19 CS-DR SE {num(a_se[0])[0]}")

for f in ["README.md", "sources/arxiv-2609.27732v1.txt", "sources/arxiv-2606.14621v1.txt"]:
    with open(os.path.join("data", f), "rb") as fh:
        res["n_input_rows"][f] = fh.read().count(b"\n")
scalar_keys = [k for k in res if k.startswith(("dcr_", "ri_"))]
res["n_analysed"] = len(scalar_keys)
res["printed_strings"] = {
    "dcr_att_5km": att[col6["5 km"]], "dcr_att_5km_se": se[col6["5 km"]],
    "dcr_att_3km": att[col6["3 km"]], "dcr_att_3km_se": se[col6["3 km"]],
    "dcr_obs_5km": obs[col6["5 km"]],
    "table7_narcotics": " | ".join(cells([l for l in t7 if l.startswith("Narcotics")][0])),
    "table7_fatal": " | ".join(cells([l for l in t7 if l.startswith("Drug-related (fatal)")][0])),
    "table7_hiv": " | ".join(cells([l for l in t7 if l.startswith("HIV")][0])),
    "ri_wage_ate_cell": rows2["Wage"][col2[ATE]], "ri_wage_control_cell": rows2["Wage"][col2["Control"]],
    "ri_reemployment_ate_cell": rows2["Reemployment"][col2[ATE]],
    "ri_weeks_ui_ate_cell": rows2["Weeks on UI"][col2[ATE]],
    "note": "JSON numbers cannot carry trailing zeros (e.g. printed 0.210 is stored as 0.21); the printed strings are kept here verbatim.",
}
os.makedirs("out", exist_ok=True)
json.dump(res, open("out/results.json", "w"), indent=2, ensure_ascii=False)
print(json.dumps({k: res[k] for k in scalar_keys}, indent=1))