"""F4-conf-01: extract printed values from two arXiv extracts and derive two quantities.
Reads only data/; writes out/results.json deterministically. No randomness (random_seed = null).
Prints every matched source line so values can be traced."""
import json, os, re, math

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

def lines_of(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 lit(s):
    """Printed literal with $, commas, stars, units removed (keeps trailing zeros)."""
    s = s.strip().replace("$", "").replace(",", "")
    s = re.sub(r"\*+", "", s)
    return re.search(r"-?\d+(?:\.\d+)?", s).group(0)

def num(s):
    """Parse a printed number: strip $, thousands commas, stars, %, units. Keep sign."""
    s = s.strip().replace("$", "").replace(",", "")
    s = re.sub(r"\*+", "", s)
    m = re.search(r"-?\d+(?:\.\d+)?", s)
    v = m.group(0)
    return int(v) if re.fullmatch(r"-?\d+", v) else float(v)

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

def find_table(lines, title_prefix):
    for i, l in enumerate(lines):
        if l.startswith(title_prefix):
            return i
    raise ValueError(title_prefix)

def dcr_table(lines, title_prefix):
    """Return header columns and dict row_label -> (coef cells, se cells or None, line numbers)."""
    i = find_table(lines, title_prefix)
    j = i + 1
    while not lines[j].strip():
        j += 1
    header = cells(lines[j])[1:]  # first cell is row-label column ('' or 'Outcome')
    rows, k, nonblank = {}, j + 1, []
    while k < len(lines) and not lines[k].startswith("Note:"):
        if lines[k].strip():
            nonblank.append(k)
        k += 1
    for idx, ln in enumerate(nonblank):
        c = cells(lines[ln])
        if c[0] == "":
            continue  # SE line, attached below
        se, se_ln = None, None
        if idx + 1 < len(nonblank):
            nxt = cells(lines[nonblank[idx + 1]])
            if nxt[0] == "" and nxt[1].startswith("("):
                se, se_ln = nxt[1:], nonblank[idx + 1]
        rows[c[0]] = (c[1:], se, ln + 1, (se_ln + 1) if se_ln is not None else None)
    return header, rows

def stars(cell):
    m = re.search(r"(\*+)\s*$", cell.strip())
    return len(m.group(1)) if m else 0

dl, rl = lines_of(DCR), lines_of(RI)
h6, t6 = dcr_table(dl, "Table 6:")
h7, t7 = dcr_table(dl, "Table 7:")
hA, tA = dcr_table(dl, "Table A19:")

def get(tab, hdr, row, col, se=False):
    coef, ses, ln, sln = tab[row]
    j = hdr.index(col)
    cell = (ses if se else coef)[j]
    print(f"  matched line {sln if se else ln}: row '{row}' col '{col}' -> '{cell}'")
    v = num(cell.strip("()"))
    return v, cell

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

R, L = {}, {}
print("DCR Table 6 / Table 7 extraction:")
for key, tab, hdr, t, row, col in [
    ("dcr_att_5km", 6, h6, t6, "ATT", "5 km"),
    ("dcr_att_3km", 6, h6, t6, "ATT", "3 km"),
    ("dcr_narcotics_att_3km", 7, h7, t7, "Narcotics", "3 km"),
    ("dcr_fatal_att_1km", 7, h7, t7, "Drug-related (fatal)", "1 km"),
    ("dcr_hiv_att_5km", 7, h7, t7, "HIV", "5 km"),
]:
    R[key], cell = get(t, hdr, row, col)
    L[key] = loc(DCR_ID, tab, row, col)
    R[key + "_se"], _ = get(t, hdr, row, col, se=True)
    L[key + "_se"] = loc(DCR_ID, tab, row, col, " | standard error in parentheses on the line below the coefficient")
    if key == "dcr_hiv_att_5km":
        R["dcr_hiv_att_5km_significant_10pct"] = stars(cell) >= 1
R["dcr_obs_5km"], _ = get(t6, h6, "Observations", "5 km")
L["dcr_obs_5km"] = loc(DCR_ID, 6, "Observations", "5 km")

# exact p-value of HIV 5 km: search Table 7 block and whole extract for a printed p-value
dcr_text = "\n".join(dl)
has_p = re.search(r"\bp\s*=\s*0?\.\d+", dcr_text) is not None
R["dcr_hiv_att_5km_p_value"] = None  # Table 7 prints stars only
L["dcr_hiv_att_5km_p_value"] = (f"not reported in {DCR_ID} (Table 7 prints only the coefficient, SE and "
                                "significance stars; no exact p-value is printed anywhere in the extract)")
assert not has_p

# derived 95% CI from Table 6 SE
b, se = R["dcr_att_5km"], R["dcr_att_5km_se"]
R["dcr_att_5km_ci_low"] = round(b - Z * se, 4)
R["dcr_att_5km_ci_high"] = round(b + Z * se, 4)
for k in ("dcr_att_5km_ci_low", "dcr_att_5km_ci_high"):
    L[k] = f'{DCR_ID} | Table 6 | row "ATT" (and SE line below) | column "5 km" | derived: estimate -/+ 1.96 x SE'

print("RESEA Table 2 extraction:")
def ri_table(lines, header_start):
    i = [n for n, l in enumerate(lines) if l.startswith(header_start)][0]
    hdr = cells(lines[i])
    rows, k = {}, i + 1
    while lines[k].strip():
        c = cells(lines[k]); rows[c[0]] = (c, k + 1); k += 1
    return hdr, rows
h2, t2 = ri_table(rl, "Outcome | Treatment | Control | Treatment-control difference")
ATE = "Average treatment effect estimate"

def ate_cell(row):
    c, ln = t2[row]
    cell = c[h2.index(ATE)]
    print(f"  matched line {ln}: row '{row}' col '{ATE}' -> '{cell}'")
    est = num(cell.split(";")[0])
    lo, hi = re.search(r"\[(.+?),\s*(-?[\d,\.]+)\]", cell).groups()
    return est, num(lo), num(hi), cell

R["ri_wage_ate"], R["ri_wage_ate_ci_low"], R["ri_wage_ate_ci_high"], wcell = ate_cell("Wage")
L["ri_wage_ate"] = loc(RI_ID, 2, "Wage", ATE, " | point estimate")
L["ri_wage_ate_ci_low"] = loc(RI_ID, 2, "Wage", ATE, " | 95% CI lower bound")
L["ri_wage_ate_ci_high"] = loc(RI_ID, 2, "Wage", ATE, " | 95% CI upper bound")
R["ri_wage_ate_se_reported"] = ("SE" in wcell.replace("SD", "")) or ("(" in wcell)
R["ri_wage_ate_se_derived"] = round((R["ri_wage_ate_ci_high"] - R["ri_wage_ate_ci_low"]) / (2 * Z), 4)
L["ri_wage_ate_se_derived"] = (f'{RI_ID} | Table 2 | row "Wage" | column "{ATE}" | derived: '
                               f'(CI upper - CI lower) / (2 x 1.96); SE not reported in {RI_ID}')
c, ln = t2["Wage"]; cc = c[h2.index("Control")]
print(f"  matched line {ln}: row 'Wage' col 'Control' -> '{cc}'")
R["ri_control_mean_wage"] = num(cc.split("(")[0]); L["ri_control_mean_wage"] = loc(RI_ID, 2, "Wage", "Control", " | mean (SD in parentheses not used)")
R["ri_reemployment_ate_pp"], *_ = ate_cell("Reemployment")
L["ri_reemployment_ate_pp"] = loc(RI_ID, 2, "Reemployment", ATE, " | point estimate (pp)")
R["ri_weeks_ui_ate"], R["ri_weeks_ui_ate_ci_low"], R["ri_weeks_ui_ate_ci_high"], _ = ate_cell("Weeks on UI")
L["ri_weeks_ui_ate"] = loc(RI_ID, 2, "Weeks on UI", ATE, " | point estimate")
L["ri_weeks_ui_ate_ci_low"] = loc(RI_ID, 2, "Weeks on UI", ATE, " | 95% CI lower bound")
L["ri_weeks_ui_ate_ci_high"] = loc(RI_ID, 2, "Weeks on UI", ATE, " | 95% CI upper bound")

# ---------- cross-checks used to document data issues (not sources for keys) ----------
a19_se, _ = get(tA, hA, "ATT", "CS-DR", se=True)
a19_b, _ = get(tA, hA, "ATT", "CS-DR")
obs = [num(x) for x in t6["Observations"][0]]
treated = [num(x) for x in t6["Treated regions"][0]]
per640 = [o / (t + 640) for o, t in zip(obs, treated)]
per608 = [o / (t + 608) for o, t in zip(obs, treated)]
h3, t3 = ri_table(rl, "Outcome | Difference in means")
t3_wage = t3["Wage"][0][h3.index("Difference in means")]
wage_mid = (R["ri_wage_ate_ci_low"] + R["ri_wage_ate_ci_high"]) / 2
from statistics import NormalDist
p_implied = 2 * (1 - NormalDist().cdf(abs(R["ri_wage_ate"]) / R["ri_wage_ate_se_derived"]))
z_dcr = abs(b) / se
p_dcr = 2 * (1 - NormalDist().cdf(z_dcr))
z_hiv = abs(R["dcr_hiv_att_5km"]) / R["dcr_hiv_att_5km_se"]

issues = [
    {"issue": f"Same 5 km CS-DR ATT ({a19_b}) has SE {t6["ATT"][1][h6.index("5 km")].strip("()")} in Table 6 but SE {tA["ATT"][1][hA.index("CS-DR")].strip("()")} in Table A19 of {DCR_ID}.",
     "handling": "Keys name Table 6, so the Table 6 SE (0.210) is used for dcr_att_5km_se and the derived CI; the discrepancy is reported, not resolved."},
    {"issue": (f"Control-region count inconsistent in {DCR_ID}: Tables 6/7 and Fig. 6 note say 608 never-treated regions, Table A19 note says 640. "
               f"Table 6 observations divided by (treated + 640) = {per640}; by (treated + 608) = {[round(x,3) for x in per608]}."),
     "handling": "dcr_obs_5km reported as printed (17250); the arithmetic suggests the panel uses 640 controls x 25 periods, but this is not asserted as fact."},
    {"issue": f"Narrative of {DCR_ID} is inconsistent with itself: abstract says fatal hospitalisations fall by 35%, Section 6.1 says ~30% at 5 km (Table 7 5 km ATT -0.354 -> exp(-0.354)-1 = {math.exp(-0.354)-1:.3f}); text calls HIV estimates 'close to zero' although 5 km ATT is -0.340 (not significant).",
     "handling": "Only table values are extracted; narrative claims are not used."},
    {"issue": f"{RI_ID} narrative reports the weeks-on-UI CI as unsigned magnitudes [1.75, 2.24] while Table 2 prints [-2.24, -1.75].",
     "handling": "Signed Table 2 values used, lower bound first."},
    {"issue": f"Rendering defect in {RI_ID} extract line 39: 'p ¡ 0.001' (should read p < 0.001).",
     "handling": "No extracted key affected; ignored."},
    {"issue": f"{RI_ID}: Table 2 labels the column 'Average treatment effect estimate' while the Table 3 caption calls the same numbers ITT estimates; Table 2 caption does not state the CI construction beyond 'methodology described in the main text' (omitted from the extract).",
     "handling": "Values taken from Table 2 as required; SE derivation assumes a symmetric normal-approximation interval (critical value 1.96), which is an assumption."},
]

n_items = len(R)  # = 25: value keys assessed (13 DCR printed/marked, 2 DCR derived, 10 RESEA incl. 1 derived, 1 boolean)
out = dict(R)
out.update({
    "locators": L,
    "derivation_notes": {
        "dcr_att_5km_ci": "95% CI = ATT -/+ 1.96 x SE (normal approximation, critical value 1.96) using the printed Table 6 values ATT=-0.431, SE=0.210; rounded to 4 decimals. Table 6 does not print a CI. Sensitivity with critical value 2.0: see baseline_dcr_att_5km_ci_low_z2 / _high_z2.",
        "ri_wage_ate_se": "SE = (upper - lower) / (2 x 1.96) using the printed Table 2 95% CI [166, 2,141] (normal approximation, critical value 1.96, assumes a symmetric Wald-type interval); rounded to 4 decimals. The source prints no SE.",
    },
    "task_id": "F4-conf-01",
    "feasibility": "complete",
    "feasibility_reason": None,
    "data_issues": issues,
    "n_input_rows": {"README.md": wc_l(README), DCR_ID + ".txt": wc_l(DCR), RI_ID + ".txt": wc_l(RI)},
    "n_analysed": n_items,
    "random_seed": None,
    "baseline_dcr_att_5km_ci_low_z2": round(b - 2.0 * se, 4),
    "baseline_dcr_att_5km_ci_high_z2": round(b + 2.0 * se, 4),
    "baseline_ri_wage_ate_se_derived_z2": round((R["ri_wage_ate_ci_high"] - R["ri_wage_ate_ci_low"]) / 4.0, 4),
    "check_ri_wage_ci_midpoint": wage_mid,
    "check_ri_wage_p_implied_by_derived_se": round(p_implied, 4),
    "check_dcr_att_5km_p_implied_normal": round(p_dcr, 4),
    "check_dcr_hiv_5km_abs_z": round(z_hiv, 4),
    "check_dcr_a19_se_5km": a19_se,
    "check_table3_wage_identical_to_table2": t3_wage == wcell,
})
# Printed literals (to preserve trailing zeros such as 0.210 in the JSON number text)
PRINTED = {
    "dcr_att_5km": lit(t6["ATT"][0][h6.index("5 km")]), "dcr_att_5km_se": lit(t6["ATT"][1][h6.index("5 km")]),
    "dcr_att_3km": lit(t6["ATT"][0][h6.index("3 km")]), "dcr_att_3km_se": lit(t6["ATT"][1][h6.index("3 km")]),
    "dcr_obs_5km": lit(t6["Observations"][0][h6.index("5 km")]),
    "dcr_narcotics_att_3km": lit(t7["Narcotics"][0][h7.index("3 km")]), "dcr_narcotics_att_3km_se": lit(t7["Narcotics"][1][h7.index("3 km")]),
    "dcr_fatal_att_1km": lit(t7["Drug-related (fatal)"][0][h7.index("1 km")]), "dcr_fatal_att_1km_se": lit(t7["Drug-related (fatal)"][1][h7.index("1 km")]),
    "dcr_hiv_att_5km": lit(t7["HIV"][0][h7.index("5 km")]), "dcr_hiv_att_5km_se": lit(t7["HIV"][1][h7.index("5 km")]),
    "ri_wage_ate": lit(wcell.split(";")[0]), "ri_control_mean_wage": lit(cc.split("(")[0]),
    "ri_reemployment_ate_pp": lit(t2["Reemployment"][0][h2.index(ATE)].split(";")[0]),
    "ri_weeks_ui_ate": lit(t2["Weeks on UI"][0][h2.index(ATE)].split(";")[0]),
}
for k, v in PRINTED.items():
    assert float(v) == out[k], (k, v, out[k])
os.makedirs(os.path.join(ROOT, "out"), exist_ok=True)
with open(os.path.join(ROOT, "out", "results.json"), "w") as f:
    txt = json.dumps(out, indent=2, ensure_ascii=False)
    for k, v in PRINTED.items():
        txt, n = re.subn(r'^(  "%s": )[^,\n]+' % re.escape(k), lambda m: m.group(1) + v, txt, flags=re.M)
        assert n == 1, k
    assert json.loads(txt) == json.loads(json.dumps(out, ensure_ascii=False))
    f.write(txt + "\n")
print("Wrote out/results.json with", len(R), "value keys")