#!/usr/bin/env python3
"""K19 - independent check of the Greenland-Longnecker (GL) dose-response slopes in update2026.

Written without importing any pipeline module (gl_doseresponse, catsets, mvmeta, derive_per100).
It re-implements, from the documented rules, (1) the input reconstruction (doses, cases,
denominators) and (2) three GL covariance variants:

  TEXTBOOK  Greenland & Longnecker 1992 as implemented by Orsini et al. 2006 / dosresmeta:
            diagonal = published var(log RR); off-diagonal = pseudo-count covariance
              ir (person-time):      1/A0*
              ci (n participants):   1/A0* - 1/N0          (risk-ratio model)
              cc (controls):         1/A0* + 1/B0*         (odds model, margins M_i = cases_i + controls_i)
  PIPELINE  the variant the pipeline actually codes (gl_doseresponse.gl_trend / mvmeta.gl_cov):
            pseudo-counts as above but n_participants AND n_controls are both put through the odds
            model with category total N_i := the printed denominator; off-diagonal =
            r_ij * sd_i * sd_j with r_ij = c0 / sqrt(pv_i pv_j) taken from the pseudo-counts
            (pv_i = pseudo variance), i.e. pseudo correlation rescaled to the published SDs.
  PIPE_FIXED_CC  as PIPELINE, but case-control margins M_i = cases_i + controls_i (bug fix only).
  PIPE_FIXED_ALL as PIPE_FIXED_CC, and n_participants through the risk-ratio model
            (c0 = 1/A0* - 1/N0, pv_i = 1/A_i* - 1/N_i + c0) as the docstring's "cumulative incidence" promises.

Slope: GLS through the origin, x = dose_i - dose_ref, y = log RR_i (natural log, natural dose),
SE = sqrt(1/(x' S^-1 x)), 95 % CI with z = 1.959964 (per 100 g: exp(100*b +- z*100*se)).

Run:  python3 gl_check.py [ROOT]   (prints markdown tables; writes gl_check_out.json next to this file)
"""
import csv, json, math, os, sys
import numpy as np
from scipy.optimize import brentq

ROOT = sys.argv[1] if len(sys.argv) > 1 else "/home/claude/db"
Z = 1.959964


def load(p):
    with open(p, newline="", encoding="utf-8") as f:
        return list(csv.DictReader(f))


def num(x):
    x = (x or "").strip().replace(",", ".")
    try:
        return float(x)
    except ValueError:
        return None


EST = load(os.path.join(ROOT, "database", "estimates.csv"))
EST_BY_ID = {e["estimate_id"]: e for e in EST}
CONV = load(os.path.join(ROOT, "database", "dose_conversions.csv"))
if os.path.exists(os.path.join(ROOT, "database", "dose_conversions_cup.csv")):
    CONV += load(os.path.join(ROOT, "database", "dose_conversions_cup.csv"))
DER = load(os.path.join(ROOT, "database", "derived_per100g.csv"))
STU = load(os.path.join(ROOT, "database", "studies.csv"))
CUPX = load(os.path.join(ROOT, "database", "cup_extraction.csv"))
JS = json.load(open(os.path.join(ROOT, "kontroll", "reanalyser", "update2026.json"), encoding="utf-8"))

# ------------------------------------------------------------------ documented input rules
PRI = {"paper": 0, "cup": 1, "us_assumed": 2}


def conversion(sid):
    """Rule: g_day_assigned from dose_conversions(_cup); priority paper > cup > us_assumed > other;
    rows with midpoint_rule == as_used are a separate variant, not the default."""
    best = None
    for c in CONV:
        if c["category_set_id"] != sid or not c.get("g_day_assigned") or c.get("midpoint_rule") == "as_used":
            continue
        p = PRI.get(c.get("portion_source"), 3)
        if best is None or p < best[0]:
            best = (p, c)
    return best[1] if best else None


def total_cases_lookup(study, outcome, sex):
    """Rule as coded: studies.csv cases_<outcome> (stored under sex 'all'), then cup_extraction
    cases_correct / cup_cases (first row wins). Returns (value, source, sex_specific?)."""
    for s in STU:
        if s["study_id"] == study and num(s.get("cases_" + outcome)):
            return num(s["cases_" + outcome]), "studies.csv cases_" + outcome, False
    for c in CUPX:
        if c["study_id"] == study and not c.get("cup_row_suffix") and c["outcome"] == outcome:
            v = num(c.get("cases_correct")) or num(c.get("cup_cases"))
            if v:
                return v, "cup_extraction " + c.get("cup_row_id", ""), False
    return None, None, False


QUANT = ("q1", "quintile", "quartile", "tertile", "t1", "quarter")


def build(sid):
    rows = sorted([e for e in EST if e["category_set_id"] == sid and e["contrast_type"] == "category" and not e.get("subgroup")],
                  key=lambda r: int(num(r["category_rank"]) or 0))
    e0 = rows[0]
    d = dict(set=sid, study=e0["study_id"], sex=e0["sex"], outcome=e0["outcome"], K=len(rows), ids=[r["estimate_id"] for r in rows],
             labels=[r["category_label"] for r in rows])
    if e0["dose_unit"] == "g_day" and all(num(r["dose_point"]) is not None for r in rows):
        d["dose"] = [num(r["dose_point"]) for r in rows]
        d["dose_src"] = "printed per-category " + (e0.get("dose_point_source") or "median/mean") + " (g/day)"
    else:
        c = conversion(sid)
        d["dose"] = [num(x) for x in c["g_day_assigned"].split(";")]
        d["dose_src"] = f"assigned: {c['original_unit']} -> g/day, rule={c['midpoint_rule']}, portion={c['portion_g'] or '-'} g ({c['portion_source'] or 'no portion'})"
    d["rr"] = [num(r["point"]) for r in rows]
    d["lo"] = [num(r["ci_low"]) for r in rows]
    d["hi"] = [num(r["ci_high"]) for r in rows]
    cases = [num(r["n_cases"]) for r in rows]
    if any(c is None for c in cases):
        tot, src, _ = total_cases_lookup(e0["study_id"], e0["outcome"], e0["sex"])
        d["cases"] = [tot * v / sum(d["rr"]) for v in d["rr"]]
        d["cases_src"] = f"ASSUMED: total {tot:g} ({src}, not sex-specific) split proportional to RR (i.e. equal person-time)"
        d["cases_assumed"] = True
    else:
        d["cases"] = cases; d["cases_src"] = "printed"; d["cases_assumed"] = False
    dt = e0.get("denominator_type") or "none"
    den = None
    if dt == "person_years":
        den = [num(r["person_years"]) for r in rows]
    elif dt in ("n_participants", "n_controls"):
        den = [num(r["n_in_category"]) for r in rows]
    if den and any(x is None for x in den):
        den, dt = None, "none"
    quant = any(q in (e0.get("category_label") or "").lower() for q in QUANT)
    if den is not None:
        d["den_src"] = f"printed {dt}"
    elif not d["cases_assumed"] and not quant:
        den = [c / v for c, v in zip(d["cases"], d["rr"])]; dt = "person_years"
        d["den_src"] = "ASSUMED: person-time proportional to cases/RR (crude RR = adjusted RR)"
    else:
        dt = "none"
        d["den_src"] = "ASSUMED: equal denominators (" + ("quantile categories" if quant else "cases assumed") + ")"
    d["den"], d["dt"] = den, dt
    return d


# ------------------------------------------------------------------ GL core (own implementation)
def pseudo(cases, den, rr, model):
    """Fitted pseudo-cells. model: 'ir' (rate), 'ci' (risk ratio), 'odds' (M = den as given).
    Returns (A*, B* or None)."""
    A = sum(cases); K = len(cases)
    if model in ("ir", "ci"):
        A0 = A / sum(rr[i] * den[i] / den[0] for i in range(K))
        return [rr[i] * A0 * den[i] / den[0] for i in range(K)], None
    M = den

    def f(A0):
        B0 = M[0] - A0
        return A0 + sum(rr[i] * A0 * M[i] / (B0 + rr[i] * A0) for i in range(1, K)) - A
    A0 = brentq(f, 1e-12, min(M[0], A) - 1e-12, xtol=1e-14)
    B0 = M[0] - A0
    Af = [A0] + [rr[i] * A0 * M[i] / (B0 + rr[i] * A0) for i in range(1, K)]
    return Af, [M[i] - Af[i] for i in range(K)]


def covariance(s, variant):
    rr, lo, hi, cases = s["rr"], s["lo"], s["hi"], s["cases"]
    K = len(rr); n = K - 1
    v = np.array([((math.log(hi[i]) - math.log(lo[i])) / (2 * Z)) ** 2 for i in range(1, K)])
    dt = s["dt"]; den = s["den"]
    if dt == "none":
        den, dt = [1.0] * K, "person_years"
    if variant == "TEXTBOOK":
        if dt == "person_years":
            Af, _ = pseudo(cases, den, rr, "ir"); c0 = 1 / Af[0]
        elif dt == "n_participants":
            Af, _ = pseudo(cases, den, rr, "ci"); c0 = 1 / Af[0] - 1 / den[0]
        else:  # n_controls: margins = cases + controls
            M = [a + b for a, b in zip(cases, den)]
            Af, Bf = pseudo(cases, M, rr, "odds"); c0 = 1 / Af[0] + 1 / Bf[0]
        S = np.full((n, n), c0); np.fill_diagonal(S, v)
        return S, Af
    # pipeline-type: correlation from pseudo-counts, rescaled to published SDs
    if dt == "person_years":
        Af, Bf = pseudo(cases, den, rr, "ir")
    elif variant == "PIPE_FIXED_ALL" and dt == "n_participants":
        Af, _ = pseudo(cases, den, rr, "ci"); c0 = 1 / Af[0] - 1 / den[0]
        pv = [1 / Af[i] - 1 / den[i] + c0 for i in range(1, K)]
        S = np.empty((n, n))
        for i in range(n):
            for j in range(n):
                S[i, j] = (1.0 if i == j else c0 / math.sqrt(pv[i] * pv[j])) * math.sqrt(v[i] * v[j])
        return S, Af
    else:
        M = den if variant == "PIPELINE" or dt == "n_participants" else [a + b for a, b in zip(cases, den)]
        Af, Bf = pseudo(cases, M, rr, "odds")
    if Bf is None:
        c0 = 1 / Af[0]; pv = [1 / Af[i] + c0 for i in range(1, K)]
    else:
        c0 = 1 / Af[0] + 1 / Bf[0]; pv = [1 / Af[i] + 1 / Bf[i] + c0 for i in range(1, K)]
    S = np.empty((n, n))
    for i in range(n):
        for j in range(n):
            r = 1.0 if i == j else c0 / math.sqrt(pv[i] * pv[j])
            S[i, j] = r * math.sqrt(v[i] * v[j])
    return S, Af


def slope(s, variant):
    S, Af = covariance(s, variant)
    y = np.log(np.array(s["rr"][1:]))
    x = np.array(s["dose"][1:]) - s["dose"][0]
    Si = np.linalg.inv(S)
    var = 1 / float(x @ Si @ x); b = var * float(x @ Si @ y)
    return b, math.sqrt(var), Af


def per100(b, se):
    return math.exp(100 * b), math.exp(100 * (b - Z * se)), math.exp(100 * (b + Z * se))


def fixed(rows):
    """Inverse-variance fixed combination of (logRR, var) pairs."""
    w = [1 / v for _, v in rows]; m = sum(wi * y for wi, (y, _) in zip(w, rows)) / sum(w); se = math.sqrt(1 / sum(w))
    return math.exp(m), math.exp(m - Z * se), math.exp(m + Z * se)


def dl(y, v):
    y = np.asarray(y); v = np.asarray(v); w = 1 / v; k = len(y)
    yf = (w * y).sum() / w.sum(); Q = (w * (y - yf) ** 2).sum(); C = w.sum() - (w ** 2).sum() / w.sum()
    t2 = max(0.0, (Q - (k - 1)) / C); wr = 1 / (v + t2); m = (wr * y).sum() / wr.sum(); se = math.sqrt(1 / wr.sum())
    return m, se, t2, Q


# ------------------------------------------------------------------ which rows does update2026 use? (D() selection)
def D(study, exposure, outcome, sex="all", method=None, cal="no", main=None):
    c = [d for d in DER if d["study_id"] == study and d["exposure"] == exposure and d["outcome"] == outcome and d["sex"] == sex
         and (d["is_calibrated"] or "no") == cal and (not method or d["method"] == method) and (main is None or d["is_main_model"] == main)
         and "as_used" not in d["dose_source"]]
    c.sort(key=lambda d: d["is_main_model"] != "yes")
    return c[0]


def set_of(derrow):
    return EST_BY_ID[derrow["source_estimate_ids"].split(";")[0]]["category_set_id"]


TARGETS = [  # (label, role, list of D() calls; >1 = sex strata combined fixed)
    ("ATBC (Pietinen 1999)", "main pool", [("pietinen1999", "red_unprocessed", "colorectal", "men")]),
    ("Netherlands MP (Tiemersma 2002)", "main pool", [("tiemersma2002", "red_unprocessed", "colorectal", "all")]),
    ("SMC (Larsson 2005)", "main pool", [("larsson2005", "red_main_dish", "colorectal", "women")]),
    ("MEC (Ollberding 2012)", "main pool", [("ollberding2012", "red_unprocessed", "colorectal", "all")]),
    ("UK Biobank (Knuppel 2020)", "main pool", [("knuppel2020", "red_unprocessed", "colorectal", "all", "yes")]),
    ("Alberta's Tomorrow (Al Rajabi 2021)", "main pool", [("alrajabi2021", "red_unprocessed", "colorectal", "men", "yes"), ("alrajabi2021", "red_unprocessed", "colorectal", "women", "yes")]),
    ("BWHS (Yiannakou 2022)", "main pool", [("yiannakou2022", "red_unprocessed", "colorectal", "women", "yes")]),
    ("Wada 2017 (Takayama only)", "multiverse alt.", [("wada2017", "red_unprocessed", "colorectal", "men"), ("wada2017", "red_unprocessed", "colorectal", "women")]),
    ("Bradbury 2020 (UK Biobank)", "multiverse alt.", [("bradbury2020", "red_unprocessed", "colorectal", "all", "yes")]),
    ("Vulcan 2017 (Malmo DCS)", "sensitivity", [("vulcan2017", "red_unprocessed", "colorectal", "men", "yes"), ("vulcan2017", "red_unprocessed", "colorectal", "women", "yes")]),
]
SPL_SETS = [u["set"] for u in JS["spline"]["sets"]]
VARIANTS = ("PIPELINE", "PIPE_FIXED_CC", "PIPE_FIXED_ALL", "TEXTBOOK")


def match(a, b, dp):
    return all(round(x, dp) == round(y, dp) for x, y in zip(a, b))


def fmt(t):
    return f"{t[0]:.4f} ({t[1]:.4f}-{t[2]:.4f})"


def main():
    out = {"sets": {}, "targets": [], "linear": {}}
    all_sets = []
    for lab, role, calls in TARGETS:
        for c in calls:
            der = D(c[0], c[1], c[2], c[3], "GL_categories", main=(c[4] if len(c) > 4 else None))
            all_sets.append((set_of(der), der))
    for sid in SPL_SETS:
        if sid not in [a for a, _ in all_sets]:
            all_sets.append((sid, None))
    print("## Per category set: pipeline vs own GL (per 100 g/day)\n")
    print("| set | pipeline (derived_per100g) | own PIPELINE variant | 2dp | 4dp | own PIPE_FIXED_CC | own PIPE_FIXED_ALL | own TEXTBOOK GL |")
    print("|---|---|---|---|---|---|---|---|")
    for sid, der in all_sets:
        s = build(sid)
        res = {}
        for vv in VARIANTS:
            b, se, Af = slope(s, vv); res[vv] = (b, se, per100(b, se), Af)
        pipe = (num(der["rr_per100"]), num(der["ci_low"]), num(der["ci_high"])) if der else None
        s["res"] = {k: dict(beta=v[0], se=v[1], per100=v[2], fitted=v[3]) for k, v in res.items()}
        s["pipeline_derived"] = pipe
        if der:
            s["pipeline_dose_check"] = der["notes"]; s["pipeline_den_note"] = der["denominator"]
        out["sets"][sid] = s
        own = res["PIPELINE"][2]
        print(f"| {sid} | {fmt(pipe) if pipe else '(spline only)'} | {fmt(own)} | {match(pipe, own, 2) if pipe else '-'} | {match(pipe, own, 4) if pipe else '-'} | {fmt(res['PIPE_FIXED_CC'][2])} | {fmt(res['PIPE_FIXED_ALL'][2])} | {fmt(res['TEXTBOOK'][2])} |")
    # target rows (with fixed combination of sex strata) vs update2026.json
    print("\n## update2026 rows built from GL\n")
    jsrows = {r["key"]: (r["rr"], r["lo"], r["hi"]) for r in JS["set"]}
    jsrows["Vulcan 2017 (Malmo DCS)"] = tuple(JS["vulcan"]["row"])
    print("| row | role | update2026 value | own PIPELINE (unrounded strata) | own PIPELINE (strata rounded 4 dp, as pipeline) | 2dp | 4dp | own PIPE_FIXED_ALL | own TEXTBOOK |")
    print("|---|---|---|---|---|---|---|---|---|")
    for lab, role, calls in TARGETS:
        sids = [set_of(D(c[0], c[1], c[2], c[3], "GL_categories", main=(c[4] if len(c) > 4 else None))) for c in calls]
        comb = {}
        for vv in VARIANTS:
            parts = []
            for sid in sids:
                p = out["sets"][sid]["res"][vv]["per100"]
                # combine exactly as update2026.fx/pool.combine_fixed does: from the rounded per-100 g row and its CI
                y = math.log(p[0]); se = (math.log(p[2]) - math.log(p[1])) / (2 * Z); parts.append((y, se ** 2))
            comb[vv] = fixed(parts) if len(parts) > 1 else out["sets"][sids[0]]["res"][vv]["per100"]
        pub = jsrows.get(lab)
        if pub is None:  # multiverse alternatives are not printed in the json; compare to the derived rows
            parts = []
            for sid in sids:
                p = out["sets"][sid]["pipeline_derived"]; parts.append((math.log(p[0]), ((math.log(p[2]) - math.log(p[1])) / (2 * Z)) ** 2))
            pub = fixed(parts) if len(parts) > 1 else out["sets"][sids[0]]["pipeline_derived"]
        out["targets"].append(dict(label=lab, role=role, sets=sids, pipeline=pub, own={k: v for k, v in comb.items()}))
        # the pipeline combines sex strata from derived_per100g.csv, which is rounded to 4 dp; do the same
        if len(sids) > 1:
            parts = []
            for sid in sids:
                p = [round(x, 4) for x in out["sets"][sid]["res"]["PIPELINE"]["per100"]]
                parts.append((math.log(p[0]), ((math.log(p[2]) - math.log(p[1])) / (2 * Z)) ** 2))
            comb["PIPELINE_rounded"] = fixed(parts)
        else:
            comb["PIPELINE_rounded"] = tuple(round(x, 4) for x in comb["PIPELINE"])
        out["targets"][-1]["own"] = comb
        print(f"| {lab} | {role} | {fmt(pub)} | {fmt(comb['PIPELINE'])} | {fmt(comb['PIPELINE_rounded'])} | {match(pub, comb['PIPELINE_rounded'], 2)} | {match(pub, comb['PIPELINE_rounded'], 4)} | {fmt(comb['PIPE_FIXED_ALL'])} | {fmt(comb['TEXTBOOK'])} |")
    # linear two-stage slope on the 12 spline sets
    print("\n## Linear two-stage slope, 12 spline sets (DL)\n")
    pub = tuple(JS["spline"]["linear_same_sets"])
    for vv in VARIANTS:
        ys, vs = [], []
        for sid in SPL_SETS:
            r = out["sets"][sid]["res"][vv]; ys.append(100 * r["beta"]); vs.append((100 * r["se"]) ** 2)
        m, se, t2, Q = dl(ys, vs)
        # update2026 uses z = 1.96 (not 1.959964) for this CI
        est = (math.exp(m), math.exp(m - 1.96 * se), math.exp(m + 1.96 * se))
        out["linear"][vv] = dict(est=est, tau2=t2, Q=Q, per_set=dict(zip(SPL_SETS, [math.exp(y) for y in ys])))
        print(f"- {vv}: {fmt(est)}  tau2={t2:.4f} Q={Q:.2f}  vs pipeline {fmt(pub)}  2dp={match(pub, est, 2)} 4dp={match(pub, est, 4)}")
    # impact on the main pool: replace the 7 GL rows by each variant (DL, z = 1.959964 as pool.py)
    print("\n## Main pool (k = 20, DL) with the GL rows from each variant\n")
    base = {r["key"]: (r["rr"], r["lo"], r["hi"]) for r in JS["set"]}
    tg = {t["label"]: t for t in out["targets"] if t["role"] == "main pool"}
    for vv in ("json",) + VARIANTS:
        rows = [tg[k]["own"][vv] if (vv != "json" and k in tg) else base[k] for k in base]
        y = [math.log(r[0]) for r in rows]; v = [((math.log(r[2]) - math.log(r[1])) / (2 * Z)) ** 2 for r in rows]
        m, se, t2, Q = dl(y, v); est = (math.exp(m), math.exp(m - Z * se), math.exp(m + Z * se))
        out["linear"]["mainpool_" + vv] = est
        print(f"- {vv}: {fmt(est)}  tau2={t2:.4f}  (json main: {fmt(JS['main']['re'])})")
    json.dump(out, open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "gl_check_out.json"), "w"), indent=1, default=float)
    return out


if __name__ == "__main__":
    main()