#!/usr/bin/env python3
"""Offline analysis (runs in Docker with no network): NeuroVault small-study inflation of whole-brain effect size.
Inputs (relative to the bundle root):
  data/included_maps.csv  frozen list (sha256 preregistered)
  data/map_pins.csv       per-map sha256 pinned at first fetch (status ok / download_failed / sha_mismatch)
  data/maps/<id>.nii[.gz] the map files
  data/mask_gm_2mm.nii.gz fixed grey-matter mask (sha256 in plan/parameters.json)
  plan/parameters.json    locked constants
Outputs: results/R1.json (primary), results/secondary.json, results/per_map.csv, results/exclusions.json"""
import csv, hashlib, json, math, os, re, sys, warnings
import numpy as np, pandas as pd, nibabel as nib
from scipy import stats
import statsmodels.api as sm
from statsmodels.regression.mixed_linear_model import MixedLM
import patsy
warnings.filterwarnings("ignore")
ROOT = sys.argv[1] if len(sys.argv) > 1 else "."
# Harness containers carry only code/, env/, data/: read the locked parameters from data/parameters.json, a byte-identical
# copy of plan/parameters.json (checked against plan/ whenever plan/ is present).
_pp, _dp = os.path.join(ROOT, "plan/parameters.json"), os.path.join(ROOT, "data/parameters.json")
if os.path.exists(_pp) and os.path.exists(_dp): assert open(_pp, "rb").read() == open(_dp, "rb").read(), "parameters copy differs from plan"
P = json.load(open(_dp if os.path.exists(_dp) else _pp))
OUT = os.path.join(ROOT, "results"); os.makedirs(OUT, exist_ok=True)

def r6(x):
    if x is None: return None
    if isinstance(x, (bool, np.bool_)): return bool(x)
    if isinstance(x, (int, np.integer)): return int(x)
    x = float(x)
    if not math.isfinite(x): return None
    return float(f"{x:.6g}") if abs(x) < 1e-3 and x != 0 else round(x, 6)

def sha(path): return hashlib.sha256(open(path, "rb").read()).hexdigest()

# ---------- inputs ----------
lst_path = os.path.join(ROOT, "data/included_maps.csv")
assert sha(lst_path) == P["frozen_list_sha256"], "frozen list sha256 mismatch"
mask_path = os.path.join(ROOT, "data/mask_gm_2mm.nii.gz")
assert sha(mask_path) == P["mask_sha256"], "mask sha256 mismatch"
L = pd.read_csv(lst_path, dtype={"df_meta": str, "fwhm_mm": str, "upload_year": str, "name": str, "contrast_definition": str},
                keep_default_na=False)
only = os.environ.get("NV_ONLY_IDS")  # pilot/testing only: restrict to these image ids
if only: L = L[L.image_id.isin([int(x) for x in only.split(",")])]
pins = {int(r["image_id"]): r for r in csv.DictReader(open(os.path.join(ROOT, "data/map_pins.csv")))}
mimg = nib.load(mask_path); M = np.asarray(mimg.dataobj) > 0
mijk = np.argwhere(M)  # C-order, deterministic
mw = (mimg.affine @ np.c_[mijk, np.ones(len(mijk))].T)  # 4 x n world coords
NMASK = len(mijk)

def load_masked(path):
    """Nearest-neighbour sample of the map at every mask voxel centre (world-coordinate lookup). Returns float64 vector with NaN where outside the map."""
    img = nib.load(path)
    data = np.asarray(img.dataobj, dtype=np.float64)
    while data.ndim > 3 and data.shape[-1] == 1: data = data[..., 0]
    if data.ndim != 3: raise ValueError("not_3d")
    A = img.affine
    if not np.all(np.isfinite(A)) or abs(np.linalg.det(A[:3, :3])) < 1e-9: raise ValueError("bad_affine")
    x = np.linalg.inv(A) @ mw
    ijk = np.floor(x[:3] + 0.5 + 1e-6).astype(np.int64).T
    inside = np.all((ijk >= 0) & (ijk < np.array(data.shape)), axis=1)
    v = np.full(NMASK, np.nan)
    v[inside] = data[ijk[inside, 0], ijk[inside, 1], ijk[inside, 2]]
    v[~np.isfinite(v)] = np.nan
    v[v == 0] = np.nan
    return v

def path_of(iid):
    for ext in (".nii.gz", ".nii"):
        p = os.path.join(ROOT, "data/maps", f"{iid}{ext}")
        if os.path.exists(p): return p
    return None

BASE_RE = re.compile(P["baseline_contrast_regex"], re.I)
rows = []
for r in L.sort_values("image_id").itertuples(index=False):
    iid, N, mt = int(r.image_id), int(r.N), r.map_type
    rec = {"image_id": iid, "collection_id": int(r.collection_id), "N": N, "map_type": mt,
           "fwhm_mm": r.fwhm_mm, "upload_year": r.upload_year, "task": r.task,
           "baseline_contrast": int(bool(BASE_RE.search(f"{r.name} | {r.contrast_definition}"))),
           "status": "ok", "coverage": None, "n_vox": None, "mean_sq": None, "q10_abs": None, "c10_null": None,
           "p999_abs": None, "delta2": None, "content_hash": ""}
    pin = pins.get(iid); p = path_of(iid)
    if pin is None or pin["status"] != "ok" or p is None:
        rec["status"] = "X1_download_failed"; rows.append(rec); continue
    if sha(p) != pin["sha256"]:
        rec["status"] = "X2_sha_mismatch"; rows.append(rec); continue
    try:
        v = load_masked(p)
    except Exception as e:
        rec["status"] = "X3_unreadable_" + str(e)[:20].replace(",", " "); rows.append(rec); continue
    ok = np.isfinite(v); cov = ok.mean(); rec["coverage"] = cov; rec["n_vox"] = int(ok.sum())
    if ok.sum() < 1000 or cov < P["qc_min_coverage"]:
        rec["status"] = "X4_coverage"; rows.append(rec); continue
    x = v[ok]; ax = np.abs(x)
    nu = N - 1
    m2 = float(np.mean(x * x))
    q10 = float(np.quantile(ax, 0.10)); p999 = float(np.quantile(ax, 0.999))
    c10 = float(stats.t.ppf(0.55, nu)) if mt == "T" else float(stats.norm.ppf(0.55))
    d2 = (m2 * (nu - 2) / nu - 1) / N if mt == "T" else (m2 - 1) / N
    rec.update(mean_sq=m2, q10_abs=q10, c10_null=c10, p999_abs=p999, delta2=d2,
               content_hash=hashlib.sha256(np.round(x, 3).astype(np.float32).tobytes() + ok.tobytes()).hexdigest()[:16])
    if q10 < P["qc_core_ratio"] * c10: rec["status"] = "X5_core_inconsistent_with_t"
    elif p999 > P["qc_extreme_a"] * math.sqrt(N) + P["qc_extreme_b"]: rec["status"] = "X6_implausible_extreme"
    rows.append(rec)

D = pd.DataFrame(rows)
# content duplicates among QC-passing maps: keep the lowest image id
qc_ok = D.status == "ok"
dup = D[qc_ok].duplicated("content_hash", keep="first")
D.loc[dup[dup].index, "status"] = "X7_duplicate_content"
D["fwhm_missing"] = (D.fwhm_mm == "").astype(int)
D["fwhm_c"] = pd.to_numeric(D.fwhm_mm.replace("", "0"))
D["log_N"] = np.log(D.N)

def logd(d2, floor): return 0.5 * np.log(np.maximum(d2, floor))
D["log_d"] = np.where(D.delta2.notna(), logd(D.delta2.fillna(0).astype(float), P["delta2_floor"]), np.nan)
D["delta_rms"] = np.exp(D.log_d)

def design(df, rhs):
    X = patsy.dmatrix(rhs, df, return_type="dataframe")
    X = X.loc[:, X.std(ddof=0).gt(0) | (X.columns == "Intercept")]
    # drop linearly dependent columns deterministically (keep earliest)
    keep = []
    for c in X.columns:
        trial = X[keep + [c]].to_numpy()
        if np.linalg.matrix_rank(trial) == len(keep) + 1: keep.append(c)
    return X[keep]

RHS_FULL = "log_N + C(map_type) + fwhm_c + fwhm_missing + C(upload_year) + C(task)"

def fit_mixed(df, y="log_d", rhs=RHS_FULL):
    df = df.sort_values("image_id").reset_index(drop=True)
    X = design(df, rhs)
    out = {"n_maps": int(len(df)), "n_collections": int(df.collection_id.nunique()), "n_fixed_effects": int(X.shape[1])}
    if df.collection_id.nunique() < 3 or len(df) < X.shape[1] + 3:
        out["error"] = "too_few"; return out
    res = None
    for method in P["mixedlm_methods"]:  # first converged fit in the preregistered order; else the last successful fit
        try:
            r_ = MixedLM(df[y].to_numpy(), X.to_numpy(), groups=df.collection_id.to_numpy()).fit(
                reml=True, method=method, maxiter=5000)
        except Exception:
            continue
        res = r_
        if res.converged: break
    if res is None:
        out["error"] = "fit_failed"; return out
    j = list(X.columns).index("log_N")
    b, se = res.fe_params[j], res.bse_fe[j]
    z = stats.norm.ppf(0.975)
    out.update(beta_logN=b, se=se, ci95_low=b - z * se, ci95_high=b + z * se, converged=bool(res.converged),
               method=method, collection_var=float(res.cov_re.iloc[0, 0]) if hasattr(res.cov_re, "iloc") else float(res.cov_re[0, 0]),
               resid_var=float(res.scale))
    return {k: r6(v) if not isinstance(v, str) else v for k, v in out.items()}

def fit_ols(df, y="log_d", rhs=RHS_FULL, cluster=False):
    df = df.sort_values("image_id").reset_index(drop=True)
    X = design(df, rhs)
    if len(df) < X.shape[1] + 3: return {"n": int(len(df)), "error": "too_few"}
    m = sm.OLS(df[y].to_numpy(), X.to_numpy())
    res = m.fit(cov_type="cluster", cov_kwds={"groups": df.collection_id.to_numpy()}) if cluster else m.fit(cov_type="HC3")
    j = list(X.columns).index("log_N"); b, se = res.params[j], res.bse[j]; z = stats.norm.ppf(0.975)
    return {k: r6(v) for k, v in dict(n=len(df), n_collections=df.collection_id.nunique(), beta_logN=b, se=se,
                                       ci95_low=b - z * se, ci95_high=b + z * se).items()}

def decide(f):
    if "beta_logN" not in f: return "inconclusive"
    if f["beta_logN"] <= P["support_beta_max"] and f["ci95_high"] < 0: return "support"
    if f["ci95_low"] > P["refute_ci_low_min"]: return "refute"
    return "inconclusive"

A = D[D.status == "ok"].copy()
prim = fit_mixed(A)
dec = decide(prim)
Nq = A.N.quantile([0, .1, .25, .5, .75, .9, 1]).tolist() if len(A) else []
R1 = {"n_listed": int(len(D)), "n_analyzed": int(len(A)), "n_collections": int(A.collection_id.nunique()),
      "N_quantiles_0_10_25_50_75_90_100": [r6(x) for x in Nq],
      "delta_rms_median": r6(A.delta_rms.median()) if len(A) else None,
      "n_delta2_nonpositive": int((A.delta2 <= 0).sum()), "n_delta2_below_floor": int((A.delta2 < P["delta2_floor"]).sum()),
      **{k: v for k, v in prim.items()}, "decision": dec, "supports_small_study_inflation": dec == "support"}
json.dump(R1, open(os.path.join(OUT, "R1.json"), "w"), indent=1, sort_keys=True)

# ---------- secondaries (descriptive; not part of the decision rule) ----------
S = {}
B = A[A.delta2 > 0].copy(); B["log_d"] = 0.5 * np.log(B.delta2.astype(float))
S["S1_exclude_nonpositive"] = fit_mixed(B)
C2 = A.copy(); C2["log_d"] = logd(C2.delta2.astype(float), P["delta2_floor_alt"]); S["S2_floor_alt"] = fit_mixed(C2)
rng = np.random.default_rng(P["seed"])
def one_per_collection(df, rng):
    g = df.sort_values(["collection_id", "image_id"]).groupby("collection_id").image_id.apply(list)
    pick = [ids[rng.integers(len(ids))] for ids in g]
    return df[df.image_id.isin(pick)]
S["S3_one_map_per_collection"] = fit_ols(one_per_collection(A, rng), rhs="log_N + C(map_type) + fwhm_c + fwhm_missing + C(upload_year)")
draws = []
for _ in range(P["one_per_collection_draws"]):
    f = fit_ols(one_per_collection(A, rng), rhs="log_N + C(map_type) + fwhm_c + fwhm_missing + C(upload_year)")
    if "beta_logN" in f: draws.append(f["beta_logN"])
if draws:
    S["S3b_one_map_per_collection_draws"] = {"n_draws": len(draws), "beta_median": r6(np.median(draws)),
        "beta_p2_5": r6(np.quantile(draws, .025)), "beta_p97_5": r6(np.quantile(draws, .975))}
cnt = A.collection_id.map(A.collection_id.value_counts())
for k in P["max_maps_per_collection"]:
    S[f"S4_collections_le_{k}_maps"] = fit_mixed(A[cnt <= k])
G = A.groupby("collection_id").agg(log_d=("log_d", "mean"), N=("N", "median"), image_id=("image_id", "min"),
                                    fracT=("map_type", lambda s: (s == "T").mean())).reset_index()
G["log_N"] = np.log(G.N)
S["S5_collection_mean"] = fit_ols(G, rhs="log_N + fracT")
S["S6_T_maps_only"] = fit_mixed(A[A.map_type == "T"])
S["S7_Z_maps_only"] = fit_mixed(A[A.map_type == "Z"])
S["S8_N_10_to_100"] = fit_mixed(A[A.N <= 100])
S["S9_baseline_contrast_covariate"] = fit_mixed(A, rhs=RHS_FULL + " + baseline_contrast")
S["S10_ols_cluster_collection"] = fit_ols(A, cluster=True)
S["S11_no_paradigm_fe"] = fit_mixed(A, rhs="log_N + C(map_type) + fwhm_c + fwhm_missing + C(upload_year)")
S["S12_known_paradigm_only"] = fit_mixed(A[A.task != "unknown"])
S["S13_keep_content_duplicates"] = fit_mixed(D[D.status.isin(["ok", "X7_duplicate_content"])])
S["S14_qc_load_and_coverage_only"] = fit_mixed(D[D.status.isin(["ok", "X5_core_inconsistent_with_t", "X6_implausible_extreme"])])

# S15: leave-collection-out ROI Cohen's d for paradigms with >= K collections
K = P["roi_min_collections"]
kc = A[A.task != "unknown"].groupby("task").collection_id.nunique()
elig = sorted(kc[kc >= K].index)
roi_rows = []
nroi = int(round(P["roi_fraction"] * NMASK))
for task in elig:
    T = A[A.task == task].sort_values("image_id")
    sums, cnts = {}, {}
    for r in T.itertuples(index=False):
        v = load_masked(path_of(r.image_id)) / math.sqrt(r.N)
        s = sums.setdefault(r.collection_id, np.zeros(NMASK)); c = cnts.setdefault(r.collection_id, np.zeros(NMASK))
        ok = np.isfinite(v); s[ok] += v[ok]; c[ok] += 1
    cm = {c: np.where(cnts[c] > 0, sums[c] / np.maximum(cnts[c], 1), np.nan) for c in sorted(sums)}
    tot = np.zeros(NMASK); tn = np.zeros(NMASK)
    for c in sorted(cm):
        ok = np.isfinite(cm[c]); tot[ok] += cm[c][ok]; tn[ok] += 1
    for r in T.itertuples(index=False):
        own = cm[r.collection_id]; ok = np.isfinite(own)
        st = tot.copy(); sn = tn.copy(); st[ok] -= own[ok]; sn[ok] -= 1
        cons = np.where(sn >= (len(cm) - 1) / 2, st / np.maximum(sn, 1), 0.0)
        roi = np.argsort(-np.abs(cons), kind="stable")[:nroi]
        v = load_masked(path_of(r.image_id))[roi]
        sg = np.sign(cons[roi]); okv = np.isfinite(v)
        d = float(np.mean(sg[okv] * v[okv]) / math.sqrt(r.N)) if okv.sum() >= nroi / 2 else np.nan
        roi_rows.append({"image_id": r.image_id, "roi_d": d})
if roi_rows:
    R = A.merge(pd.DataFrame(roi_rows), on="image_id").dropna(subset=["roi_d"])
    S["S15_roi_cohens_d"] = {"paradigms": elig, **fit_mixed(R, y="roi_d"),
                             "roi_d_median": r6(R.roi_d.median())}
    D = D.merge(pd.DataFrame(roi_rows), on="image_id", how="left")
else:
    S["S15_roi_cohens_d"] = {"paradigms": [], "error": "no_paradigm_with_enough_collections"}
    D["roi_d"] = np.nan
for k, f in S.items():
    if isinstance(f, dict) and "beta_logN" in f: f["direction_rule"] = decide(f)
json.dump(S, open(os.path.join(OUT, "secondary.json"), "w"), indent=1, sort_keys=True)
json.dump(dict(sorted(D.status.value_counts().items())), open(os.path.join(OUT, "exclusions.json"), "w"), indent=1)
cols = ["image_id", "collection_id", "N", "map_type", "fwhm_mm", "upload_year", "task", "baseline_contrast", "status",
        "coverage", "n_vox", "mean_sq", "q10_abs", "c10_null", "p999_abs", "delta2", "log_d", "roi_d", "content_hash"]
D = D.sort_values("image_id")[cols]
D["n_vox"] = D.n_vox.map(lambda x: "" if pd.isna(x) else str(int(x)))
for c in ["coverage", "mean_sq", "q10_abs", "c10_null", "p999_abs", "delta2", "log_d", "roi_d"]:
    D[c] = D[c].map(lambda x: "" if pd.isna(x) else f"{float(x):.6g}")
D.to_csv(os.path.join(OUT, "per_map.csv"), index=False, lineterminator="\n")
print(json.dumps(R1, indent=1))