"""Shared loader for the Stage-2 member pool (registry v3) from stored LOYO predictions.

Assembles, per held-out year and lead block, every member's out-of-sample calibrated P(yEHDE)
and its leakage-safe binary vote (P >= the member's per-region x block threshold), pooling three
prediction sources that share an identical (block, cell) row ordering:

  outputs/loyo_pred_v2/pred_v2_*.npz       7 sklearn families x 3 cascade cutoffs  (21 members)
  outputs/loyo_pred_v2_cnn/pred_v2cnn_*.npz    cnn_c1, ann_torch_c1                (2 members)
  outputs/loyo_pred_v3/pred_v3_*.npz       stage-20 expansion configs (_tele, _bal) (if present)

Registry v2's per-week active mask (results/tables/v2_active_members_by_week.csv) marks which
members are ACTIVE at each block; `active_only=True` drops the consensus-pruned member-weeks so the
probabilities in items 4/7 use the pruned pool. If v3 preds are absent (SLURM not finished) the
loader silently returns the v2+cnn pool, so downstream stages run on whatever exists and refresh
when v3 lands. Everything here is out-of-sample; no fitting.
"""
from __future__ import annotations
import os, glob, re
import numpy as np
import pandas as pd


def family_of(mid):
    """Base family of a member id (configs within a family are correlated)."""
    for suf in ("_tele_c1", "_bal_c1"):
        if mid.endswith(suf):
            return mid[:-len(suf)]
    return re.sub(r"_c[0-9.]+$", "", mid)


def family_balanced_weights(members, active):
    """Weight vector (sums to 1): equal weight per ACTIVE family, split across its active members.
    Returns (w, active_families). Handles the family-correlation problem in <member_definition>."""
    fams = [family_of(m) for m in members]
    w = np.zeros(len(members))
    act_fams = sorted({f for f, a in zip(fams, active) if a})
    for f in act_fams:
        idx = [i for i in range(len(members)) if fams[i] == f and active[i]]
        for i in idx:
            w[i] = 1.0 / len(act_fams) / len(idx)
    return w, act_fams


def _thr_lookup(decthr):
    d = {}
    for row in decthr:
        d[(str(row["tag"]), str(row["block"]), int(row["region"]))] = float(row["thr"])
    return d


def discover(cfg):
    """Return dict source -> {year: path} for the three prediction sources."""
    root = cfg["paths"]["outputs"]
    def m(sub, pat):
        return {int(os.path.basename(p).split("_")[-1][:-4]): p
                for p in glob.glob(os.path.join(root, sub, pat))}
    return {"v2": m("loyo_pred_v2", "pred_v2_*.npz"),
            "cnn": m("loyo_pred_v2_cnn", "pred_v2cnn_*.npz"),
            "v3": m("loyo_pred_v3", "pred_v3_*.npz")}


def member_list(cfg, include_v3=True):
    """Member ids present, in a stable order: v2 sklearn, cnn/ann, then v3 expansion."""
    src = discover(cfg)
    ids = []
    if src["v2"]:
        d = np.load(next(iter(src["v2"].values())), allow_pickle=True)
        ids += [k[2:] for k in d.files if k.startswith("p_")]
    if src["cnn"]:
        d = np.load(next(iter(src["cnn"].values())), allow_pickle=True)
        ids += [k[2:] for k in d.files if k.startswith("p_")]
    if include_v3 and src["v3"]:
        d = np.load(next(iter(src["v3"].values())), allow_pickle=True)
        ids += [k[2:] for k in d.files if k.startswith("p_")]
    return ids


def active_mask(cfg, members, blocks):
    """block -> boolean array over `members` (True = active). Uses registry v2 mask if present;
    members absent from the mask (e.g. v3 expansion) default to active everywhere."""
    path = os.path.join(cfg["paths"]["results"], "tables", "v2_active_members_by_week.csv")
    act = {b: np.ones(len(members), dtype=bool) for b in blocks}
    if not os.path.exists(path):
        return act
    df = pd.read_csv(path)
    lut = {(r.member_id, r.block): bool(r.active) for r in df.itertuples()}
    for b in blocks:
        act[b] = np.array([lut.get((m, b), True) for m in members])
    return act


def load_pool(cfg, blocks, include_v3=True):
    """Return (data, members) where data[year][block] = dict(P[n,M], V[n,M] uint8, y, cell, region).

    P = member probabilities, V = binary votes at each member's leakage-safe per-(block,region)
    threshold. Member order is `members` (from member_list)."""
    src = discover(cfg)
    years = sorted(src["v2"])
    members = member_list(cfg, include_v3=include_v3)
    block_ids = {b: i + 1 for i, b in enumerate(blocks)}

    # map member -> (source, prob_col)
    where = {}
    for m in members:
        col = "p_" + m
        for s in ("v3", "cnn", "v2") if include_v3 else ("cnn", "v2"):
            if src[s] and col in np.load(next(iter(src[s].values())), allow_pickle=True).files:
                where[m] = (s, col); break

    data = {}
    for y in years:
        loaded = {}
        thr = {}
        for s in ("v2", "cnn", "v3"):
            if src[s] and y in src[s]:
                dd = np.load(src[s][y], allow_pickle=True)
                loaded[s] = dd
                thr.update(_thr_lookup(dd["decthr"]))
        base = loaded["v2"]
        bid = base["block_id"].astype(int); reg = base["region"].astype(int)
        cell = base["cell"].astype(int); yv = base["y"].astype(float)
        data[y] = {}
        for b in blocks:
            bcode = block_ids[b]
            mm = bid == bcode
            reg_b = reg[mm]; cell_b = cell[mm]; y_b = yv[mm]
            P = np.full((mm.sum(), len(members)), np.nan, np.float32)
            V = np.zeros((mm.sum(), len(members)), np.uint8)
            rmax = int(reg_b.max()) if reg_b.size else 7
            for j, mid in enumerate(members):
                s, col = where[mid]
                dd = loaded[s]
                sb = dd["block_id"].astype(int) == bcode
                # sources share (block,cell) ordering
                p = dd[col][sb].astype(float)
                P[:, j] = p
                # vectorised per-region threshold lookup (region ids 1..rmax)
                lut = np.full(rmax + 1, 0.5)
                for rr in range(1, rmax + 1):
                    lut[rr] = thr.get((mid, b, rr), 0.5)
                V[:, j] = (p >= lut[reg_b]).astype(np.uint8)
            data[y][b] = dict(P=P, V=V, y=y_b, cell=cell_b, region=reg_b)
    return data, members