"""F4-conf-01: extract printed values from two arXiv extracts; derive two quantities.
Reads only data/; writes out/results.json deterministically. No randomness."""
import json, re, os, hashlib

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 two-sided 95% intervals


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


def wc_l(text):
    return text.count("\n")


readme = read("README.md")
dcr = read(f"sources/{DCR_ID}.txt")
ri = read(f"sources/{RI_ID}.txt")
n_input_rows = {"README.md": wc_l(readme), f"sources/{DCR_ID}.txt": wc_l(dcr),
                f"sources/{RI_ID}.txt": wc_l(ri)}

# verify SHA-256 against README provenance table
for name, txt in [(DCR_ID, dcr), (RI_ID, ri)]:
    h = hashlib.sha256(txt.encode("utf-8")).hexdigest()
    assert h in readme, f"SHA mismatch for {name}"
    print(f"[sha256 ok] {name} {h}")

dcr_lines = [l for l in dcr.split("\n") if l.strip()]


def table_block(lines, start_pat, end_pat):
    i = next(k for k, l in enumerate(lines) if l.startswith(start_pat))
    j = next(k for k, l in enumerate(lines) if k > i and l.startswith(end_pat))
    return lines[i:j]


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


def num(s):
    """Parse a printed number: strip stars, parentheses, $, commas, %; keep sign."""
    t = s.replace("$", "").replace(",", "").replace("*", "").strip("() %")
    return float(t) if "." in t else int(t)  # integers printed without decimals stay integers


def coef_and_se(block, row_label, col_label, header_idx=0):
    header = cells(block[header_idx])
    col = header.index(col_label)
    k = next(k for k, l in enumerate(block) if cells(l)[0] == row_label)
    row, se_row = cells(block[k]), cells(block[k + 1])
    assert se_row[0] == "" and se_row[col].startswith("("), se_row
    print(f"[match] row: {block[k]!r}\n        se : {block[k+1]!r}  -> col {col_label!r}: {row[col]} / {se_row[col]}")
    return row[col], se_row[col]


# ---- DCR Table 6
t6 = table_block(dcr_lines, "Table 6:", "Note:")[1:]
att5_s, se5_s = coef_and_se(t6, "ATT", "5 km")
att3_s, se3_s = coef_and_se(t6, "ATT", "3 km")
hdr6 = cells(t6[0])
obs_line = next(l for l in t6 if cells(l)[0] == "Observations")
print(f"[match] {obs_line!r}")
obs5 = int(cells(obs_line)[hdr6.index("5 km")].replace(",", ""))

# ---- DCR Table 7
t7 = table_block(dcr_lines, "Table 7:", "Note:")[1:]
narc_s, narc_se_s = coef_and_se(t7, "Narcotics", "3 km")
fatal_s, fatal_se_s = coef_and_se(t7, "Drug-related (fatal)", "1 km")
hiv_s, hiv_se_s = coef_and_se(t7, "HIV", "5 km")
hiv_sig10 = "*" in hiv_s  # any star = p<0.10 per table note
# exact p-value: does the extract print one for Table 7 cells?
t7_has_p = any(re.search(r"p\s*=\s*0?\.\d", l) for l in t7)
hiv_p = None
print(f"[check] Table 7 prints exact p-values: {t7_has_p}")
assert not t7_has_p

# ---- Table A19 (only used for consistency check)
ta19 = table_block(dcr_lines, "Table A19:", "Note:")[1:]
a19_att, a19_se = coef_and_se(ta19, "ATT", "CS-DR")

# ---- RI Table 2
ri_lines = ri.split("\n")
hdr_idx = next(k for k, l in enumerate(ri_lines) if l.startswith("Outcome | Treatment | Control"))
hdr2 = cells(ri_lines[hdr_idx])
t2 = {}
for l in ri_lines[hdr_idx + 1:]:
    if not l.strip():
        break
    c = cells(l)
    t2[c[0]] = dict(zip(hdr2, c))
ATE = "Average treatment effect estimate"


def parse_effect(cell):
    m = re.match(r"^(-?\$?-?[\d,]*\d(?:\.\d+)?)\s*(?:pp|weeks|wks)?;\s*p\s*([=<])\s*([\d.]+);\s*95% CI:\s*\[(-?[\d.]+(?:,\d{3})*),\s*(-?[\d.]+(?:,\d{3})*)\]$", cell)
    assert m, cell
    return m.group(1), m.group(2) + m.group(3), m.group(4), m.group(5)


wage_cell = t2["Wage"][ATE]
reemp_cell = t2["Reemployment"][ATE]
weeks_cell = t2["Weeks on UI"][ATE]
ctrl_cell = t2["Wage"]["Control"]
for lab in ["Wage", "Reemployment", "Weeks on UI"]:
    print(f"[match] Table 2 {lab!r} ATE cell: {t2[lab][ATE]!r}")
print(f"[match] Table 2 'Wage' Control cell: {ctrl_cell!r}")
w_pt, w_p, w_lo, w_hi = parse_effect(wage_cell)
r_pt, r_p, r_lo, r_hi = parse_effect(reemp_cell)
u_pt, u_p, u_lo, u_hi = parse_effect(weeks_cell)
ctrl_mean = num(re.match(r"^(\$[\d,]+)", ctrl_cell).group(1))
wage_se_reported = bool(re.search(r"\bSE\b|standard error|\(\s*[\d,.]+\s*\)", wage_cell))

# ---- derived
att5, se5 = num(att5_s), num(se5_s)
ci_lo = round(att5 - Z * se5, 4)
ci_hi = round(att5 + Z * se5, 4)
wlo, whi = num(w_lo), num(w_hi)
wage_se = round((whi - wlo) / (2 * Z), 2)

res = {
    "task_id": "F4-conf-01",
    "dcr_att_5km": att5, "dcr_att_5km_se": se5,
    "dcr_att_3km": num(att3_s), "dcr_att_3km_se": num(se3_s),
    "dcr_obs_5km": obs5,
    "dcr_narcotics_att_3km": num(narc_s), "dcr_narcotics_att_3km_se": num(narc_se_s),
    "dcr_fatal_att_1km": num(fatal_s), "dcr_fatal_att_1km_se": num(fatal_se_s),
    "dcr_hiv_att_5km": num(hiv_s), "dcr_hiv_att_5km_se": num(hiv_se_s),
    "dcr_hiv_att_5km_p_value": hiv_p,
    "dcr_hiv_att_5km_significant_10pct": hiv_sig10,
    "dcr_att_5km_ci_low": ci_lo, "dcr_att_5km_ci_high": ci_hi,
    "ri_wage_ate": num(w_pt), "ri_wage_ate_ci_low": wlo, "ri_wage_ate_ci_high": whi,
    "ri_wage_ate_se_reported": wage_se_reported,
    "ri_wage_ate_se_derived": wage_se,
    "ri_control_mean_wage": ctrl_mean,
    "ri_reemployment_ate_pp": num(r_pt),
    "ri_weeks_ui_ate": num(u_pt), "ri_weeks_ui_ate_ci_low": num(u_lo), "ri_weeks_ui_ate_ci_high": num(u_hi),
}

L6 = f"{DCR_ID} | Table 6"
L7 = f"{DCR_ID} | Table 7"
L2 = f"{RI_ID} | Table 2"
res["locators"] = {
    "dcr_att_5km": f'{L6} | row "ATT" | column "5 km"',
    "dcr_att_5km_se": f'{L6} | row "ATT" (standard-error line in parentheses) | column "5 km"',
    "dcr_att_3km": f'{L6} | row "ATT" | column "3 km"',
    "dcr_att_3km_se": f'{L6} | row "ATT" (standard-error line in parentheses) | column "3 km"',
    "dcr_obs_5km": f'{L6} | row "Observations" | column "5 km"',
    "dcr_narcotics_att_3km": f'{L7} | row "Narcotics" | column "3 km"',
    "dcr_narcotics_att_3km_se": f'{L7} | row "Narcotics" (standard-error line in parentheses) | column "3 km"',
    "dcr_fatal_att_1km": f'{L7} | row "Drug-related (fatal)" | column "1 km"',
    "dcr_fatal_att_1km_se": f'{L7} | row "Drug-related (fatal)" (standard-error line in parentheses) | column "1 km"',
    "dcr_hiv_att_5km": f'{L7} | row "HIV" | column "5 km"',
    "dcr_hiv_att_5km_se": f'{L7} | row "HIV" (standard-error line in parentheses) | column "5 km"',
    "dcr_hiv_att_5km_p_value": f"not reported in {DCR_ID} (Table 7 prints only significance stars, no exact p-values; an exact p-value could only be computed under an additional distributional assumption)",
    "dcr_att_5km_ci_low": f'{L6} | row "ATT" | column "5 km" | derived from printed ATT and SE (ATT - 1.96*SE)',
    "dcr_att_5km_ci_high": f'{L6} | row "ATT" | column "5 km" | derived from printed ATT and SE (ATT + 1.96*SE)',
    "ri_wage_ate": f'{L2} | row "Wage" | column "Average treatment effect estimate" (point estimate)',
    "ri_wage_ate_ci_low": f'{L2} | row "Wage" | column "Average treatment effect estimate" (95% CI lower bound)',
    "ri_wage_ate_ci_high": f'{L2} | row "Wage" | column "Average treatment effect estimate" (95% CI upper bound)',
    "ri_wage_ate_se_derived": f'{L2} | row "Wage" | column "Average treatment effect estimate" | derived from printed 95% CI ((high - low)/(2*1.96)); SE not reported in {RI_ID}',
    "ri_control_mean_wage": f'{L2} | row "Wage" | column "Control" (mean)',
    "ri_reemployment_ate_pp": f'{L2} | row "Reemployment" | column "Average treatment effect estimate" (point estimate)',
    "ri_weeks_ui_ate": f'{L2} | row "Weeks on UI" | column "Average treatment effect estimate" (point estimate)',
    "ri_weeks_ui_ate_ci_low": f'{L2} | row "Weeks on UI" | column "Average treatment effect estimate" (95% CI lower bound)',
    "ri_weeks_ui_ate_ci_high": f'{L2} | row "Weeks on UI" | column "Average treatment effect estimate" (95% CI upper bound)',
}
res["derivation_notes"] = {
    "dcr_att_5km_ci": "Normal (Wald) approximation with critical value 1.96: ATT +/- 1.96*SE using the Table 6 printed ATT -0.431 and clustered SE 0.210; bounds rounded to 4 decimals. Assumes approximate normality of the Callaway-Sant'Anna estimator; the source prints no CI for the overall ATT.",
    "ri_wage_ate_se": "SE = (upper - lower)/(2*1.96) from the Table 2 printed 95% CI [166, 2,141]; rounded to 2 decimals. Assumes the CI is a symmetric normal-approximation (Wald) interval; the printed CI is nearly symmetric about 1,153 (987 below, 988 above). With N = 23,549 a t critical value is practically identical to 1.96.",
}
# consistency checks -> data issues
issues = []
if (a19_att.replace("*", ""), a19_se) != (att5_s.replace("*", ""), se5_s):
    issues.append({"issue": f"Internal inconsistency in {DCR_ID}: the 5 km CS-DR ATT {att5_s} has SE {se5_s} in Table 6 but SE {a19_se} in Appendix Table A19 (same estimator, buffer and N=17,250).",
                   "handling": "Used Table 6 as required by the key definitions; the derived CI uses the Table 6 SE 0.210. Flagged; not resolvable from the extract."})
if "640 regions" in dcr:
    issues.append({"issue": f"Internal inconsistency in {DCR_ID}: Tables 6/7 and the Figure 6 note report 608 never-treated control regions, but the Table A19 note says 640 never-treated comparison regions.",
                   "handling": "No requested key depends on it; flagged only."})
if "fall by 35%" in dcr and "approximately 30% at 5 km" in dcr:
    issues.append({"issue": f"Text inconsistency in {DCR_ID}: the abstract says fatal hospitalisations fall by 35%, Section 6.1 says ~30% at 5 km (Table 7 fatal 5 km ATT -0.354 implies exp(-0.354)-1 = -29.8%). Section 6 text also calls HIV estimates 'close to zero' although the 5 km point estimate is -0.340.",
                   "handling": "Narrative only; all keys taken from the tables. Flagged."})
if "[1.75 weeks, 2.24 weeks]" in ri:
    issues.append({"issue": f"Sign convention differs in {RI_ID}: the text reports the weeks-on-UI CI as magnitudes of a reduction [1.75, 2.24], while Table 2 prints the ATE -1.99 with CI [-2.24, -1.75].",
                   "handling": "Used the Table 2 signed values as printed (lower bound first)."})
if "p ¡ 0.001" in ri:
    issues.append({"issue": f"Rendering defect in {RI_ID} text: 'p ¡ 0.001' (should read p < 0.001).",
                   "handling": "Does not affect any key; flagged."})
t3w = next(l for l in ri_lines if l.startswith("Wage | $1,153"))
assert cells(t3w)[1] == wage_cell, "Table 3 / Table 2 wage mismatch"
issues.append({"issue": f"{RI_ID} Table 2 was re-set by the extract editor (one row per outcome; p-values reattached to columns), per the extract header.",
               "handling": "Accepted the editor's layout; the Table 3 'Difference in means' wage cell duplicates the Table 2 wage ATE cell exactly (checked in code), consistent with correct reattachment."})
res["data_issues"] = issues
res["feasibility"] = "complete"
res["feasibility_reason"] = None
res["n_input_rows"] = n_input_rows
value_keys = [k for k in res if k.startswith(("dcr_", "ri_"))]
res["n_analysed"] = len(value_keys)
res["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: res[k] for k in value_keys}, indent=1))