"""MP1 · paso 20: empirical Bayes jerárquico + modelos H (logística) y G (boosting monotónico), elección, recalibración,
intervalos, explicación por caso y deriva, para un desenlace (ou = oferta única, des = desierto, mod = modificación temprana).

Pre-registro: prereg/DEFINICIONES_MP1.md (md5 3da2be07…) §2-§7. Imagen bquantagents-worker (sklearn 1.9.1).
Uso: python3 20_modelo.py ou|des|mod [--fichero lotes.csv] [--muestra 0.1] [--sin-intervalos]
Lee /w/interno/<fichero>; escribe /w/interno/mp1_<h>/ (predicciones por lote, modelos) y /w/salidas/20_<h>.json.
"""
import sys, os, json, time, math, pickle, resource
import numpy as np, pandas as pd
import sklearn, scipy
from scipy.special import gammaln, expit, logit
from scipy.optimize import minimize_scalar
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.isotonic import IsotonicRegression
from sklearn.metrics import roc_auc_score

T0 = time.time()
H = sys.argv[1]
assert H in ("ou", "des", "mod")
FICH = sys.argv[sys.argv.index("--fichero") + 1] if "--fichero" in sys.argv else "lotes.csv"
MUESTRA = float(sys.argv[sys.argv.index("--muestra") + 1]) if "--muestra" in sys.argv else None
SIN_INT = "--sin-intervalos" in sys.argv
SEM = 20261026
OUT = f"/w/interno/mp1_{H}" + ("_muestra" if MUESTRA else "") + ("_pert" if "perturbado" in FICH else "")
os.makedirs(OUT, exist_ok=True)
inf = {"desenlace": H, "versiones": {"python": sys.version.split()[0], "sklearn": sklearn.__version__, "numpy": np.__version__,
       "pandas": pd.__version__, "scipy": scipy.__version__}, "fichero": FICH, "muestra": MUESTRA}
def log(*a):
    print(f"[{time.time() - T0:7.1f}s]", *a, flush=True)

# --- ventanas (pre-registro §5.3) ----------------------------------------------------------------------------------
VENT = {"ou": (("2018-01-01", "2024-12-31"), ("2025-01-01", "2025-12-31"), ("2026-01-01", "2026-04-30")),
        "des": (("2018-01-01", "2024-12-31"), ("2025-01-01", "2025-12-31"), ("2026-01-01", "2026-04-30")),
        "mod": (("2019-01-01", "2023-12-31"), ("2024-01-01", "2024-12-31"), ("2025-01-01", "2025-04-15"))}
HS = ("ou", "des", "mod")
NIV = ["glob", "mkt", "cpv", "terr", "toc", "org", "orgl", "orgcpv"]
PADRE = {"mkt": "glob", "cpv": "glob", "terr": "cpv", "toc": "glob", "org": "toc", "orgl": "org", "orgcpv": "org"}
CLAVES = {"glob": [], "mkt": ["cpv_cel", "proc_nat"], "cpv": ["cpv_cel"], "terr": ["ccaa_c", "cpv_cel"], "toc": ["toc_c", "ccaa_c"],
          "org": ["organo_clave"], "orgl": ["organo_clave"], "orgcpv": ["organo_clave", "cpv_cel"]}
NIV_VAR = ["org", "orgl", "orgcpv", "mkt", "terr", "toc"]           # niveles que entran como variables (§4)
MONO = ["org", "orgl", "orgcpv", "mkt", "terr"]                     # monotonía +1 en G (§5.2)
CATS = ["proc_nat", "tram", "tipo_contrato", "cpv_div", "tipo_organo", "ccaa", "mes_pub"]

# --- lectura -------------------------------------------------------------------------------------------------------
dt_str = {c: "string" for c in ["conjunto", "licitacion_id", "lote_u", "organo_clave", "cpv_cel", "cpv_div", "proc_nat", "tram",
                                "tipo_contrato", "tipo_organo", "ccaa", "organo_nom"]}
df = pd.read_csv(f"/w/interno/{FICH}", dtype=dt_str, parse_dates=["fecha_pub", "ffo", "fecha_adj"], low_memory=False)
log("leído", df.shape)
for c in ["fpub_ok", "todo_copia", "resuelto", "adjudicado", "des", "desistido", "y_mod", "y_mod10", "sara", "ue", "es_lote",
          "suman100", "pcap", "ppt"]:
    if c in df:
        df[c] = df[c].map({True: True, False: False, "true": True, "false": False, "True": True, "False": False})
df["ccaa_c"] = df["ccaa"].fillna("??")
df["toc_c"] = df["tipo_organo"].fillna("??")

# --- población P y universo del desenlace (§2, §3) -----------------------------------------------------------------
hist = df["conjunto"] == "hist"
P = hist & df["fpub_ok"].fillna(False) & ~df["todo_copia"].fillna(False) & df["plazo"].notna() & df["ffo"].notna() \
    & df["resuelto"].fillna(False) & (df["dias_res"] <= 175)
if H == "ou":
    U_ = P & (df["num_ofertas"] >= 1); y = (df["num_ofertas"] == 1)
elif H == "des":
    U_ = P; y = df["des"].fillna(False)
else:
    U_ = P & df["adjudicado"].fillna(False); y = df["y_mod"].fillna(False)
df["y"] = y.astype(int)
fp = df["fecha_pub"]
def en(v):
    return (fp >= pd.Timestamp(v[0])) & (fp <= pd.Timestamp(v[1]))
vt, vr, vp = VENT[H]
df["ventana"] = np.select([U_ & en(vt), U_ & en(vr), U_ & en(vp), df["conjunto"] == "pred"], ["train", "recal", "test", "pred"], "")
inf["poblacion"] = {"hist": int(hist.sum()), "P": int(P.sum()), "universo": int(U_.sum()),
                    "por_ventana": {k: int(v) for k, v in df["ventana"].value_counts().items() if k},
                    "positivos": {k: int(df.loc[df["ventana"] == k, "y"].sum()) for k in ("train", "recal", "test")}}
log("población", inf["poblacion"])
d = df[df["ventana"] != ""].copy()
del df
if MUESTRA:
    rng0 = np.random.default_rng(SEM)
    keep = (d["ventana"] != "train") | (rng0.random(len(d)) < MUESTRA)
    d = d[keep].copy()
    log("muestra", len(d))
d.reset_index(drop=True, inplace=True)
TR, RC, TE, PR = (d["ventana"] == w for w in ("train", "recal", "test", "pred"))

# --- 5.1 empirical Bayes jerárquico ---------------------------------------------------------------------------------
def bb_ll(m, k, n, pi):
    a = m * pi; b = m * (1 - pi)
    return np.sum(gammaln(k + a) + gammaln(n - k + b) - gammaln(n + a + b) - gammaln(a) - gammaln(b) + gammaln(a + b))

def estimar_m(k, n, pi):
    ok = n > 0
    k, n, pi = k[ok], n[ok], np.clip(pi[ok], 1e-6, 1 - 1e-6)
    if len(k) < 5:
        return 20.0
    grid = np.exp(np.linspace(np.log(0.5), np.log(2000), 120))
    ll = np.array([bb_ll(m, k, n, pi) for m in grid])
    i = int(np.argmax(ll))
    lo, hi = np.log(grid[max(i - 1, 0)]), np.log(grid[min(i + 1, len(grid) - 1)])
    r = minimize_scalar(lambda lm: -bb_ll(np.exp(lm), k, n, pi), bounds=(lo, hi), method="bounded")
    return float(np.exp(r.x))

EB_M = {}
post = {}       # (h, nivel) -> media posterior por fila
alfa_beta = {}  # (h, nivel) -> (a, b) de la posterior por fila
for h in HS:
    for niv in NIV:
        k = d[f"{niv}_{h}_k"].to_numpy(float); n = d[f"{niv}_{h}_n"].to_numpy(float)
        if niv == "glob":
            p = (k + 1) / (n + 2); post[(h, niv)] = p; alfa_beta[(h, niv)] = (k + 1, n - k + 1); EB_M[f"{h}.{niv}"] = None
            continue
        pi = post[(h, PADRE[niv])]
        # m con las celdas-año distintas de la ventana de entrenamiento (§5.1)
        claves = CLAVES[niv] + ["T"]
        sub = d.loc[TR, claves].copy(); sub["k"] = k[TR.to_numpy()]; sub["n"] = n[TR.to_numpy()]; sub["pi"] = pi[TR.to_numpy()]
        sub = sub.dropna(subset=[c for c in CLAVES[niv]]).drop_duplicates(subset=claves)
        m = estimar_m(sub["k"].to_numpy(), sub["n"].to_numpy(), sub["pi"].to_numpy())
        EB_M[f"{h}.{niv}"] = {"m": round(m, 4), "celdas": int(len(sub)), "celdas_n_pos": int((sub["n"] > 0).sum())}
        a = k + m * pi; b = n - k + m * (1 - pi)
        post[(h, niv)] = a / (a + b); alfa_beta[(h, niv)] = (a, b)
inf["eb_m"] = EB_M
log("EB", {k: (v["m"] if v else None) for k, v in EB_M.items()})

# --- variables (§4) ---------------------------------------------------------------------------------------------------
def num_basicas(x):
    f = pd.DataFrame(index=x.index)
    pbl = x["pbl"]; ve = x["ve_exp"]
    f["log_pbl"] = np.log10(np.maximum(pbl.fillna(1), 1)); f["pbl_na"] = pbl.isna().astype(float)
    f["log_ve"] = np.log10(np.maximum(ve.fillna(1), 1)); f["ve_na"] = ve.isna().astype(float)
    f["dur"] = np.where(x["dur"].isna(), -1, np.minimum(x["dur"].fillna(0), 120)); f["dur_na"] = x["dur"].isna().astype(float)
    f["plazo"] = np.minimum(x["plazo"].fillna(-1), 120); f["plazo_lt15"] = (x["plazo"] < 15).astype(float)
    f["log_nlotes"] = np.log1p(x["n_lotes_an"].fillna(1)); f["es_lote"] = x["es_lote"].fillna(False).astype(float)
    f["sara"] = x["sara"].fillna(False).astype(float); f["sara_na"] = x["sara"].isna().astype(float)
    f["ue"] = x["ue"].fillna(False).astype(float)
    f["cpv_n_div"] = x["cpv_n_div"].fillna(-1); f["cpv_n_div_na"] = x["cpv_n_div"].isna().astype(float)
    ok_tipo = x["suman100"].fillna(False) & (x["n_sin_tipo"].fillna(1) == 0)
    ps = x["peso_suma"].replace(0, np.nan)
    f["sin_crit"] = x["n_crit"].isna().astype(float); f["n_crit"] = np.minimum(x["n_crit"].fillna(0), 15)
    f["un_solo_crit"] = (x["n_crit"] == 1).astype(float)
    f["peso_precio"] = np.where(ok_tipo, (x["peso_precio"].fillna(0) / ps).fillna(0), -1)
    pjv = np.where(ok_tipo, (x["peso_jv"].fillna(0) / ps).fillna(0), -1)
    f["peso_jv"] = pjv; f["peso_jv_gt50"] = (pjv > 0.5).astype(float)
    f["pcap"] = x["pcap"].fillna(False).astype(float); f["ppt"] = x["ppt"].fillna(False).astype(float)
    obras = (x["tipo_contrato"] == "Obras").fillna(False)
    for nom, (u, ob) in {"bajo_60k": (60000, False), "bajo_139k": (139000, False), "bajo_215k": (215000, False),
                         "bajo_80k_obras": (80000, True), "bajo_2M_obras": (2000000, True), "bajo_5M_obras": (5382000, True)}.items():
        r = pbl.fillna(0) / u
        f[nom] = ((obras if ob else ~obras) & (r >= 0.85) & (r < 1.0)).astype(float)
    f["org_nsp"] = (x["nsp_n"] + 20 * 0.10) / (x["nsp_adj"] + 20); f["log_org_adj"] = np.log1p(x["nsp_adj"])
    return f

F = num_basicas(d)
GR = {}  # grupo de explicación de cada variable
for c in ["log_pbl", "pbl_na", "log_ve", "ve_na", "bajo_60k", "bajo_139k", "bajo_215k", "bajo_80k_obras", "bajo_2M_obras", "bajo_5M_obras"]:
    GR[c] = "importe y umbrales"
for c in ["dur", "dur_na"]: GR[c] = "duración"
for c in ["plazo", "plazo_lt15"]: GR[c] = "plazo"
for c in ["log_nlotes", "es_lote"]: GR[c] = "lotes"
for c in ["sara", "sara_na", "ue"]: GR[c] = "procedimiento y tramitación"
for c in ["cpv_n_div", "cpv_n_div_na"]: GR[c] = "objeto"
for c in ["sin_crit", "n_crit", "un_solo_crit", "peso_precio", "peso_jv", "peso_jv_gt50", "pcap", "ppt"]: GR[c] = "criterios y pliegos"
for c in ["org_nsp", "log_org_adj"]: GR[c] = "historial del órgano"
for h in HS:
    for niv in NIV_VAR:
        p = np.clip(post[(h, niv)], 1e-4, 1 - 1e-4)
        F[f"eb_{h}_{niv}"] = logit(p); F[f"ln_{h}_{niv}"] = np.log1p(d[f"{niv}_{h}_n"].to_numpy(float))
        g = ("historial del órgano" if niv in ("org", "orgl", "orgcpv") else "mercado y territorio") if h == H else "otros desenlaces"
        GR[f"eb_{h}_{niv}"] = g; GR[f"ln_{h}_{niv}"] = g
GRC = {"proc_nat": "procedimiento y tramitación", "tram": "procedimiento y tramitación", "tipo_contrato": "objeto",
       "cpv_div": "objeto", "tipo_organo": "órgano (tipo, CCAA)", "ccaa": "órgano (tipo, CCAA)", "mes_pub": "mes"}
NUM = list(F.columns)
# columnas que leen las variables (test estático, §12 (6))
LEIDAS = sorted(set(["pbl", "ve_exp", "dur", "plazo", "n_lotes_an", "es_lote", "sara", "ue", "cpv_n_div", "suman100", "n_sin_tipo",
                     "peso_suma", "n_crit", "peso_precio", "peso_jv", "pcap", "ppt", "tipo_contrato", "nsp_n", "nsp_adj"] + CATS +
                    [f"{niv}_{h}_{s}" for h in HS for niv in NIV for s in ("k", "n")] + ["cpv_cel", "proc_nat", "ccaa", "tipo_organo",
                    "organo_clave", "T"]))
inf["columnas_leidas"] = LEIDAS

# categóricas: niveles con >= 30 lotes de entrenamiento; resto «otros»; NULL nivel propio «∅»
CAT_NIV = {}
Ccat = pd.DataFrame(index=d.index)
for c in CATS:
    v = d[c].astype("string").fillna("∅") if c != "mes_pub" else d[c].astype("Int64").astype("string").fillna("∅")
    cnt = v[TR].value_counts()
    niv = sorted([x for x, n in cnt.items() if n >= 30])
    CAT_NIV[c] = niv
    Ccat[c] = v.where(v.isin(niv), "otros")

def disenno_H(idx, Fx=None, Cx=None):
    Fx = F.loc[idx] if Fx is None else Fx
    Cx = Ccat.loc[idx] if Cx is None else Cx
    Z = (Fx[NUM].to_numpy(float) - MU) / SD
    D = [np.column_stack([(Cx[c].to_numpy() == lv).astype(float) for lv in CAT_NIV[c] + ["otros"]]) for c in CATS]
    return np.column_stack([Z] + D)

MU = F.loc[TR, NUM].mean().to_numpy(); SD = F.loc[TR, NUM].std(ddof=0).to_numpy(); SD[SD == 0] = 1
COLS_H = NUM + [f"{c}={lv}" for c in CATS for lv in CAT_NIV[c] + ["otros"]]
GRUPO_H = [GR[c] for c in NUM] + [GRC[c] for c in CATS for lv in CAT_NIV[c] + ["otros"]]

ytr = d.loc[TR, "y"].to_numpy(); yrc = d.loc[RC, "y"].to_numpy()
XH_tr = disenno_H(TR)
log("diseño H", XH_tr.shape)
mH = LogisticRegression(C=1.0, solver="lbfgs", max_iter=2000)
mH.fit(XH_tr, ytr)
inf["H"] = {"iteraciones": int(mH.n_iter_[0]), "n_coef": int(XH_tr.shape[1])}
del XH_tr
log("H ajustado", inf["H"])

# G: numéricas + categóricas como códigos
def disenno_G(idx, Fx=None, Cx=None):
    Fx = F.loc[idx] if Fx is None else Fx
    Cx = Ccat.loc[idx] if Cx is None else Cx
    cols = [Fx[NUM].to_numpy(float)]
    for c in CATS:
        mapa = {lv: i for i, lv in enumerate(CAT_NIV[c] + ["otros"])}
        cols.append(Cx[c].map(mapa).to_numpy(float)[:, None])
    return np.column_stack(cols)
mono = [1 if (c.startswith(f"eb_{H}_") and c.split("_")[-1] in MONO) else 0 for c in NUM] + [0] * len(CATS)
catmask = [False] * len(NUM) + [True] * len(CATS)
mG = HistGradientBoostingClassifier(learning_rate=0.05, max_iter=600, max_leaf_nodes=63, min_samples_leaf=200,
                                    l2_regularization=1.0, early_stopping=True, validation_fraction=0.1, n_iter_no_change=30,
                                    random_state=SEM, categorical_features=catmask, monotonic_cst=mono)
XG_tr = disenno_G(TR)
mG.fit(XG_tr, ytr)
inf["G"] = {"iteraciones": int(mG.n_iter_), "monotonia": [c for c, m in zip(NUM, mono) if m]}
del XG_tr
log("G ajustado", inf["G"])

def p_H(idx, Fx=None, Cx=None):
    return mH.predict_proba(disenno_H(idx, Fx, Cx))[:, 1]
def p_G(idx, Fx=None, Cx=None):
    return mG.predict_proba(disenno_G(idx, Fx, Cx))[:, 1]

# --- elección en la ventana de recalibración (§5.2) -------------------------------------------------------------------
pr_H = p_H(RC); pr_G = p_G(RC)
aH, aG = roc_auc_score(yrc, pr_H), roc_auc_score(yrc, pr_G)
ELEG = "G" if aG - aH >= 0.002 else "H"
inf["eleccion"] = {"auc_recal_H": aH, "auc_recal_G": aG, "elegido": ELEG}
log("elección", inf["eleccion"])
praw = p_G if ELEG == "G" else p_H
pr_raw = pr_G if ELEG == "G" else pr_H

# --- recalibración --------------------------------------------------------------------------------------------------
npos = int(yrc.sum())
if npos >= 1000:
    cal = IsotonicRegression(out_of_bounds="clip", y_min=1e-6, y_max=1 - 1e-6).fit(pr_raw, yrc); CAL = "isotonica"
    def calibrar(p): return np.clip(cal.predict(p), 1e-6, 1 - 1e-6)
else:
    lr = LogisticRegression(C=1e6, solver="lbfgs").fit(logit(np.clip(pr_raw, 1e-6, 1 - 1e-6))[:, None], yrc); CAL = "platt"
    def calibrar(p): return lr.predict_proba(logit(np.clip(p, 1e-6, 1 - 1e-6))[:, None])[:, 1]
inf["calibrador"] = {"tipo": CAL, "positivos_recal": npos}

# --- predicciones centrales -------------------------------------------------------------------------------------------
res = d.loc[RC | TE | PR, ["oid", "ventana", "licitacion_id", "lote_u", "organo_clave", "cpv_cel", "proc_nat", "tipo_contrato",
                           "ccaa", "tipo_organo", "fecha_pub", "ffo", "y", "pbl", "organo_nom"]].copy()
idx = res.index
pH = p_H(idx); pG = p_G(idx)
res["p_H"] = pH; res["p_G"] = pG
res["p_raw"] = pG if ELEG == "G" else pH
res["p"] = calibrar(res["p_raw"].to_numpy())
# B1 «órgano × CPV» (receta orgcpv_rate de E1 sobre los mismos historiales), P0 = tasa del entrenamiento
P0 = float(ytr.mean())
ko, no = d.loc[idx, f"org_{H}_k"].to_numpy(float), d.loc[idx, f"org_{H}_n"].to_numpy(float)
kc, nc = d.loc[idx, f"orgcpv_{H}_k"].to_numpy(float), d.loc[idx, f"orgcpv_{H}_n"].to_numpy(float)
org_rate = (ko + 20 * P0) / (no + 20)
res["b1"] = (kc + 10 * org_rate) / (nc + 10)
# B3 «solo importe» (log VE, si falta log PBL): logística de una variable en el entrenamiento
def x_imp(ix):
    ve = d.loc[ix, "ve_exp"]; pbl = d.loc[ix, "pbl"]
    return np.log10(np.maximum(ve.fillna(pbl).fillna(1), 1)).to_numpy()[:, None]
b3 = LogisticRegression(C=1e6).fit(x_imp(TR), ytr)
res["b3"] = b3.predict_proba(x_imp(idx))[:, 1]
inf["P0"] = P0
log("predicciones centrales")

# --- 7.1 intervalos: simulación de las posteriores + λ fijado en recalibración ---------------------------------------
S = 200
if not SIN_INT:
    rng = np.random.default_rng(SEM)
    sub_idx = idx
    n_s = len(sub_idx)
    ETA = np.empty((n_s, S), dtype=np.float32)
    Fbase = F.loc[sub_idx].copy()
    ii = sub_idx.to_numpy()
    # Una extracción por parámetro: los lotes de la misma celda-año (claves del nivel + T) comparten la del parámetro de su
    # celda (corrección del 26-oct, vista en la factibilidad: extraer por lote rompía la correlación dentro de la celda).
    # Sin clave (lote sin órgano), cada lote es su propio parámetro.
    GRP = {}
    for niv in NIV_VAR:
        cl = d.loc[sub_idx, CLAVES[niv] + ["T"]].astype(str)
        nulo = d.loc[sub_idx, CLAVES[niv]].isna().any(axis=1).to_numpy()
        g, _ = pd.factorize(cl.agg("|".join, axis=1))
        g = g.astype(np.int64); g[nulo] = g.max() + 1 + np.arange(nulo.sum())
        first = pd.Series(np.arange(len(g))).groupby(g).first()
        GRP[niv] = (g, first.index.to_numpy(), first.to_numpy())
    for s in range(S):
        Fs = Fbase.copy()
        for h in HS:
            for niv in NIV_VAR:
                a, b = alfa_beta[(h, niv)]
                g, gids, rep_row = GRP[niv]
                ag = np.maximum(a[ii][rep_row], 1e-9); bg = np.maximum(b[ii][rep_row], 1e-9)
                dg = np.empty(g.max() + 1); dg[gids] = rng.beta(ag, bg)
                Fs[f"eb_{h}_{niv}"] = logit(np.clip(dg[g], 1e-4, 1 - 1e-4))
        ps = calibrar(praw(sub_idx, Fs, Ccat.loc[sub_idx]))
        ETA[:, s] = logit(np.clip(ps, 1e-6, 1 - 1e-6))
        if s % 50 == 0:
            log("simulación", s)
    mu = logit(np.clip(res["p"].to_numpy(), 1e-6, 1 - 1e-6)).astype(np.float32)
    DEV = ETA - mu[:, None]
    np.save(f"{OUT}/dev.npy", DEV)
    # celdas órgano × CPV con >= 10 lotes por ventana
    def cobertura(ventana, lam, R=20, semilla=SEM, guardar=False):
        m_ = (res["ventana"] == ventana).to_numpy() & res["organo_clave"].notna().to_numpy()
        cel = (res.loc[m_, "organo_clave"].astype(str) + "|" + res.loc[m_, "cpv_cel"].astype(str))
        vc = cel.value_counts(); buenas = vc[vc >= 10].index
        sel = cel.isin(buenas).to_numpy()
        pos = np.where(m_)[0][sel]
        if len(pos) == 0:
            return None, 0
        codes, uniq = pd.factorize(cel[sel])
        P_ = expit(mu[pos][:, None] + lam * DEV[pos, :])          # n × S
        obs = np.bincount(codes, weights=res["y"].to_numpy()[pos], minlength=len(uniq))
        rg = np.random.default_rng(semilla)
        sims = np.empty((len(uniq), S * R), dtype=np.float32)
        for r in range(R):
            B = (rg.random(P_.shape, dtype=np.float32) < P_).astype(np.float32)
            for s in range(S):
                sims[:, r * S + s] = np.bincount(codes, weights=B[:, s], minlength=len(uniq))
        lo = np.quantile(sims, 0.05, axis=1); hi = np.quantile(sims, 0.95, axis=1)
        if guardar:
            pd.DataFrame({"celda": np.asarray(uniq), "n": np.bincount(codes, minlength=len(uniq)), "observados": obs,
                          "lo": lo, "hi": hi}).to_csv(f"{OUT}/celdas_{ventana}.csv", index=False)
        return float(np.mean((obs >= lo) & (obs <= hi))), int(len(uniq))
    grid = np.arange(0.25, 6.01, 0.25)
    LAM = None; cov_grid = {}
    for lam in grid:
        c, ncel = cobertura("recal", lam)
        cov_grid[float(lam)] = c
        if c is not None and c >= 0.90 and LAM is None:
            LAM = float(lam)
            break
    if LAM is None:
        LAM = float(grid[-1])
    cov_test, ncel_test = cobertura("test", LAM, guardar=True)
    cov_recal, ncel_recal = cobertura("recal", LAM, guardar=True)
    inf["intervalos"] = {"S": S, "lambda": LAM, "cobertura_recal": cov_recal, "celdas_recal": ncel_recal,
                         "cobertura_test": cov_test, "celdas_test": ncel_test, "rejilla": cov_grid}
    q = expit(mu[:, None] + LAM * DEV)
    res["p_lo"] = np.quantile(q, 0.05, axis=1); res["p_hi"] = np.quantile(q, 0.95, axis=1)
    log("intervalos", inf["intervalos"])

# --- 7.2 explicación por caso (prueba y predicción) ----------------------------------------------------------------
ex_idx = res.index[(res["ventana"].isin(["test", "pred"])).to_numpy()]
if ELEG == "H":
    X = disenno_H(ex_idx)
    beta = mH.coef_[0]
    Xm = np.concatenate([np.zeros(len(NUM)), disenno_H(d.index[TR][:200000])[:, len(NUM):].mean(axis=0)])
    contrib = (X - Xm) * beta
    grupos = sorted(set(GRUPO_H))
    G_ = np.column_stack([contrib[:, [i for i, g in enumerate(GRUPO_H) if g == gname]].sum(axis=1) for gname in grupos])
    base = mH.intercept_[0] + float(Xm @ beta)
    chk = np.max(np.abs(base + G_.sum(axis=1) - mH.decision_function(X)))
    inf["explicacion"] = {"metodo": "contribucion al logit (H)", "comprobacion_suma_max_abs": float(chk)}
else:
    grupos = sorted(set(list(GR.values()) + list(GRC.values())))
    Fx = F.loc[ex_idx]; Cx = Ccat.loc[ex_idx]
    p0 = p_G(ex_idx, Fx, Cx)
    G_ = []
    ref_num = F.loc[TR, NUM].median(); ref_cat = {c: Ccat.loc[TR, c].mode().iloc[0] for c in CATS}
    for gname in grupos:
        F2 = Fx.copy(); C2 = Cx.copy()
        for c in NUM:
            if GR[c] == gname: F2[c] = ref_num[c]
        for c in CATS:
            if GRC[c] == gname: C2[c] = ref_cat[c]
        G_.append(p0 - p_G(ex_idx, F2, C2))
    G_ = np.column_stack(G_)
    inf["explicacion"] = {"metodo": "oclusión por grupo (G): Δp bruta frente al valor de referencia del entrenamiento"}
orden = np.argsort(-np.abs(G_), axis=1)[:, :3]
for j in range(3):
    res.loc[ex_idx, f"exp{j + 1}"] = [grupos[orden[i, j]] for i in range(len(ex_idx))]
    res.loc[ex_idx, f"exp{j + 1}_v"] = G_[np.arange(len(ex_idx)), orden[:, j]]
log("explicación")

# --- 7.3 deriva -------------------------------------------------------------------------------------------------------
def psi(a, b, bins):
    pa = np.histogram(a, bins)[0] / max(len(a), 1); pb = np.histogram(b, bins)[0] / max(len(b), 1)
    pa = np.clip(pa, 1e-4, None); pb = np.clip(pb, 1e-4, None)
    return float(np.sum((pb - pa) * np.log(pb / pa)))
der = {}
for c in [x for x in NUM if not x.startswith("ln_")]:
    a = F.loc[TR, c].to_numpy(); b = F.loc[TE, c].to_numpy()
    qs = np.unique(np.quantile(a, np.linspace(0, 1, 11)))
    if len(qs) < 3:
        continue
    qs[0], qs[-1] = -np.inf, np.inf
    der[c] = psi(a, b, qs)
for c in CATS:
    a = Ccat.loc[TR, c].value_counts(normalize=True); b = Ccat.loc[TE, c].value_counts(normalize=True)
    lv = sorted(set(a.index) | set(b.index))
    pa = np.clip(a.reindex(lv).fillna(0).to_numpy(), 1e-4, None); pb = np.clip(b.reindex(lv).fillna(0).to_numpy(), 1e-4, None)
    der[c] = float(np.sum((pb - pa) * np.log(pb / pa)))
ptr = praw(d.index[TR][:300000]); pte = res.loc[res["ventana"] == "test", "p_raw"].to_numpy()
qs = np.unique(np.quantile(ptr, np.linspace(0, 1, 11))); qs[0], qs[-1] = -np.inf, np.inf
der["p_bruta"] = psi(ptr, pte, qs)
mes = res.loc[res["ventana"] == "test"].assign(mes=lambda x: x["fecha_pub"].dt.strftime("%Y-%m")).groupby("mes").agg(
    n=("y", "size"), observada=("y", "mean"), predicha=("p", "mean"))
inf["deriva"] = {"psi": dict(sorted(der.items(), key=lambda kv: -kv[1])), "alarmas_psi_ge_025": [k for k, v in der.items() if v >= 0.25],
                 "calibracion_por_mes": mes.reset_index().to_dict(orient="records")}
log("deriva")

# --- guardado ---------------------------------------------------------------------------------------------------------
res.drop(columns=["organo_nom"]).to_csv(f"{OUT}/predicciones.csv", index=False)
res.loc[res["ventana"] == "pred"].to_csv(f"{OUT}/prediccion_en_plazo.csv", index=False)
with open(f"{OUT}/modelos.pkl", "wb") as fh:
    pickle.dump({"H": mH, "G": mG, "elegido": ELEG, "calibrador": CAL, "MU": MU, "SD": SD, "NUM": NUM, "CATS": CATS,
                 "CAT_NIV": CAT_NIV, "EB_M": EB_M, "P0": P0}, fh)
coef = {"intercepto": float(mH.intercept_[0]), "coeficientes": dict(zip(COLS_H, map(float, mH.coef_[0]))),
        "medias": dict(zip(NUM, map(float, MU))), "escalas": dict(zip(NUM, map(float, SD)))}
json.dump(coef, open(f"{OUT}/modelo_H.json", "w"), indent=0, ensure_ascii=False)
inf["segundos"] = round(time.time() - T0, 1)
inf["ru_maxrss_mb"] = round(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024)
os.makedirs("/w/salidas", exist_ok=True)
json.dump(inf, open(f"/w/salidas/20_{H}{'_muestra' if MUESTRA else ''}{'_pert' if 'perturbado' in FICH else ''}.json", "w"),
          indent=1, ensure_ascii=False, default=str)
log("fin", inf["segundos"], "s", inf["ru_maxrss_mb"], "MB")