#!/usr/bin/env python3
"""Samleanalyse (DerSimonian–Laird random effects og fixed effects) med
kohort-uniqueness-sjekk, leave-one-out og scenarier.

Bruk som modul:
  from pool import pool, loo
  r = pool([(1.21,1.02,1.43,'norat2005','epic'), ...])        # (rr, lo, hi, study_id, cohort_id)
  r['re'] -> (rr, lo, hi), r['fe'], r['I2'], r['tau2'], r['weights_re']
"""
import math, json, sys, os, csv

Z = 1.959964

def _prep(rows):
    out = []
    for rr, lo, hi, sid, cid in rows:
        y = math.log(rr); se = (math.log(hi) - math.log(lo)) / (2 * Z)
        out.append((y, se * se, sid, cid))
    return out

def load_cohort_links(cohorts_csv):
    """Read cohorts.csv -> (parents, overlaps). parents: cohort -> parent cohort (sub-cohort of).
    overlaps: cohort -> set of cohorts it shares participants with (column overlaps_with, ';'-separated).
    Overlap is treated as symmetric. Added in R100 after the ATBC duplicate (Pietinen 1999 + Tammi 2024)
    passed a parent-only check."""
    parents, overlaps = {}, {}
    PARTIAL_LINKS.clear()
    with open(cohorts_csv, newline="", encoding="utf-8") as f:
        for c in csv.DictReader(f):
            cid = c["cohort_id"]
            parents[cid] = c.get("parent_cohort_id") or None
            for x in (c.get("overlaps_with") or "").split(";"):
                x = x.strip()
                if not x: continue
                if x.endswith("~"):  # partial overlap (some participants shared): reported, not refused (R104)
                    x = x[:-1]; PARTIAL_LINKS.add((cid, x)); PARTIAL_LINKS.add((x, cid))
                else:  # symmetric: registering the overlap on either cohort gives the same result (K29)
                    overlaps.setdefault(cid, set()).add(x); overlaps.setdefault(x, set()).add(cid)
    return parents, overlaps

def _root(cid, parents):
    seen = set()
    while cid in parents and parents[cid] and cid not in seen:
        seen.add(cid); cid = parents[cid]
    return cid

def _chain(cid, parents):
    out = [cid]
    while cid in parents and parents[cid] and parents[cid] not in out:
        cid = parents[cid]; out.append(cid)
    return out

PARTIAL_OVERLAPS = {}  # (study, study) -> description; filled by check_unique_cohorts (R104)
PARTIAL_LINKS = set()  # (cohort, cohort) pairs marked partial ('~') in cohorts.csv; filled by load_cohort_links

def check_unique_cohorts(rows, parents=None, overlaps=None, partial=None):
    """rows: list of (rr,lo,hi,study_id,cohort_id). parents: dict cohort->parent.
    overlaps: dict cohort->set of cohorts sharing participants (from cohorts.csv overlaps_with).
    Raises ValueError if two rows come from the same root cohort, or if one row's cohort (or any of
    its ancestors) is recorded as overlapping the other row's cohort (or any of its ancestors)."""
    parents = parents or {}; overlaps = overlaps or {}
    seen = {}
    for *_, sid, cid in rows:
        root = _root(cid, parents)
        if root in seen:
            raise ValueError(f"kohort {root} to ganger: {seen[root]} og {sid}")
        seen[root] = sid
    info = []
    for *_, sid, cid in rows:
        ch = _chain(cid, parents)
        ov = set()
        for c in ch:
            for o in overlaps.get(c, ()): ov.update(_chain(o, parents))
        info.append((sid, cid, set(ch), ov))
    for i in range(len(info)):
        for j in range(i + 1, len(info)):
            si, ci, chi, ovi = info[i]; sj, cj, chj, ovj = info[j]
            if (ovi & chj) or (ovj & chi):
                raise ValueError(f"overlappende kohorter: {ci} ({si}) og {cj} ({sj}) deler deltakere (cohorts.csv overlaps_with)")
    # Partial overlap (R104/K29): cohorts.csv marks links where only some participants are shared with "~"
    # (e.g. EPIC-Oxford, inside EPIC, and the Million Women Study). These are reported, not refused; the
    # sensitivity analyses without either cohort are in update2026.json.
    links = PARTIAL_LINKS if partial is None else partial
    if not links: return
    kids = {}
    for c in parents:
        for anc in _chain(c, parents)[1:]: kids.setdefault(anc, set()).add(c)
    for i in range(len(info)):
        for j in range(i + 1, len(info)):
            si, ci, chi, _ = info[i]; sj, cj, chj, _ = info[j]
            fam_i = chi | kids.get(ci, set()); fam_j = chj | kids.get(cj, set())
            hit = next(((a, b) for a in fam_i for b in fam_j if (a, b) in links), None)
            if hit:
                key = (si, sj)
                if key not in PARTIAL_OVERLAPS:
                    PARTIAL_OVERLAPS[key] = f"{hit[0]} ({si}) og {hit[1]} ({sj}) deler noen deltakere"
                    print(f"merk: delvis overlapp – {PARTIAL_OVERLAPS[key]}", file=sys.stderr)

def pool(rows, parents=None, allow_duplicate_cohorts=False, overlaps=None):
    if not allow_duplicate_cohorts:
        check_unique_cohorts(rows, parents, overlaps)
    d = _prep(rows)
    k = len(d)
    w = [1 / v for _, v, _, _ in d]
    sw = sum(w)
    yf = sum(wi * y for wi, (y, *_ ) in zip(w, d)) / sw
    Q = sum(wi * (y - yf) ** 2 for wi, (y, *_ ) in zip(w, d))
    C = sw - sum(wi * wi for wi in w) / sw
    tau2 = max(0.0, (Q - (k - 1)) / C) if k > 1 else 0.0
    wr = [1 / (v + tau2) for _, v, _, _ in d]
    swr = sum(wr)
    yr = sum(wi * y for wi, (y, *_ ) in zip(wr, d)) / swr
    se_r = math.sqrt(1 / swr); se_f = math.sqrt(1 / sw)
    I2 = max(0.0, (Q - (k - 1)) / Q) * 100 if Q > 0 else 0.0
    return {
        "k": k,
        "re": (math.exp(yr), math.exp(yr - Z * se_r), math.exp(yr + Z * se_r)),
        "fe": (math.exp(yf), math.exp(yf - Z * se_f), math.exp(yf + Z * se_f)),
        "Q": Q, "I2": I2, "tau2": tau2,
        "weights_re": {sid: wi / swr * 100 for wi, (_, _, sid, _) in zip(wr, d)},
        "weights_fe": {sid: wi / sw * 100 for wi, (_, _, sid, _) in zip(w, d)},
    }

def loo(rows, parents=None, overlaps=None):
    out = {}
    for i in range(len(rows)):
        sub = rows[:i] + rows[i+1:]
        out[rows[i][3]] = pool(sub, parents, overlaps=overlaps)["re"]
    return out

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

def combine_fixed(rows):
    """Slå sammen flere rader fra samme studie (f.eks. SMC distal+proksimal) med fixed effect til én rad."""
    d = _prep(rows)
    w = [1 / v for _, v, _, _ in d]
    y = sum(wi * yy for wi, (yy, *_ ) in zip(w, d)) / sum(w)
    se = math.sqrt(1 / sum(w))
    return (math.exp(y), math.exp(y - Z * se), math.exp(y + Z * se), rows[0][3], rows[0][4])

def _selftest():
    """Duplicate/overlap checks (R100). Run: python3 pool.py"""
    par = {"sub": "big", "atbc": None, "finnish_pooled": None, "x": None}
    ovl = {"finnish_pooled": {"atbc", "fmches"}}
    a, b = (1.1, 1.0, 1.2, "pietinen1999", "atbc"), (1.1, 1.0, 1.2, "tammi2024", "finnish_pooled")
    for rows, p_, o_, should_fail in (([a, b], par, ovl, True), ([b, a], par, ovl, True), ([a, b], par, None, False),
                                      ([(1.1, 1.0, 1.2, "s1", "sub"), (1.1, 1.0, 1.2, "s2", "big")], par, None, True),
                                      ([b, (1.1, 1.0, 1.2, "s3", "x")], par, ovl, False)):
        try: check_unique_cohorts(rows, p_, o_); failed = False
        except ValueError: failed = True
        assert failed == should_fail, (rows, should_fail)
    for links in ({("oxf", "mwsx"), ("mwsx", "oxf")},):  # partial overlap, either registration order: reported, not refused (R104, K29)
        PARTIAL_OVERLAPS.clear()
        check_unique_cohorts([(1.1, 1.0, 1.2, "papadimitriou2022", "epicp"), (1.1, 1.0, 1.2, "papier2025", "mwsx")],
                             {"epicp": None, "oxf": "epicp", "mwsx": None}, {}, partial=links)
        assert len(PARTIAL_OVERLAPS) == 1
        PARTIAL_OVERLAPS.clear()
        check_unique_cohorts([(1.1, 1.0, 1.2, "papier2025", "mwsx"), (1.1, 1.0, 1.2, "papadimitriou2022", "epicp")],
                             {"epicp": None, "oxf": "epicp", "mwsx": None}, {}, partial=links)
        assert len(PARTIAL_OVERLAPS) == 1
    PARTIAL_OVERLAPS.clear()
    return True

if __name__ == "__main__":
    # selvtest mot kjente CUP-tall gjøres i reproduce_cup.py
    r = pool([(1.21,1.02,1.43,'a','c1'), (1.01,0.90,1.14,'b','c2')])
    print(fmt(r["re"]), f"I2={r['I2']:.0f}%")
    print("selvtest overlapp:", _selftest())