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

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DCR_ID, RI_ID = "arxiv-2609.27732v1", "arxiv-2606.14621v1"
DCR_F = os.path.join(ROOT, "data", "sources", DCR_ID + ".txt")
RI_F = os.path.join(ROOT, "data", "sources", RI_ID + ".txt")
README_F = os.path.join(ROOT, "data", "README.md")
Z = 1.96  # normal critical value for two-sided 95%

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

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

def num(s):
    """Parse a printed number string (keep sign; strip $, commas, stars, %, units)."""
    s = s.strip().replace("$", "").replace(",", "")
    m = re.search(r"-?\d+(?:\.\d+)?", s)
    return m.group(0)

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

def table_block(text, title_prefix):
    lines = text.splitlines()
    start = next(i for i, l in enumerate(lines) if l.startswith(title_prefix))
    block = []
    for l in lines[start + 1:]:
        if l.startswith("Note:"):
            break
        if l.strip():
            block.append(l)
    return block

def parse_reg_table(block):
    """Return header cols, {row_label: {col: (coef_str, se_str or None, raw lines)}}."""
    header = [c.strip() for c in block[0].split("|")]
    cols = [c for c in header if c and c != "Outcome"]
    rows, last = {}, None
    for l in block[1:]:
        cells = [c.strip() for c in l.split("|")]
        if cells[0] == "" and last is not None and all(c.startswith("(") for c in cells[1:] if c):
            for c, v in zip(cols, cells[1:]):
                coef, _, raw = rows[last][c]
                rows[last][c] = (coef, v, raw + [l])
            continue
        label = cells[0]
        rows[label] = {c: (v, None, [l]) for c, v in zip(cols, cells[1:])}
        last = label
    return cols, rows

dcr = read(DCR_F); ri = read(RI_F)
t6_cols, t6 = parse_reg_table(table_block(dcr, "Table 6:"))
t7_cols, t7 = parse_reg_table(table_block(dcr, "Table 7:"))

R, LOC, PRINTED, MATCHED = {}, {}, {}, {}

def loc(src, tab, row, col, extra=""):
    return f'{src} | Table {tab} | row "{row}" | column "{col}"' + extra

def put_dcr(key, tab, rows, row, col, what):
    coef, se, raw = rows[row][col]
    s = coef if what == "coef" else se
    p = num(s)
    R[key] = to_val(p); PRINTED[key] = s
    LOC[key] = loc(DCR_ID, tab, row, col, "" if what == "coef" else " (standard error in parentheses on the line below)")
    MATCHED[key] = raw
    return coef

put_dcr("dcr_att_5km", 6, t6, "ATT", "5 km", "coef")
put_dcr("dcr_att_5km_se", 6, t6, "ATT", "5 km", "se")
put_dcr("dcr_att_3km", 6, t6, "ATT", "3 km", "coef")
put_dcr("dcr_att_3km_se", 6, t6, "ATT", "3 km", "se")
put_dcr("dcr_obs_5km", 6, t6, "Observations", "5 km", "coef")
put_dcr("dcr_narcotics_att_3km", 7, t7, "Narcotics", "3 km", "coef")
put_dcr("dcr_narcotics_att_3km_se", 7, t7, "Narcotics", "3 km", "se")
put_dcr("dcr_fatal_att_1km", 7, t7, "Drug-related (fatal)", "1 km", "coef")
put_dcr("dcr_fatal_att_1km_se", 7, t7, "Drug-related (fatal)", "1 km", "se")
hiv_coef = put_dcr("dcr_hiv_att_5km", 7, t7, "HIV", "5 km", "coef")
put_dcr("dcr_hiv_att_5km_se", 7, t7, "HIV", "5 km", "se")

# HIV exact p-value: search the extract for any printed p-value (p= / p-value) -- none expected
pv_hits = [l for l in dcr.splitlines() if re.search(r"p\s*=\s*0?\.\d|p-value", l, re.I)]
R["dcr_hiv_att_5km_p_value"] = None
LOC["dcr_hiv_att_5km_p_value"] = (f"not reported in {DCR_ID} (Table 7 prints only the coefficient, SE and "
                                  f"significance stars; {len(pv_hits)} exact p-value strings found in the extract)")
stars = len(hiv_coef) - len(hiv_coef.rstrip("*"))
R["dcr_hiv_att_5km_significant_10pct"] = stars >= 1  # * p<0.10 or stronger

# Derived 95% CI from Table 6 printed values
b, se = R["dcr_att_5km"], R["dcr_att_5km_se"]
R["dcr_att_5km_ci_low"] = round(b - Z * se, 6)
R["dcr_att_5km_ci_high"] = round(b + Z * se, 6)
for k in ("dcr_att_5km_ci_low", "dcr_att_5km_ci_high"):
    LOC[k] = (f'{DCR_ID} | Table 6 | row "ATT" | column "5 km" | derived from printed ATT and SE')

# RESEA Table 2
lines = ri.splitlines()
hdr_i = next(i for i, l in enumerate(lines) if l.startswith("Outcome | Treatment | Control"))
t2_cols = [c.strip() for c in lines[hdr_i].split("|")]
t2 = {}
for l in lines[hdr_i + 1:]:
    if not l.strip():
        break
    cells = [c.strip() for c in l.split("|")]
    t2[cells[0]] = dict(zip(t2_cols[1:], cells[1:])) | {"_raw": l}
ATE = "Average treatment effect estimate"

def parse_effect(cell):
    est = cell.split(";")[0]
    m = re.search(r"95% CI:\s*\[([^\]]+)\]", cell)
    lo, hi = [x.strip() for x in m.group(1).split(", ")]
    return est, lo, hi

def put_ri(key, row, col, s, part):
    R[key] = to_val(num(s)); PRINTED[key] = s
    LOC[key] = loc(RI_ID, 2, row, col, part); MATCHED[key] = [t2[row]["_raw"]]

w_est, w_lo, w_hi = parse_effect(t2["Wage"][ATE])
put_ri("ri_wage_ate", "Wage", ATE, w_est, " (point estimate)")
put_ri("ri_wage_ate_ci_low", "Wage", ATE, w_lo, " (95% CI lower bound)")
put_ri("ri_wage_ate_ci_high", "Wage", ATE, w_hi, " (95% CI upper bound)")
# Search the wage row and the paper body (after the editor's header) for a printed standard error.
ri_body = ri.split("=" * 20, 1)[1]
wage_row = t2["Wage"]["_raw"]
R["ri_wage_ate_se_reported"] = bool(re.search(r"\bSE\b|s\.e\.|standard error", wage_row, re.I)
                                    or re.search(r"standard error|\bSE\b", ri_body, re.I))
R["ri_wage_ate_se_derived"] = round((R["ri_wage_ate_ci_high"] - R["ri_wage_ate_ci_low"]) / (2 * Z), 6)
LOC["ri_wage_ate_se_derived"] = (f'{RI_ID} | Table 2 | row "Wage" | column "{ATE}" | derived from printed 95% CI bounds')
ctrl = t2["Wage"]["Control"].split("(")[0].strip()
put_ri("ri_control_mean_wage", "Wage", "Control", ctrl, " (mean)")
r_est, _, _ = parse_effect(t2["Reemployment"][ATE])
put_ri("ri_reemployment_ate_pp", "Reemployment", ATE, r_est, " (point estimate)")
u_est, u_lo, u_hi = parse_effect(t2["Weeks on UI"][ATE])
put_ri("ri_weeks_ui_ate", "Weeks on UI", ATE, u_est, " (point estimate)")
put_ri("ri_weeks_ui_ate_ci_low", "Weeks on UI", ATE, u_lo, " (95% CI lower bound)")
put_ri("ri_weeks_ui_ate_ci_high", "Weeks on UI", ATE, u_hi, " (95% CI upper bound)")

# Cross-table check values used only for data_issues
a19_cols, a19 = parse_reg_table(table_block(dcr, "Table A19:"))
a19_se = a19["ATT"]["CS-DR"][1]

derivation_notes = {
    "dcr_att_5km_ci": (f"Normal approximation, critical value 1.96: CI = ATT +/- 1.96*SE using Table 6 printed "
                       f"ATT {PRINTED['dcr_att_5km']} and SE {PRINTED['dcr_att_5km_se']} (clustered at MedStat region level); "
                       f"bounds rounded to 6 decimals. Not printed in the source."),
    "ri_wage_ate_se": (f"SE = (upper - lower)/(2*1.96) from Table 2 printed 95% CI [{w_lo}, {w_hi}], assuming a symmetric "
                       f"normal-approximation interval (the CI method is 'described in the main text', which is not in the extract); "
                       f"rounded to 6 decimals. The source prints no standard error."),
}

data_issues = [
    {"issue": f"{DCR_ID}: the 5 km CS-DR ATT -0.431 has SE {PRINTED['dcr_att_5km_se']} in Table 6 but {a19_se} in Table A19 (same estimator, buffer and window).",
     "handling": "Keys name Table 6, so the Table 6 SE (0.210) is used for dcr_att_5km_se and the derived CI; discrepancy reported."},
    {"issue": f"{DCR_ID}: Table A19 note states 640 never-treated comparison regions whereas Tables 6/7 and the Figure 6 note state 608.",
     "handling": "No requested key depends on it; reported only."},
    {"issue": f"{DCR_ID}: abstract says fatal drug-related hospitalisations fall by 35%, Section 6.1 says approximately 30% at 5 km; Table 7 5 km fatal ATT is -0.354 (exp(-0.354)-1 = -29.8%).",
     "handling": "No requested key depends on the prose; table values used; reported only."},
    {"issue": f"{DCR_ID}: text calls HIV estimates 'insignificant and close to zero' while the Table 7 5 km HIV ATT is -0.340 (SE 0.294), similar in size to significant estimates.",
     "handling": "Significance flag taken from the table stars (none); exact p-value not printed -> null."},
    {"issue": f"{DCR_ID}: Figure 7 note and Table 7 note describe four sub-categories while Table 7 shows five outcome rows (fatal has no ICD code listed).",
     "handling": "Reported only."},
    {"issue": f"{RI_ID}: prose gives the weeks-on-UI CI as positive magnitudes [1.75 weeks, 2.24 weeks] while Table 2 prints [-2.24, -1.75].",
     "handling": "Keys use Table 2 as printed (negative, lower bound first)."},
    {"issue": f"{RI_ID}: rendering defect 'p ¡ 0.001' in the short-term results paragraph (intended 'p < 0.001').",
     "handling": "No requested key affected; reported only."},
    {"issue": f"{RI_ID}: Tables 2 and 3 were re-set by the extract editor (one row per outcome); Table 3 column 1 duplicates Table 2's ATE column.",
     "handling": "Values taken from Table 2 as the keys require; identical to Table 3 'Difference in means'."},
]

value_keys = [k for k in R]
results = dict(R)
results["locators"] = {k: LOC[k] for k in R if k not in ("dcr_hiv_att_5km_significant_10pct", "ri_wage_ate_se_reported")}
results["printed_strings"] = PRINTED
results["derivation_notes"] = derivation_notes
results["task_id"] = "F4-conf-01"
results["feasibility"] = "complete"
results["feasibility_reason"] = None
results["data_issues"] = data_issues
results["n_input_rows"] = {"README.md": wc_l(README_F), DCR_ID + ".txt": wc_l(DCR_F), RI_ID + ".txt": wc_l(RI_F)}
results["n_analysed"] = len(value_keys)
results["random_seed"] = None

os.makedirs(os.path.join(ROOT, "out"), exist_ok=True)
with open(os.path.join(ROOT, "out", "results.json"), "w", encoding="utf-8") as f:
    json.dump(results, f, indent=2, ensure_ascii=False)

print("Matched source lines:")
for k, raw in MATCHED.items():
    print(f"  {k} = {R[k]!r}  <- {PRINTED[k]!r}  from: {' // '.join(r.strip() for r in raw)}")
for k in ("dcr_hiv_att_5km_p_value", "dcr_hiv_att_5km_significant_10pct", "dcr_att_5km_ci_low",
          "dcr_att_5km_ci_high", "ri_wage_ate_se_reported", "ri_wage_ate_se_derived"):
    print(f"  {k} = {R[k]!r}")
print(f"n_analysed = {len(value_keys)}; wrote out/results.json")