"""F4-conf-01: extract printed values from two source extracts and 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
from scipy.stats import norm

DCR = "data/sources/arxiv-2609.27732v1.txt"
RI = "data/sources/arxiv-2606.14621v1.txt"
SID_D, SID_R = "arxiv-2609.27732v1", "arxiv-2606.14621v1"

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

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

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

def block(L, title, end_marker):
    i = next(k for k, s in enumerate(L) if s.startswith(title))
    j = next(k for k in range(i + 1, len(L)) if L[k].startswith(end_marker))
    return [s for s in L[i:j] if s.strip()], i

def num(c):
    """Return the printed number string (sign + decimals kept), stars stripped."""
    m = re.search(r"-?[\d,]*\.?\d+", c.replace("$", ""))
    return m.group(0)

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

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

D = lines(DCR)
R = lines(RI)

# ---- Table 6
t6, _ = block(D, "Table 6:", "Note:")
hdr6 = cells(t6[1])  # ['', '1 km', '3 km', '5 km']
def t6row(label):
    for k, s in enumerate(t6):
        if cells(s)[0] == label:
            print("[T6 match]", s)
            return cells(s), (cells(t6[k + 1]) if label == "ATT" else None)
att6, se6 = t6row("ATT")
obs6, _ = t6row("Observations")
print("[T6 SE line]", " | ".join(se6))
col6 = {h: i for i, h in enumerate(hdr6) if h}

# ---- Table 7
t7, _ = block(D, "Table 7:", "Note:")
hdr7 = cells(t7[1])
col7 = {h: i for i, h in enumerate(hdr7)}
def t7cell(row, col):
    for k, s in enumerate(t7):
        if cells(s)[0] == row:
            print("[T7 match]", s, "  ||  SE line:", t7[k + 1])
            return cells(s)[col7[col]], cells(t7[k + 1])[col7[col]]
    raise KeyError(row)

res = {"task_id": "F4-conf-01"}
loc = {}
def L6(row, col, extra=""):
    return f'{SID_D} | Table 6 | row "{row}" | column "{col}"' + extra
def L7(row, col, extra=""):
    return f'{SID_D} | Table 7 | row "{row}" | column "{col}"' + extra

for km in ["5", "3"]:
    c = col6[f"{km} km"]
    res[f"dcr_att_{km}km"] = to_num(num(att6[c]))
    res[f"dcr_att_{km}km_se"] = to_num(num(se6[c]))
    loc[f"dcr_att_{km}km"] = L6("ATT", f"{km} km")
    loc[f"dcr_att_{km}km_se"] = L6("ATT", f"{km} km", " (standard error in parentheses on the line below the ATT row)")
res["dcr_obs_5km"] = to_num(num(obs6[col6["5 km"]]))
loc["dcr_obs_5km"] = L6("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")]:
    v, se = t7cell(row, col)
    res[key] = to_num(num(v)); res[key + "_se"] = to_num(num(se))
    loc[key] = L7(row, col)
    loc[key + "_se"] = L7(row, col, " (standard error in parentheses on the line below)")
    if key == "dcr_hiv_att_5km":
        hiv_stars = stars(v)

# HIV p-value: the extract prints only stars, no exact p-values
res["dcr_hiv_att_5km_p_value"] = None
loc["dcr_hiv_att_5km_p_value"] = (f"not reported in {SID_D} (Table 7 prints only significance stars; "
                                  f'cell row "HIV" column "5 km" carries no star)')
res["dcr_hiv_att_5km_significant_10pct"] = hiv_stars >= 1

# derived CI
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)
for k in ["dcr_att_5km_ci_low", "dcr_att_5km_ci_high"]:
    loc[k] = f'{SID_D} | Table 6 | row "ATT" (estimate and SE) | column "5 km" | derived'

# ---- RI Table 2
i2 = next(k for k, s in enumerate(R) if s.startswith("Outcome | Treatment | Control"))
hdr2 = cells(R[i2]); col2 = {h: i for i, h in enumerate(hdr2)}
def t2(row):
    for s in R[i2 + 1:]:
        if cells(s)[0] == row:
            print("[RI T2 match]", s)
            return cells(s)
ATE = "Average treatment effect estimate"
def parse_effect(c):
    parts = [p.strip() for p in c.split(";")]
    est = num(parts[0])
    ci = re.search(r"\[(.*?)\]", c).group(1)
    lo, hi = [x.strip() for x in re.split(r",\s(?=-?\d)", ci)]
    return to_num(est), to_num(num(lo)), to_num(num(hi)), parts

def L2(row, col, extra=""):
    return f'{SID_R} | Table 2 | row "{row}" | column "{col}"' + extra

w = t2("Wage")
est, lo, hi, parts = parse_effect(w[col2[ATE]])
res["ri_wage_ate"], res["ri_wage_ate_ci_low"], res["ri_wage_ate_ci_high"] = est, lo, hi
loc["ri_wage_ate"] = L2("Wage", ATE, " (point estimate)")
loc["ri_wage_ate_ci_low"] = L2("Wage", ATE, " (95% CI, lower bound)")
loc["ri_wage_ate_ci_high"] = L2("Wage", ATE, " (95% CI, upper bound)")
res["ri_wage_ate_se_reported"] = any(re.search(r"\bSE\b|standard error", p, re.I) for p in parts)
res["ri_wage_ate_se_derived"] = round((hi - lo) / (2 * z), 2)
loc["ri_wage_ate_se_derived"] = L2("Wage", ATE, " (95% CI bounds) | derived")
res["ri_control_mean_wage"] = to_num(num(w[col2["Control"]].split("(")[0]))
loc["ri_control_mean_wage"] = L2("Wage", "Control", " (mean; SD in parentheses not used)")

r = t2("Reemployment")
res["ri_reemployment_ate_pp"] = parse_effect(r[col2[ATE]])[0]
loc["ri_reemployment_ate_pp"] = L2("Reemployment", ATE, " (point estimate)")

u = t2("Weeks on UI")
est, lo, hi, _ = parse_effect(u[col2[ATE]])
res["ri_weeks_ui_ate"], res["ri_weeks_ui_ate_ci_low"], res["ri_weeks_ui_ate_ci_high"] = est, lo, hi
loc["ri_weeks_ui_ate"] = L2("Weeks on UI", ATE, " (point estimate)")
loc["ri_weeks_ui_ate_ci_low"] = L2("Weeks on UI", ATE, " (95% CI, lower bound)")
loc["ri_weeks_ui_ate_ci_high"] = L2("Weeks on UI", ATE, " (95% CI, upper bound)")

# ---- consistency checks (printed for the report, not stored as keys)
wl, wh, we = res["ri_wage_ate_ci_low"], res["ri_wage_ate_ci_high"], res["ri_wage_ate"]
print(f"[check] wage CI half-widths: {we-wl} below, {wh-we} above (symmetric -> normal-type CI)")
print(f"[check] implied z for wage ATE = {we/res['ri_wage_ate_se_derived']:.3f}, "
      f"normal two-sided p = {2*norm.sf(abs(we/res['ri_wage_ate_se_derived'])):.4f} (printed p=0.022)")
print(f"[check] Table 6 5km SE {res['dcr_att_5km_se']} vs Table A19 CS-DR SE line:",
      [s for s in D if s.strip().startswith("| (0.218)")])
for n_ctl in (608, 640):
    print(f"[check] (treated + {n_ctl}) * 25 years:", [(t + n_ctl) * 25 for t in (7, 35, 50)])
import math
print(f"[check] exp(-0.354)-1 = {math.exp(-0.354)-1:.3f}; exp(0.390)-1 = {math.exp(0.390)-1:.3f}; "
      f"exp(-0.431)-1 = {math.exp(-0.431)-1:.3f}")

res["locators"] = loc
res["derivation_notes"] = {
    "dcr_att_5km_ci": ("Normal approximation with critical value 1.96: CI = ATT +/- 1.96*SE = -0.431 +/- 1.96*0.210, "
                       "using the Table 6 printed values (estimate -0.431, clustered SE 0.210); rounded to 4 decimals. "
                       "The paper's own inference may use a different (e.g. bootstrap/t) critical value; not stated in the extract."),
    "ri_wage_ate_se": ("SE = (upper - lower)/(2*1.96) = (2141 - 166)/3.92, using the printed Table 2 95% CI; assumes a "
                       "symmetric normal-approximation (Wald) interval, which is consistent with the printed bounds "
                       "(half-widths 987 and 988) and with the printed p=0.022 (implied normal p ~ 0.022). Rounded to cents."),
}
res["feasibility"] = "complete"
res["feasibility_reason"] = None
res["data_issues"] = [
    {"issue": "Table A19 (CS-DR column) prints SE (0.218) for the same 5 km ATT (-0.431) for which Table 6 prints SE (0.210).",
     "handling": "Key definitions name Table 6, so 0.210 is used for dcr_att_5km_se and the derived CI; discrepancy documented."},
    {"issue": "Number of control regions is internally inconsistent: Tables 6/7 and Figure 6 note say 608 never-treated regions, the Table A19 note says 640; the printed Observations (16,175 / 16,875 / 17,250) equal (treated + 640) x 25 exactly, not (treated + 608) x 25.",
     "handling": "Observations reported as printed (17250); inconsistency flagged; nothing re-estimated."},
    {"issue": "Abstract says fatal hospitalisations fall by 35%, but Table 7 5 km fatal ATT -0.354 implies exp(-0.354)-1 = -29.8% (Section 6.1 says ~30%). Text also calls HIV estimates 'close to zero' though the 3/5 km point estimates are -0.239/-0.340 (insignificant).",
     "handling": "Narrative inconsistencies noted; table values used as required."},
    {"issue": "RESEA text reports the weeks-on-UI CI as magnitudes '[1.75 weeks, 2.24 weeks]' while Table 2 prints [-2.24, -1.75].",
     "handling": "Table 2 signed values used (lower bound -2.24 first)."},
    {"issue": "Rendering defect in the RESEA extract: 'p ¡ 0.001' (should read p < 0.001); Tables 2-3 were re-set by the extract editor.",
     "handling": "No keyed value affected; values taken from the re-set Table 2 as verbatim."},
    {"issue": "DCR Table 7 reports only significance stars, no exact p-values; RESEA Table 2 reports CIs and p-values but no standard errors.",
     "handling": "dcr_hiv_att_5km_p_value = null (not reported); RESEA wage SE derived from CI with stated convention; ri_wage_ate_se_reported = false."},
]
files = ["data/README.md", DCR, RI]
res["n_input_rows"] = {os.path.basename(p): wc_l(p) for p in files}
res["n_analysed"] = 25  # 25 value keys (15 DCR incl. p-value check, significance flag and 2 derived; 10 RESEA incl. SE-reported flag and 1 derived)
res["random_seed"] = None

with open("out/results.json", "w") as f:
    json.dump(res, f, indent=2)
print(json.dumps({k: v for k, v in res.items() if k not in ("locators", "derivation_notes", "data_issues")}, indent=1))