#!/usr/bin/env python
"""F4-conf-01: extract printed treatment effects/uncertainty from two source extracts.

Reads only data/; writes out/results.json deterministically. No randomness (random_seed = null).
Values are parsed from the text (not hard-coded); matched source lines are printed.
"""
import json, math, os, re, hashlib
from statistics import NormalDist

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


def read_lines(p):
    with open(p, encoding="utf-8") as fh:
        return fh.read().split("\n")


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


def num(s):
    """Parse a printed number: strip $, commas, stars, %, units; keep sign."""
    s = s.strip().replace("$", "").replace(",", "").replace("*", "")
    m = re.search(r"-?\d+(?:\.\d+)?", s)
    if not m:
        raise ValueError(f"no number in {s!r}")
    t = m.group(0)
    return float(t) if "." in t else int(t)  # keep integers printed without decimals as integers


def stars(s):
    return len(re.search(r"\**$", s.strip()).group(0))


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


# ---------------- DCR tables ----------------
def dcr_table(lines, title_prefix):
    """Return (header, rows) for a DCR table; rows: label -> {col: (cell, se_cell, lineno, se_lineno)}."""
    start = next(i for i, l in enumerate(lines) if l.startswith(title_prefix))
    body = []
    for i in range(start + 1, len(lines)):
        l = lines[i]
        if l.startswith("Note:"):
            break
        if l.strip():
            body.append((i + 1, l))
    header = cells(body[0][1])
    cols = header[1:]
    rows, last = {}, None
    for ln, l in body[1:]:
        c = cells(l)
        if c[0] == "" and last is not None and l.strip().startswith("|") and "(" in l:
            for j, col in enumerate(cols):
                v = c[j + 1] if j + 1 < len(c) else ""
                if v:
                    cell, _, cln, _ = rows[last][col]
                    rows[last][col] = (cell, v, cln, ln)
            continue
        last = c[0]
        rows[last] = {col: ((c[j + 1] if j + 1 < len(c) else ""), None, ln, None) for j, col in enumerate(cols)}
    return cols, rows


dcr = read_lines(DCR_F)
t6_cols, t6 = dcr_table(dcr, "Table 6:")
t7_cols, t7 = dcr_table(dcr, "Table 7:")
ta19_cols, ta19 = dcr_table(dcr, "Table A19:")

matched = []


def grab(tab, tname, row, col, src=DCR_ID, lines=dcr):
    cell, se, ln, sln = tab[row][col]
    matched.append(f"{src} L{ln}: {lines[ln-1]}")
    if sln:
        matched.append(f"{src} L{sln}: {lines[sln-1]}")
    return cell, se


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


res, locs = {}, {}
c, se = grab(t6, "6", "ATT", "5 km")
res["dcr_att_5km"], res["dcr_att_5km_se"] = num(c), num(se)
locs["dcr_att_5km"] = loc(DCR_ID, 6, "ATT", "5 km")
locs["dcr_att_5km_se"] = loc(DCR_ID, 6, "ATT", "5 km", "standard error in parentheses on the line below")
c, se = grab(t6, "6", "ATT", "3 km")
res["dcr_att_3km"], res["dcr_att_3km_se"] = num(c), num(se)
locs["dcr_att_3km"] = loc(DCR_ID, 6, "ATT", "3 km")
locs["dcr_att_3km_se"] = loc(DCR_ID, 6, "ATT", "3 km", "standard error in parentheses on the line below")
c, _ = grab(t6, "6", "Observations", "5 km")
res["dcr_obs_5km"] = int(num(c))
locs["dcr_obs_5km"] = loc(DCR_ID, 6, "Observations", "5 km")
for key, row, col in [("dcr_narcotics_att_3km", "Narcotics", "3 km"),
                      ("dcr_fatal_att_1km", "Drug-related (fatal)", "1 km"),
                      ("dcr_hiv_att_5km", "HIV", "5 km")]:
    c, se = grab(t7, "7", row, col)
    res[key], res[key + "_se"] = num(c), num(se)
    locs[key] = loc(DCR_ID, 7, row, col)
    locs[key + "_se"] = loc(DCR_ID, 7, row, col, "standard error in parentheses on the line below")
    if key == "dcr_hiv_att_5km":
        hiv_cell = c

# exact p-value: Table 7 prints only significance stars, no p-values
t7_start = next(i for i, l in enumerate(dcr) if l.startswith("Table 7:"))
t7_end = next(i for i in range(t7_start, len(dcr)) if dcr[i].startswith("Figure 7 and Table 7"))
pval_printed = any(re.search(r"\bp\s*=\s*0?\.\d", l) for l in dcr[t7_start:t7_end])
res["dcr_hiv_att_5km_p_value"] = None
locs["dcr_hiv_att_5km_p_value"] = (f"not reported in {DCR_ID}: Table 7, row \"HIV\", column \"5 km\" prints only the "
                                   "coefficient, SE and significance stars (none); no exact p-value is printed")
assert not pval_printed
res["dcr_hiv_att_5km_significant_10pct"] = stars(hiv_cell) >= 1

# derived 95% CI from printed Table 6 SE (normal, z = 1.96)
b, s = res["dcr_att_5km"], res["dcr_att_5km_se"]
res["dcr_att_5km_ci_low"] = round(b - Z * s, 4)
res["dcr_att_5km_ci_high"] = round(b + Z * s, 4)
locs["dcr_att_5km_ci_low"] = loc(DCR_ID, 6, "ATT", "5 km", "derived: estimate - 1.96 x printed SE")
locs["dcr_att_5km_ci_high"] = loc(DCR_ID, 6, "ATT", "5 km", "derived: estimate + 1.96 x printed SE")

# ---------------- RESEA Table 2 ----------------
ri = read_lines(RI_F)


def ri_table(lines, marker):
    start = next(i for i, l in enumerate(lines) if l.startswith(marker))
    hdr_i = next(i for i in range(start + 1, len(lines)) if lines[i].startswith("Outcome |"))
    cols = cells(lines[hdr_i])[1:]
    rows = {}
    for i in range(hdr_i + 1, len(lines)):
        if not lines[i].strip():
            break
        c = cells(lines[i])
        rows[c[0]] = {col: (c[j + 1], i + 1) for j, col in enumerate(cols)}
    return cols, rows


def parse_effect(cell):
    parts = [p.strip() for p in cell.split(";")]
    est = num(parts[0])
    p = parts[1]
    ci = re.search(r"\[(.+?),\s*(-?[\d,\.]+)\]", cell)
    lo, hi = num(ci.group(1)), num(ci.group(2))
    return est, p, lo, hi


t2_cols, t2 = ri_table(ri, "[Table 2")
ATE = "Average treatment effect estimate"
c, ln = t2["Wage"][ATE]
matched.append(f"{RI_ID} L{ln}: {ri[ln-1]}")
w_est, w_p, w_lo, w_hi = parse_effect(c)
res["ri_wage_ate"], res["ri_wage_ate_ci_low"], res["ri_wage_ate_ci_high"] = w_est, w_lo, w_hi
locs["ri_wage_ate"] = loc(RI_ID, 2, "Wage", ATE, "point estimate")
locs["ri_wage_ate_ci_low"] = loc(RI_ID, 2, "Wage", ATE, "95% CI lower bound")
locs["ri_wage_ate_ci_high"] = loc(RI_ID, 2, "Wage", ATE, "95% CI upper bound")
res["ri_wage_ate_se_reported"] = bool(re.search(r"\bSE\b|standard error", c, re.I)) or "(" in c
res["ri_wage_ate_se_derived"] = round((w_hi - w_lo) / (2 * Z), 2)
locs["ri_wage_ate_se_derived"] = loc(RI_ID, 2, "Wage", ATE, "derived: (upper - lower) / (2 x 1.96)")
c, ln = t2["Wage"]["Control"]
res["ri_control_mean_wage"] = num(c.split("(")[0])
locs["ri_control_mean_wage"] = loc(RI_ID, 2, "Wage", "Control", "mean (SD in parentheses not used)")
c, ln = t2["Reemployment"][ATE]
matched.append(f"{RI_ID} L{ln}: {ri[ln-1]}")
res["ri_reemployment_ate_pp"] = parse_effect(c)[0]
locs["ri_reemployment_ate_pp"] = loc(RI_ID, 2, "Reemployment", ATE, "point estimate")
c, ln = t2["Weeks on UI"][ATE]
matched.append(f"{RI_ID} L{ln}: {ri[ln-1]}")
u_est, u_p, u_lo, u_hi = parse_effect(c)
res["ri_weeks_ui_ate"], res["ri_weeks_ui_ate_ci_low"], res["ri_weeks_ui_ate_ci_high"] = u_est, u_lo, u_hi
locs["ri_weeks_ui_ate"] = loc(RI_ID, 2, "Weeks on UI", ATE, "point estimate")
locs["ri_weeks_ui_ate_ci_low"] = loc(RI_ID, 2, "Weeks on UI", ATE, "95% CI lower bound")
locs["ri_weeks_ui_ate_ci_high"] = loc(RI_ID, 2, "Weeks on UI", ATE, "95% CI upper bound")

# ---------------- consistency / sensitivity (not required keys) ----------------
nd = NormalDist()
def p_from(b, s):
    return 2 * (1 - nd.cdf(abs(b / s)))
def star_from_p(p):
    return 3 if p < 0.01 else 2 if p < 0.05 else 1 if p < 0.10 else 0

star_checks = []
for tname, tab in [("6", t6), ("7", t7), ("A19", ta19)]:
    for row, d in tab.items():
        for col, (cell, se, _, _) in d.items():
            if se and cell:
                bb, ss = num(cell), num(se)
                pz = p_from(bb, ss)
                star_checks.append({"table": tname, "row": row, "col": col, "printed_stars": stars(cell),
                                    "z_implied_stars": star_from_p(pz), "z_p": round(pz, 4)})
star_mismatch = [x for x in star_checks if x["printed_stars"] != x["z_implied_stars"]]

a19_cell, a19_se, _, _ = ta19["ATT"]["CS-DR"]
w_mid = (w_lo + w_hi) / 2
sensitivity = {
    "dcr_att_5km_ci_z2.0": [round(b - 2.0 * s, 4), round(b + 2.0 * s, 4)],
    "dcr_att_5km_ci_using_tableA19_se": [round(b - Z * num(a19_se), 4), round(b + Z * num(a19_se), 4)],
    "dcr_hiv_att_5km_normal_p_approx": round(p_from(res["dcr_hiv_att_5km"], res["dcr_hiv_att_5km_se"]), 4),
    "ri_wage_ate_se_z2.0": round((w_hi - w_lo) / 4.0, 2),
    "ri_wage_ate_se_from_lower_half": round((w_est - w_lo) / Z, 2),
    "ri_wage_ate_se_from_upper_half": round((w_hi - w_est) / Z, 2),
    "ri_wage_ci_midpoint": w_mid,
    "ri_wage_printed_p": w_p,
    "ri_wage_normal_p_from_derived_se": round(p_from(w_est, (w_hi - w_lo) / (2 * Z)), 4),
    "significance_star_checks_n": len(star_checks),
    "significance_star_mismatches": star_mismatch,
}

data_issues = [
    {"issue":
              f"SE of the 5 km CS-DR ATT (-0.431) is printed as (0.210) in Table 6 but {a19_se} in Table A19 (CS-DR column, 5 km buffer, same +/-5-year window)",
     "handling": "Used Table 6 as required by the key definitions; the derived CI with the A19 SE is shown only as a sensitivity value."},
    {"issue": "Tables 6/7 state 608 never-treated control regions, the Table A19 note states 640 never-treated comparison regions",
     "handling": "Not used by any key; reported only."},
    {"issue": "DCR abstract says fatal hospitalisations 'fall by 35%', Section 6.1 says 'approximately 30% at 5 km'; exp(-0.354)-1 = -29.8%",
     "handling": "Not used by any key; text claims are not extracted; reported only."},
    {"issue": "DCR text calls HIV estimates 'close to zero' although the 5 km HIV ATT is -0.340 (SE 0.294, no star)",
     "handling": "Extracted the printed table value; significance taken from the (absent) star."},
    {"issue": "DCR Table 7 prints only significance stars, no exact p-values (and no Observations row)",
     "handling": "dcr_hiv_att_5km_p_value set to null with a 'not reported' locator; a normal-approximation p is given only as sensitivity, not as the key."},
    {"issue": "RESEA text prints the weeks-on-UI CI as positive magnitudes '[1.75 weeks, 2.24 weeks]' while Table 2 prints [-2.24, -1.75]",
     "handling": "Used Table 2 as printed (lower bound first, negative sign)."},
    {"issue": "RESEA text contains an encoding defect 'p ¡ 0.001' (for 'p < 0.001'); Table 2 cells were re-set by the extract editor",
     "handling": "Irrelevant to the keys; values taken from re-set Table 2 cells, which the extract states are verbatim."},
    {"issue": "RESEA Table 2 prints no standard errors, only p-values and 95% CIs",
     "handling": "ri_wage_ate_se_reported = false; SE derived from the printed CI with z = 1.96."},
]

results = {}
results.update(res)
results["locators"] = locs
results["derivation_notes"] = {
    "dcr_att_5km_ci": ("Normal (Wald) 95% CI: printed Table 6 ATT -0.431 +/- 1.96 x printed clustered SE 0.210; "
                       "assumes asymptotic normality of the CS-DR ATT; rounded to 4 decimals. "
                       "With 2.0 instead of 1.96: see sensitivity."),
    "ri_wage_ate_se": ("SE = (upper - lower) / (2 x 1.96) from the printed Table 2 95% CI [166, 2,141]; assumes the CI "
                       "is a symmetric normal-approximation interval (N = 23,549, so t ~ z); rounded to 2 decimals."),
}
results["sensitivity"] = sensitivity
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(os.path.join(DATA, "README.md")),
    DCR_ID + ".txt": wc_l(DCR_F),
    RI_ID + ".txt": wc_l(RI_F),
}
results["input_sha256"] = {k: hashlib.sha256(open(p, "rb").read()).hexdigest() for k, p in
                           [("README.md", os.path.join(DATA, "README.md")), (DCR_ID + ".txt", DCR_F), (RI_ID + ".txt", RI_F)]}
REQ = ["dcr_att_5km", "dcr_att_5km_se", "dcr_att_3km", "dcr_att_3km_se", "dcr_obs_5km", "dcr_narcotics_att_3km",
       "dcr_narcotics_att_3km_se", "dcr_fatal_att_1km", "dcr_fatal_att_1km_se", "dcr_hiv_att_5km", "dcr_hiv_att_5km_se",
       "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", "ri_wage_ate_ci_low", "ri_wage_ate_ci_high", "ri_wage_ate_se_reported", "ri_wage_ate_se_derived",
       "ri_control_mean_wage", "ri_reemployment_ate_pp", "ri_weeks_ui_ate", "ri_weeks_ui_ate_ci_low", "ri_weeks_ui_ate_ci_high"]
assert all(k in results for k in REQ)
results["n_analysed"] = len(REQ)
results["random_seed"] = None

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

print("Matched source lines:")
for m in matched:
    print("  " + m)
print("Extracted / derived values:")
for k in REQ:
    print(f"  {k} = {results[k]}")
print(f"Star-consistency checks: {len(star_checks)} cells, mismatches: {len(star_mismatch)}")