package_code_imports/v5_scripts/evaluate_v5_matched_baselines.py

from __future__ import annotations

import argparse
import json
import math
import sys
from collections import defaultdict
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd
from sklearn.ensemble import HistGradientBoostingClassifier
from sklearn.isotonic import IsotonicRegression
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import average_precision_score, roc_auc_score
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler


SCRIPT_DIR = Path(__file__).resolve().parent
ROOT_DIR = SCRIPT_DIR.parents[1]
V4P4_SCRIPT_DIR = ROOT_DIR / "v4p4_world_model" / "scripts"
V3P5_SCRIPT_DIR = ROOT_DIR / "v3p5_static" / "scripts"
for path in (SCRIPT_DIR, V4P4_SCRIPT_DIR, V3P5_SCRIPT_DIR):
    if str(path) not in sys.path:
        sys.path.insert(0, str(path))

from action_ontology_v5 import build_action_ontology, load_json  # noqa: E402
from build_v5_action_tensors import (  # noqa: E402
    build_action_arrays,
    build_target_trial_labels,
    medication_flags_from_stage0,
    resolve_split_path,
    validate_medication_flags,
)
from run_service_process_target_trial_v5 import pre_action_features  # noqa: E402
from scan_target_trial_support_v5 import era_tokens_np, future_event_within, load_or_build_action_arrays, unique_landmark_filter  # noqa: E402
from target_trial_estimators_v5 import effective_sample_size  # noqa: E402

try:
    from lightgbm import LGBMClassifier
except Exception:  # pragma: no cover
    LGBMClassifier = None


CAUSE_NAMES = {0: "next_contact", 1: "death", 2: "disengagement"}
PWE_BIN_EDGES_DAYS = (7.0, 14.0, 30.0, 60.0, 90.0, 365.0, math.inf)
RISK_ENDPOINTS = {
    "death": {"kind": "pwe", "cause_id": 1},
    "disengagement": {"kind": "pwe", "cause_id": 2},
    "primary_referral_relapse": {"kind": "event", "event_index": 0},
    "service_escalation": {"kind": "event", "event_index": 3},
    "clinical_deterioration": {"kind": "event", "event_index": 4},
    "high_acuity_state": {"kind": "event", "event_index": 5},
}
RISK_THRESHOLD_GRID = {
    "death": [0.005, 0.01, 0.02, 0.03, 0.05],
    "primary_referral_relapse": [0.01, 0.03, 0.05, 0.10, 0.15],
    "disengagement": [0.01, 0.02, 0.05, 0.08, 0.10],
    "service_escalation": [0.03, 0.05, 0.10, 0.15, 0.20],
    "clinical_deterioration": [0.30, 0.40, 0.50, 0.55, 0.60],
    "high_acuity_state": [0.10, 0.20, 0.30, 0.40, 0.50],
    "default": [0.01, 0.03, 0.05, 0.10, 0.15],
}


def load_npz_to_memory(path: Path) -> dict[str, np.ndarray]:
    with np.load(path, allow_pickle=True) as z:
        return {key: z[key] for key in z.files}


def json_ready(value: Any) -> Any:
    if isinstance(value, Path):
        return str(value)
    if isinstance(value, dict):
        return {str(k): json_ready(v) for k, v in value.items()}
    if isinstance(value, (list, tuple)):
        return [json_ready(v) for v in value]
    if isinstance(value, np.generic):
        return value.item()
    if isinstance(value, float) and not math.isfinite(value):
        return None
    return value


def service_era_key(service: np.ndarray, years: np.ndarray) -> np.ndarray:
    return service.astype(np.int64) * 10 + era_tokens_np(years).astype(np.int64)


def clip_prob(p: np.ndarray, eps: float = 1.0e-7) -> np.ndarray:
    return np.clip(np.asarray(p, dtype=np.float64).reshape(-1), eps, 1.0 - eps)


def weighted_mean(x: np.ndarray, sample_weight: np.ndarray | None = None) -> float:
    x = np.asarray(x, dtype=np.float64).reshape(-1)
    if not x.size:
        return math.nan
    if sample_weight is None:
        return float(np.mean(x))
    w = np.asarray(sample_weight, dtype=np.float64).reshape(-1)
    ok = np.isfinite(x) & np.isfinite(w) & (w > 0)
    if not ok.any():
        return math.nan
    return float(np.sum(w[ok] * x[ok]) / np.sum(w[ok]))


def calibration_intercept_slope(y: np.ndarray, p: np.ndarray, sample_weight: np.ndarray | None = None) -> tuple[float, float]:
    y = np.asarray(y, dtype=np.int8).reshape(-1)
    p = clip_prob(p)
    if y.size == 0 or np.unique(y).size < 2:
        return math.nan, math.nan
    try:
        x = np.log(p / (1.0 - p)).reshape(-1, 1)
        clf = LogisticRegression(C=1.0e6, solver="lbfgs", max_iter=500)
        clf.fit(x, y, sample_weight=sample_weight)
        return float(clf.intercept_[0]), float(clf.coef_[0, 0])
    except Exception:
        return math.nan, math.nan


def integrated_calibration_index(y: np.ndarray, p: np.ndarray, sample_weight: np.ndarray | None = None) -> float:
    y = np.asarray(y, dtype=np.int8).reshape(-1)
    p = clip_prob(p)
    w = None if sample_weight is None else np.asarray(sample_weight, dtype=np.float64).reshape(-1)
    if w is not None:
        ok = np.isfinite(w) & (w > 0)
        y = y[ok]
        p = p[ok]
        w = w[ok]
    if y.size == 0 or np.unique(y).size < 2:
        return math.nan
    try:
        model = IsotonicRegression(y_min=0.0, y_max=1.0, out_of_bounds="clip")
        calibrated = model.fit_transform(p, y, sample_weight=w)
        return weighted_mean(np.abs(calibrated - p), w)
    except Exception:
        return math.nan


def weighted_binary_metrics(
    y: np.ndarray,
    p: np.ndarray,
    sample_weight: np.ndarray | None = None,
    include_calibration_model: bool = True,
) -> dict[str, float]:
    y = np.asarray(y, dtype=np.int8).reshape(-1)
    p = clip_prob(p)
    w = None if sample_weight is None else np.asarray(sample_weight, dtype=np.float64).reshape(-1)
    if w is not None:
        ok = np.isfinite(w) & (w > 0)
        y = y[ok]
        p = p[ok]
        w = w[ok]
    if y.size == 0:
        return {
            "auc": math.nan,
            "average_precision": math.nan,
            "brier": math.nan,
            "ece": math.nan,
            "ici": math.nan,
            "mean_predicted": math.nan,
            "observed_rate": math.nan,
            "calibration_intercept": math.nan,
            "calibration_slope": math.nan,
        }
    auc = float(roc_auc_score(y, p, sample_weight=w)) if np.unique(y).size >= 2 else math.nan
    ap = float(average_precision_score(y, p, sample_weight=w)) if np.unique(y).size >= 2 else math.nan
    order = np.argsort(p)
    ece = 0.0
    ece_den = float(y.size) if w is None else float(w.sum())
    for idx in np.array_split(order, min(10, max(1, y.size))):
        if idx.size:
            bin_w = None if w is None else w[idx]
            bin_weight = float(idx.size) if w is None else float(bin_w.sum())
            ece += abs(weighted_mean(p[idx], bin_w) - weighted_mean(y[idx], bin_w)) * bin_weight / ece_den
    intercept, slope = calibration_intercept_slope(y, p, w) if include_calibration_model else (math.nan, math.nan)
    return {
        "auc": auc,
        "average_precision": ap,
        "brier": weighted_mean((p - y) ** 2, w),
        "ece": float(ece),
        "ici": integrated_calibration_index(y, p, w) if include_calibration_model else math.nan,
        "mean_predicted": weighted_mean(p, w),
        "observed_rate": weighted_mean(y, w),
        "calibration_intercept": intercept,
        "calibration_slope": slope,
    }


def net_benefit(y: np.ndarray, p: np.ndarray, threshold: float, sample_weight: np.ndarray | None = None) -> float:
    y = np.asarray(y, dtype=np.int8).reshape(-1)
    p = clip_prob(p)
    if y.size == 0 or threshold <= 0.0 or threshold >= 1.0:
        return math.nan
    w = np.ones(y.size, dtype=np.float64) if sample_weight is None else np.asarray(sample_weight, dtype=np.float64).reshape(-1)
    ok = np.isfinite(w) & (w > 0)
    y = y[ok]
    p = p[ok]
    w = w[ok]
    if y.size == 0 or float(w.sum()) <= 0.0:
        return math.nan
    pred_pos = p >= float(threshold)
    den = float(w.sum())
    tp = float(np.sum(w * pred_pos * (y == 1)) / den)
    fp = float(np.sum(w * pred_pos * (y == 0)) / den)
    return float(tp - fp * threshold / (1.0 - threshold))


def calibration_curve_rows(
    *,
    split: str,
    endpoint: str,
    horizon_days: float,
    model: str,
    y: np.ndarray,
    p: np.ndarray,
    n_bins: int = 10,
) -> list[dict[str, Any]]:
    y = np.asarray(y, dtype=np.int8).reshape(-1)
    p = clip_prob(p)
    if y.size == 0:
        return []
    rows = []
    for bin_idx, idx in enumerate(np.array_split(np.argsort(p), min(n_bins, max(1, y.size))), start=1):
        if idx.size == 0:
            continue
        rows.append(
            {
                "split": split,
                "endpoint": endpoint,
                "horizon_days": float(horizon_days),
                "model": model,
                "metric_family": "calibration_curve",
                "calibration_bin": int(bin_idx),
                "n": int(idx.size),
                "events": int(y[idx].sum()),
                "mean_predicted": float(p[idx].mean()),
                "observed_rate": float(y[idx].mean()),
                "predicted_min": float(p[idx].min()),
                "predicted_max": float(p[idx].max()),
            }
        )
    return rows


def bootstrap_metric_ci(
    y: np.ndarray,
    p: np.ndarray,
    patient_ids: np.ndarray,
    n_bootstrap: int,
    seed: int,
) -> dict[str, float]:
    if n_bootstrap <= 0 or y.size == 0:
        return {}
    rng = np.random.default_rng(seed)
    clusters, inv = np.unique(patient_ids.astype(str), return_inverse=True)
    tracked: dict[str, list[float]] = defaultdict(list)
    for _ in range(n_bootstrap):
        counts = rng.multinomial(clusters.size, np.full(clusters.size, 1.0 / clusters.size))
        w = counts[inv].astype(np.float64)
        metrics = weighted_binary_metrics(y, p, sample_weight=w, include_calibration_model=False)
        for key, value in metrics.items():
            if math.isfinite(float(value)):
                tracked[key].append(float(value))
    out: dict[str, float] = {"bootstrap_n": int(n_bootstrap), "bootstrap_clusters": int(clusters.size)}
    for key, vals in tracked.items():
        arr = np.asarray(vals, dtype=np.float64)
        out[f"{key}_ci_low"] = float(np.nanquantile(arr, 0.025)) if arr.size else math.nan
        out[f"{key}_ci_high"] = float(np.nanquantile(arr, 0.975)) if arr.size else math.nan
    return out


def bootstrap_net_benefit_ci(
    y: np.ndarray,
    p: np.ndarray,
    patient_ids: np.ndarray,
    threshold: float,
    n_bootstrap: int,
    seed: int,
) -> dict[str, float]:
    if n_bootstrap <= 0 or y.size == 0:
        return {}
    rng = np.random.default_rng(seed)
    clusters, inv = np.unique(patient_ids.astype(str), return_inverse=True)
    vals = []
    for _ in range(n_bootstrap):
        counts = rng.multinomial(clusters.size, np.full(clusters.size, 1.0 / clusters.size))
        vals.append(net_benefit(y, p, threshold, sample_weight=counts[inv].astype(np.float64)))
    arr = np.asarray(vals, dtype=np.float64)
    return {
        "net_benefit_ci_low": float(np.nanquantile(arr, 0.025)) if arr.size else math.nan,
        "net_benefit_ci_high": float(np.nanquantile(arr, 0.975)) if arr.size else math.nan,
    }


def active_next_mask(arrays: dict[str, np.ndarray]) -> np.ndarray:
    return arrays["valid_mask"][:, :-1].astype(bool) & arrays["valid_mask"][:, 1:].astype(bool)


def pwe_targets(arrays: dict[str, np.ndarray]) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    terminal = arrays["terminal_label"][:, 1:].astype(np.int64)
    service = arrays["service_state"][:, 1:].astype(np.int64)
    valid = active_next_mask(arrays)
    cause = np.zeros_like(terminal, dtype=np.int64)
    cause[(terminal == 1) | (service == 7)] = 1
    cause[(terminal == 2) | (terminal == 3) | (service == 6)] = 2
    censored = (terminal == 4) | (service == 5)
    days = np.expm1(arrays["delta_t_next_log"][:, :-1].astype(np.float64)).clip(min=1.0e-6)
    return cause, censored.astype(bool), valid, days


def build_conditional_specs(meta: dict[str, Any], vocab: dict[str, Any]) -> tuple[dict[str, Any], ...]:
    cat_cols = list(meta.get("cat_cols", []))

    def field_idx(name: str) -> int:
        return cat_cols.index(name) if name in cat_cols else -1

    specs = []
    tables = [
        ("是否转诊", ["是"], ["转诊类型", "转诊原因", "转诊至机构"]),
        ("治疗方式", ["住院"], ["本次入院形式", "本次住院名称", "本次入院日期", "末次出院日期"]),
    ]
    for parent, active_values, children in tables:
        parent_idx = field_idx(parent)
        child_indices = [field_idx(child) for child in children if field_idx(child) >= 0]
        active_ids = [int(vocab[f"{parent}={value}"]) for value in active_values if f"{parent}={value}" in vocab]
        if parent_idx >= 0 and child_indices and active_ids:
            specs.append({"parent_idx": parent_idx, "active_value_ids": active_ids, "child_indices": child_indices})
    return tuple(specs)


def conditional_status_np(cat_value_ids: np.ndarray, missing_ids: np.ndarray, specs: tuple[dict[str, Any], ...]) -> np.ndarray:
    status = np.zeros_like(cat_value_ids, dtype=np.int64)
    if not specs:
        return status
    observed = missing_ids.astype(np.int64) == 0
    unknown = 1
    for spec in specs:
        parent_idx = int(spec["parent_idx"])
        child_indices = [int(x) for x in spec["child_indices"]]
        active_ids = np.asarray([int(x) for x in spec["active_value_ids"]], dtype=cat_value_ids.dtype)
        if parent_idx < 0 or parent_idx >= cat_value_ids.shape[-1] or not child_indices or active_ids.size == 0:
            continue
        parent_val = cat_value_ids[:, :, parent_idx]
        parent_observed = observed[:, :, parent_idx] & (parent_val != unknown)
        parent_active = np.isin(parent_val, active_ids) & parent_observed
        parent_unknown = ~parent_observed
        for child_idx in child_indices:
            if child_idx < 0 or child_idx >= cat_value_ids.shape[-1]:
                continue
            child_val = cat_value_ids[:, :, child_idx]
            child_observed = observed[:, :, child_idx] & (child_val != unknown)
            code = np.zeros_like(parent_val, dtype=np.int64)
            code[parent_unknown] = 2
            code[(~parent_unknown) & (~parent_active) & (~child_observed)] = 1
            code[(~parent_unknown) & (~parent_active) & child_observed] = 3
            status[:, :, child_idx] = code
    return status


def mapped_missing_targets(
    arrays: dict[str, np.ndarray],
    meta: dict[str, Any],
    vocab: dict[str, Any],
    *,
    start: int,
    stop: int | None,
) -> np.ndarray:
    """Match the v4/v5 loss missingness target: 5 classes incl. structural/no-clinical."""

    raw = arrays["missing_ids"][:, start:stop, :].astype(np.int64)
    mapped = np.full_like(raw, 3, dtype=np.int64)
    mapped[raw == 0] = 0
    mapped[raw == 1] = 1
    mapped[raw == 2] = 3
    specs = build_conditional_specs(meta, vocab)
    if specs:
        cat = arrays["cat_value_ids"][:, start:stop, :].astype(np.int64)
        conditional = conditional_status_np(cat, raw, specs)
        mapped[conditional == 1] = 2
    service = arrays["service_state"][:, start:stop].astype(np.int64)
    terminal = arrays["terminal_label"][:, start:stop].astype(np.int64)
    no_clinical = (terminal > 0) | (service >= 5)
    mapped = np.where(no_clinical[:, :, None], 4, mapped)
    return mapped.clip(0, 4)


def exposure_by_bin(days: np.ndarray) -> np.ndarray:
    upper = np.asarray(PWE_BIN_EDGES_DAYS[:-1], dtype=np.float64)
    lower = np.asarray((0.0, 7.0, 14.0, 30.0, 60.0, 90.0), dtype=np.float64)
    finite = (np.minimum(days[..., None], upper) - lower).clip(min=0.0)
    tail = (days - 365.0).clip(min=0.0)[..., None]
    return np.concatenate([finite, tail], axis=-1)


def pwe_bin_index(days: np.ndarray) -> np.ndarray:
    upper = np.asarray(PWE_BIN_EDGES_DAYS[:-1], dtype=np.float64)
    return np.searchsorted(upper, days, side="left").clip(max=len(PWE_BIN_EDGES_DAYS) - 1)


def fit_semi_markov(train: dict[str, np.ndarray], smoothing: float) -> dict[str, Any]:
    cause, censored, valid, days = pwe_targets(train)
    cur_state = train["service_state"][:, :-1].astype(np.int64)
    next_state = train["service_state"][:, 1:].astype(np.int64)
    key = service_era_key(cur_state, train["visit_year"][:, :-1])
    n_states = int(max(train["service_state"].max(), 7) + 1)
    n_bins = len(PWE_BIN_EDGES_DAYS)
    trans_counts: dict[int, np.ndarray] = {}
    exposure: dict[int, np.ndarray] = {}
    events: dict[int, np.ndarray] = {}
    for k in np.unique(key[valid]):
        m = valid & (key == k)
        counts = np.bincount(next_state[m].clip(0, n_states - 1), minlength=n_states).astype(np.float64) + smoothing
        trans_counts[int(k)] = counts / counts.sum()
        ex = exposure_by_bin(days[m]).sum(axis=0)
        ev = np.zeros((3, n_bins), dtype=np.float64)
        bin_idx = pwe_bin_index(days[m])
        c = cause[m]
        z = censored[m]
        for cause_id in range(3):
            np.add.at(ev[cause_id], bin_idx[(~z) & (c == cause_id)], 1.0)
        exposure[int(k)] = ex + smoothing
        events[int(k)] = ev + smoothing
    global_counts = np.bincount(next_state[valid].clip(0, n_states - 1), minlength=n_states).astype(np.float64) + smoothing
    global_trans = global_counts / global_counts.sum()
    global_ex = exposure_by_bin(days[valid]).sum(axis=0) + smoothing
    global_ev = np.zeros((3, n_bins), dtype=np.float64) + smoothing
    bin_idx = pwe_bin_index(days[valid])
    for cause_id in range(3):
        np.add.at(global_ev[cause_id], bin_idx[(~censored[valid]) & (cause[valid] == cause_id)], 1.0)
    return {
        "trans": trans_counts,
        "exposure": exposure,
        "events": events,
        "global_trans": global_trans,
        "global_exposure": global_ex,
        "global_events": global_ev,
        "n_states": n_states,
        "n_bins": n_bins,
        "smoothing": smoothing,
    }


def semi_markov_metrics(model: dict[str, Any], arrays: dict[str, np.ndarray], split: str) -> list[dict[str, Any]]:
    cause, censored, valid, days = pwe_targets(arrays)
    cur_state = arrays["service_state"][:, :-1].astype(np.int64)
    next_state = arrays["service_state"][:, 1:].astype(np.int64).clip(0, model["n_states"] - 1)
    key = service_era_key(cur_state, arrays["visit_year"][:, :-1])
    rows = []
    key_flat = key[valid].astype(np.int64)
    y = next_state[valid]
    pred = np.empty(y.shape, dtype=np.int64)
    ce_sum = 0.0
    for k in np.unique(key_flat):
        idx = np.flatnonzero(key_flat == int(k))
        probs_k = model["trans"].get(int(k), model["global_trans"])
        pred[idx] = int(np.argmax(probs_k))
        ce_sum += float(-np.log(np.clip(probs_k[y[idx]], 1.0e-12, 1.0)).sum())
    ce = ce_sum / y.size if y.size else math.nan
    conf = np.zeros((model["n_states"], model["n_states"]), dtype=np.int64)
    if y.size:
        np.add.at(conf, (y, pred), 1)
    f1 = []
    for cls in range(model["n_states"]):
        if conf[cls].sum() <= 0:
            continue
        tp = conf[cls, cls]
        fp = conf[:, cls].sum() - tp
        fn = conf[cls].sum() - tp
        f1.append(2 * tp / max(1, 2 * tp + fp + fn))
    rows.append(
        {
            "split": split,
            "model": "Semi-Markov Care-Process",
            "metric_family": "service_transition",
            "n": int(y.size),
            "cross_entropy": float(ce),
            "accuracy": float((pred == y).mean()) if y.size else math.nan,
            "macro_f1": float(np.mean(f1)) if f1 else math.nan,
        }
    )
    cause_flat = cause[valid].astype(np.int64)
    cens_flat = censored[valid].astype(bool)
    days_flat = days[valid].astype(np.float64)
    nll_sum = 0.0
    nll_n = int(days_flat.size)
    bin_flat = pwe_bin_index(days_flat)
    for k in np.unique(key_flat):
        idx = np.flatnonzero(key_flat == int(k))
        rates = model["events"].get(int(k), model["global_events"]) / model["exposure"].get(int(k), model["global_exposure"])[None, :]
        total_rate = rates.sum(axis=0)
        ex = exposure_by_bin(days_flat[idx])
        ll = -np.sum(ex * total_rate[None, :], axis=1)
        event_idx = idx[~cens_flat[idx]]
        if event_idx.size:
            rel = np.flatnonzero(~cens_flat[idx])
            ll[rel] += np.log(np.clip(rates[cause_flat[event_idx], bin_flat[event_idx]], 1.0e-12, None))
        nll_sum += float((-ll).sum())
    rows.append(
        {
            "split": split,
            "model": "Semi-Markov Care-Process",
            "metric_family": "pwe_time_to_next_process",
            "n": nll_n,
            "pwe_nll": nll_sum / nll_n if nll_n else math.nan,
        }
    )
    return rows


def fit_empirical_grammar(train: dict[str, np.ndarray], meta: dict[str, Any], vocab_map: dict[str, Any], smoothing: float) -> dict[str, Any]:
    valid = active_next_mask(train)
    next_contact = valid & (train["service_state"][:, 1:] < 5)
    key = service_era_key(train["service_state"][:, :-1], train["visit_year"][:, :-1])
    n_missing = 5
    missing_target = mapped_missing_targets(train, meta, vocab_map, start=1, stop=None)
    vocab = int(meta.get("cat_vocab_size", int(train["cat_value_ids"].max()) + 1))
    cbe_dim = int(meta.get("cbe_dim", 5))
    out: dict[str, Any] = {"missing": {}, "cat": {}, "numeric": {}, "ordinal": {}, "global": {}, "target_contract": "v4_v5_mapped_5class_missingness"}
    for k in np.unique(key[valid]):
        m = valid & (key == k)
        miss = missing_target[m].reshape(-1).astype(int)
        out["missing"][int(k)] = (np.bincount(miss.clip(0, n_missing - 1), minlength=n_missing) + smoothing)
        out["missing"][int(k)] /= out["missing"][int(k)].sum()
    miss_global = missing_target[valid].reshape(-1).astype(int)
    out["global"]["missing"] = (np.bincount(miss_global.clip(0, n_missing - 1), minlength=n_missing) + smoothing)
    out["global"]["missing"] /= out["global"]["missing"].sum()
    for field_idx in range(train["cat_value_ids"].shape[2]):
        out["cat"][field_idx] = {}
        observed = next_contact & (missing_target[:, :, field_idx] == 0)
        vals_global = train["cat_value_ids"][:, 1:, field_idx][observed].astype(int)
        g = np.bincount(vals_global.clip(0, vocab - 1), minlength=vocab).astype(np.float64) + smoothing
        out["global"][f"cat_{field_idx}"] = g / g.sum()
        for k in np.unique(key[observed]):
            vals = train["cat_value_ids"][:, 1:, field_idx][observed & (key == k)].astype(int)
            counts = np.bincount(vals.clip(0, vocab - 1), minlength=vocab).astype(np.float64) + smoothing
            out["cat"][field_idx][int(k)] = counts / counts.sum()
    for field_idx in range(train["numeric_values"].shape[2]):
        out["numeric"][field_idx] = {}
        observed = next_contact & train["numeric_mask"][:, 1:, field_idx]
        vals = train["numeric_values"][:, 1:, field_idx][observed].astype(float)
        out["global"][f"num_{field_idx}"] = float(np.nanmean(vals)) if vals.size else 0.0
        for k in np.unique(key[observed]):
            sub = train["numeric_values"][:, 1:, field_idx][observed & (key == k)].astype(float)
            out["numeric"][field_idx][int(k)] = float(np.nanmean(sub)) if sub.size else out["global"][f"num_{field_idx}"]
    for field_idx in range(train["ordinal_cbe"].shape[2]):
        out["ordinal"][field_idx] = {}
        observed = next_contact & train["ordinal_mask"][:, 1:, field_idx]
        levels = train["ordinal_cbe"][:, 1:, field_idx, :][observed].sum(axis=1).astype(int)
        counts = np.bincount(levels.clip(0, cbe_dim), minlength=cbe_dim + 1).astype(np.float64) + smoothing
        out["global"][f"ord_{field_idx}"] = int(np.argmax(counts))
        for k in np.unique(key[observed]):
            sub = train["ordinal_cbe"][:, 1:, field_idx, :][observed & (key == k)].sum(axis=1).astype(int)
            c = np.bincount(sub.clip(0, cbe_dim), minlength=cbe_dim + 1).astype(np.float64) + smoothing
            out["ordinal"][field_idx][int(k)] = int(np.argmax(c))
    return out


def empirical_grammar_metrics(model: dict[str, Any], arrays: dict[str, np.ndarray], split: str, meta: dict[str, Any], vocab: dict[str, Any]) -> list[dict[str, Any]]:
    valid = active_next_mask(arrays)
    next_contact = valid & (arrays["service_state"][:, 1:] < 5)
    key = service_era_key(arrays["service_state"][:, :-1], arrays["visit_year"][:, :-1])
    rows = []
    missing_target = mapped_missing_targets(arrays, meta, vocab, start=1, stop=None)
    miss_y = missing_target[valid].reshape(-1).astype(int)
    miss_keys = np.repeat(key[valid], arrays["missing_ids"].shape[2])
    miss_correct = 0
    miss_brier_sum = 0.0
    for k in np.unique(miss_keys):
        idx = np.flatnonzero(miss_keys == int(k))
        probs = model["missing"].get(int(k), model["global"]["missing"])
        yk = miss_y[idx].clip(0, probs.shape[0] - 1)
        miss_correct += int((int(np.argmax(probs)) == yk).sum())
        miss_brier_sum += float((np.sum(probs * probs) + 1.0) * yk.size - 2.0 * np.sum(probs[yk]))
    rows.append(
        {
            "split": split,
            "model": "Empirical state-year grammar",
            "family": "missingness",
            "n": int(miss_y.size),
            "accuracy": float(miss_correct / miss_y.size) if miss_y.size else math.nan,
            "multiclass_brier": float(miss_brier_sum / miss_y.size) if miss_y.size else math.nan,
            "target_contract": model.get("target_contract", "v4_v5_mapped_5class_missingness"),
        }
    )
    cat_ce = []
    cat_top1 = []
    cat_top3 = []
    cat_n = 0
    for field_idx in range(arrays["cat_value_ids"].shape[2]):
        observed = next_contact & (missing_target[:, :, field_idx] == 0)
        y = arrays["cat_value_ids"][:, 1:, field_idx][observed].astype(int)
        if y.size:
            cat_n += int(y.size)
            k_obs = key[observed].astype(np.int64)
            ce_sum = 0.0
            top1_sum = 0
            top3_sum = 0
            for k in np.unique(k_obs):
                idx = np.flatnonzero(k_obs == int(k))
                probs = model["cat"][field_idx].get(int(k), model["global"][f"cat_{field_idx}"])
                yk = y[idx].clip(0, probs.shape[0] - 1)
                ce_sum += float(-np.log(np.clip(probs[yk], 1.0e-12, 1.0)).sum())
                top1 = int(np.argmax(probs))
                top1_sum += int((yk == top1).sum())
                k_top = min(3, probs.shape[0])
                top3 = np.argpartition(-probs, kth=k_top - 1)[:k_top]
                top3_sum += int(np.isin(yk, top3).sum())
            cat_ce.append(ce_sum / y.size)
            cat_top1.append(top1_sum / y.size)
            cat_top3.append(top3_sum / y.size)
    rows.append(
        {
            "split": split,
            "model": "Empirical state-year grammar",
            "family": "categorical",
            "n": int(cat_n),
            "mean_cross_entropy": float(np.nanmean(cat_ce)) if cat_ce else math.nan,
            "mean_top1_accuracy": float(np.nanmean(cat_top1)) if cat_top1 else math.nan,
            "mean_top3_accuracy": float(np.nanmean(cat_top3)) if cat_top3 else math.nan,
        }
    )
    abs_err = []
    sq_err = []
    for field_idx in range(arrays["numeric_values"].shape[2]):
        observed = next_contact & arrays["numeric_mask"][:, 1:, field_idx]
        y = arrays["numeric_values"][:, 1:, field_idx][observed].astype(float)
        pred = np.asarray([model["numeric"][field_idx].get(int(k), model["global"][f"num_{field_idx}"]) for k in key[observed]], dtype=float) if y.size else np.asarray([])
        abs_err.append(np.abs(pred - y))
        sq_err.append((pred - y) ** 2)
    ae = np.concatenate(abs_err) if abs_err else np.asarray([])
    se = np.concatenate(sq_err) if sq_err else np.asarray([])
    rows.append(
        {
            "split": split,
            "model": "Empirical state-year grammar",
            "family": "numeric",
            "n": int(ae.size),
            "mae": float(ae.mean()) if ae.size else math.nan,
            "rmse": float(math.sqrt(float(se.mean()))) if se.size else math.nan,
        }
    )
    ord_abs = []
    for field_idx in range(arrays["ordinal_cbe"].shape[2]):
        observed = next_contact & arrays["ordinal_mask"][:, 1:, field_idx]
        y = arrays["ordinal_cbe"][:, 1:, field_idx, :][observed].sum(axis=1).astype(int)
        pred = np.asarray([model["ordinal"][field_idx].get(int(k), model["global"][f"ord_{field_idx}"]) for k in key[observed]], dtype=int) if y.size else np.asarray([])
        ord_abs.append(np.abs(pred - y))
    oa = np.concatenate(ord_abs) if ord_abs else np.asarray([])
    rows.append({"split": split, "model": "Empirical state-year grammar", "family": "ordinal", "n": int(oa.size), "mae": float(oa.mean()) if oa.size else math.nan})
    return rows


def horizon_label(arrays: dict[str, np.ndarray], endpoint: str, horizon: float) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    spec = RISK_ENDPOINTS[endpoint]
    valid = active_next_mask(arrays)
    row_idx, pos_idx = np.nonzero(valid)
    if endpoint in ("death", "disengagement"):
        cause, censored, _, days = pwe_targets(arrays)
        cause_id = int(spec["cause_id"])
        y_full = ((~censored) & (cause == cause_id) & (days <= float(horizon))).astype(np.int8)
        observed = valid & ~(censored & (days <= float(horizon)))
        row_idx, pos_idx = np.nonzero(observed)
        y = y_full[row_idx, pos_idx]
    else:
        y_full = future_event_within(arrays, int(spec["event_index"]), float(horizon)).astype(np.int8)
        valid_full = arrays["valid_mask"].astype(bool)
        times = arrays["time_since_start_days"].astype(float)
        seq_len = valid_full.shape[1]
        current_len = seq_len - 1
        any_future = np.zeros((valid_full.shape[0], current_len), dtype=bool)
        for offset in range(1, seq_len):
            cur_len = min(current_len, seq_len - offset)
            if cur_len <= 0:
                break
            dt = times[:, offset : offset + cur_len] - times[:, :cur_len]
            any_future[:, :cur_len] |= valid_full[:, offset : offset + cur_len] & (dt > 0.0) & (dt <= float(horizon))
        last_time = np.where(valid_full, times, -np.inf).max(axis=1, keepdims=True)
        followup = last_time - times[:, :-1]
        observed = valid & any_future & ((y_full[:, :-1].astype(bool)) | (followup >= float(horizon)))
        row_idx, pos_idx = np.nonzero(observed)
        y = y_full[:, :-1][row_idx, pos_idx]
    if row_idx.size and "patient_ids" in arrays:
        keep = unique_landmark_filter(arrays, row_idx, pos_idx)
        row_idx = row_idx[keep]
        pos_idx = pos_idx[keep]
        y = y[keep]
    return row_idx, pos_idx, y.astype(np.int8), arrays["patient_ids"][row_idx].astype(str) if "patient_ids" in arrays else row_idx.astype(str)


def fit_classifier(X: np.ndarray, y: np.ndarray, model_kind: str, seed: int, max_train: int, require_lightgbm: bool) -> tuple[Any, str]:
    rng = np.random.default_rng(seed)
    if max_train > 0 and X.shape[0] > max_train:
        idx = rng.choice(X.shape[0], size=max_train, replace=False)
        X_fit = X[idx]
        y_fit = y[idx]
    else:
        X_fit, y_fit = X, y
    if np.unique(y_fit).size < 2:
        rate = float(y_fit.mean()) if y_fit.size else 0.0
        return rate, "constant_rate"
    if model_kind == "clinical_landmark":
        model = make_pipeline(
            StandardScaler(with_mean=True, with_std=True),
            LogisticRegression(max_iter=1000, C=0.5, solver="lbfgs"),
        )
        model.fit(X_fit, y_fit)
        return model, "Clinical Landmark logistic (L2)"
    if LGBMClassifier is not None:
        model = LGBMClassifier(
            n_estimators=400,
            learning_rate=0.03,
            num_leaves=31,
            subsample=0.85,
            colsample_bytree=0.85,
            reg_lambda=1.0,
            random_state=seed,
            n_jobs=-1,
            verbose=-1,
        )
        model.fit(X_fit, y_fit)
        return model, "LightGBM Landmark"
    if require_lightgbm:
        raise RuntimeError("LightGBM backend is required but lightgbm could not be imported.")
    model = HistGradientBoostingClassifier(
        learning_rate=0.05,
        max_iter=160,
        max_leaf_nodes=31,
        min_samples_leaf=80,
        l2_regularization=0.01,
        early_stopping=True,
        random_state=seed,
    )
    model.fit(X_fit, y_fit)
    return model, "HistGradientBoosting fallback for LightGBM"


def predict_classifier(model: Any, X: np.ndarray) -> np.ndarray:
    if isinstance(model, float):
        return np.full(X.shape[0], model, dtype=np.float64)
    return model.predict_proba(X)[:, 1].astype(np.float64)


def risk_baseline_rows(
    train: dict[str, np.ndarray],
    eval_arrays: dict[str, np.ndarray],
    split: str,
    endpoints: list[str],
    horizons: list[float],
    seed: int,
    max_train: int,
    bootstrap_clusters: int,
    require_lightgbm: bool,
    model_cache: dict[tuple[str, float, str], tuple[Any, str]] | None = None,
) -> list[dict[str, Any]]:
    rows = []
    for endpoint in endpoints:
        for horizon in horizons:
            print(json.dumps({"stage": "risk_horizon_labels", "split": split, "endpoint": endpoint, "horizon_days": float(horizon)}), flush=True)
            tr_row, tr_pos, y_train, _ = horizon_label(train, endpoint, horizon)
            ev_row, ev_pos, y_eval, patient_ids = horizon_label(eval_arrays, endpoint, horizon)
            if y_train.size == 0 or y_eval.size == 0:
                continue
            X_train, _ = pre_action_features(train, tr_row, tr_pos)
            X_eval, _ = pre_action_features(eval_arrays, ev_row, ev_pos)
            for kind in ("clinical_landmark", "lightgbm_landmark"):
                print(
                    json.dumps(
                        {
                            "stage": "risk_model_fit",
                            "split": split,
                            "endpoint": endpoint,
                            "horizon_days": float(horizon),
                            "model_kind": kind,
                            "train_n": int(y_train.size),
                            "eval_n": int(y_eval.size),
                            "max_train": int(max_train),
                        }
                    ),
                    flush=True,
                )
                cache_key = (endpoint, float(horizon), kind)
                if model_cache is not None and cache_key in model_cache:
                    model, model_label = model_cache[cache_key]
                else:
                    model, model_label = fit_classifier(X_train, y_train, kind, seed + int(horizon) + len(endpoint), max_train, require_lightgbm)
                    if model_cache is not None:
                        model_cache[cache_key] = (model, model_label)
                print(
                    json.dumps(
                        {
                            "stage": "risk_model_predict",
                            "split": split,
                            "endpoint": endpoint,
                            "horizon_days": float(horizon),
                            "model": model_label,
                        }
                    ),
                    flush=True,
                )
                p = predict_classifier(model, X_eval)
                metrics = weighted_binary_metrics(y_eval, p)
                boot = bootstrap_metric_ci(
                    y_eval,
                    p,
                    patient_ids,
                    bootstrap_clusters,
                    seed + int(horizon) * 13 + len(endpoint) * 101 + (0 if kind == "clinical_landmark" else 1),
                )
                base = {
                    "split": split,
                    "endpoint": endpoint,
                    "horizon_days": float(horizon),
                    "model": model_label,
                    "n": int(y_eval.size),
                    "events": int(y_eval.sum()),
                    **metrics,
                    **boot,
                }
                rows.append(base)
                rows.extend(
                    calibration_curve_rows(
                        split=split,
                        endpoint=endpoint,
                        horizon_days=float(horizon),
                        model=model_label,
                        y=y_eval,
                        p=p,
                    )
                )
                for threshold in RISK_THRESHOLD_GRID.get(endpoint, RISK_THRESHOLD_GRID["default"]):
                    dca_boot = bootstrap_net_benefit_ci(
                        y_eval,
                        p,
                        patient_ids,
                        float(threshold),
                        bootstrap_clusters,
                        seed + int(horizon) * 17 + int(round(float(threshold) * 10000)) + (0 if kind == "clinical_landmark" else 1),
                    )
                    rows.append(
                        {
                            "split": split,
                            "endpoint": endpoint,
                            "horizon_days": float(horizon),
                            "model": model_label,
                            "metric_family": "decision_curve",
                            "threshold": float(threshold),
                            "net_benefit": net_benefit(y_eval, p, float(threshold)),
                            "bootstrap_n": int(bootstrap_clusters),
                            "n": int(y_eval.size),
                            "events": int(y_eval.sum()),
                            **dca_boot,
                        }
                    )
    return rows


def propensity_baseline_rows(
    train: dict[str, np.ndarray],
    eval_arrays: dict[str, np.ndarray],
    train_actions: dict[str, np.ndarray],
    eval_actions: dict[str, np.ndarray],
    ontology: dict[str, Any],
    split: str,
    seed: int,
    max_train: int,
    require_lightgbm: bool,
    model_cache: dict[int, Any] | None = None,
) -> list[dict[str, Any]]:
    rows = []
    slot_names = [slot["slot_name"] for slot in ontology["slots"]]
    local = ontology.get("local_label_to_action_value_id_by_slot", {})
    valid_train = train["valid_mask"].astype(bool)
    valid_eval = eval_arrays["valid_mask"].astype(bool)
    for slot_id, slot_name in enumerate(slot_names):
        present_id = local.get(slot_name, {}).get("present")
        if present_id is None:
            positive_label = next((label for label in ("是", "有", "面访", "住院", "持续", "社区转医院") if label in local.get(slot_name, {})), None)
            if positive_label is None:
                continue
            present_id = int(local[slot_name][positive_label])
        tr_mask = valid_train & train_actions["action_mask"][:, :, slot_id].astype(bool)
        ev_mask = valid_eval & eval_actions["action_mask"][:, :, slot_id].astype(bool)
        tr_row, tr_pos = np.nonzero(tr_mask)
        ev_row, ev_pos = np.nonzero(ev_mask)
        if tr_row.size == 0 or ev_row.size == 0:
            continue
        y_train = (train_actions["action_value_ids"][:, :, slot_id][tr_row, tr_pos] == int(present_id)).astype(np.int8)
        y_eval = (eval_actions["action_value_ids"][:, :, slot_id][ev_row, ev_pos] == int(present_id)).astype(np.int8)
        X_train, _ = pre_action_features(train, tr_row, tr_pos)
        X_eval, _ = pre_action_features(eval_arrays, ev_row, ev_pos)
        for kind, label in (("empirical_state_era", "Empirical service-state x era propensity"), ("lightgbm_landmark", "LightGBM propensity")):
            if kind == "empirical_state_era":
                key_train = service_era_key(train["service_state"][tr_row, tr_pos], train["visit_year"][tr_row, tr_pos])
                key_eval = service_era_key(eval_arrays["service_state"][ev_row, ev_pos], eval_arrays["visit_year"][ev_row, ev_pos])
                rates = {}
                for k in np.unique(key_train):
                    m = key_train == k
                    rates[int(k)] = float((y_train[m].sum() + 1.0) / (m.sum() + 2.0))
                fallback = float((y_train.sum() + 1.0) / (y_train.size + 2.0))
                p = np.asarray([rates.get(int(k), fallback) for k in key_eval], dtype=np.float64)
            else:
                if model_cache is not None and slot_id in model_cache:
                    model = model_cache[slot_id]
                else:
                    model, _ = fit_classifier(X_train, y_train, "lightgbm_landmark", seed + slot_id, max_train, require_lightgbm)
                    if model_cache is not None:
                        model_cache[slot_id] = model
                p = predict_classifier(model, X_eval)
            p = np.clip(p, 1.0e-4, 1.0 - 1.0e-4)
            w = y_eval / p + (1 - y_eval) / (1.0 - p)
            rows.append(
                {
                    "split": split,
                    "model": label,
                    "slot_id": int(slot_id),
                    "slot_name": slot_name,
                    "n": int(y_eval.size),
                    "present_rate": float(y_eval.mean()) if y_eval.size else math.nan,
                    "mean_propensity": float(p.mean()) if p.size else math.nan,
                    "p01": float(np.quantile(p, 0.01)) if p.size else math.nan,
                    "p05": float(np.quantile(p, 0.05)) if p.size else math.nan,
                    "p50": float(np.quantile(p, 0.50)) if p.size else math.nan,
                    "p95": float(np.quantile(p, 0.95)) if p.size else math.nan,
                    "p99": float(np.quantile(p, 0.99)) if p.size else math.nan,
                    "extreme_propensity_rate": float(((p < 0.01) | (p > 0.99)).mean()) if p.size else math.nan,
                    "effective_sample_size_binary_ipw": effective_sample_size(w),
                    **weighted_binary_metrics(y_eval, p),
                }
            )
    return rows


def load_actions(
    split: str,
    arrays: dict[str, np.ndarray],
    action_dir: Path | None,
    tensor_metadata: dict[str, Any],
    ordinal_direction: dict[str, Any],
    cat_value_vocab: dict[str, int],
    ontology: dict[str, Any],
    stage0_dir: Path,
    max_windows: int,
) -> dict[str, np.ndarray]:
    actions, _, _ = load_or_build_action_arrays(
        split=split,
        arrays=arrays,
        action_dir=action_dir,
        tensor_metadata=tensor_metadata,
        ordinal_direction=ordinal_direction,
        cat_value_vocab=cat_value_vocab,
        ontology=ontology,
        grace_days=(30.0, 60.0, 90.0),
        stage0_dir=stage0_dir,
    )
    for key, value in actions.items():
        if max_windows > 0 and value.shape[0] > max_windows:
            actions[key] = value[:max_windows]
    return actions


def main() -> None:
    parser = argparse.ArgumentParser(description="Run matched non-neural baselines for SCTM-v5 top-journal tables.")
    parser.add_argument("--tensor-dir", type=Path, required=True)
    parser.add_argument("--stage0-dir", type=Path, required=True)
    parser.add_argument("--action-dir", type=Path, default=None)
    parser.add_argument("--out-dir", type=Path, required=True)
    parser.add_argument("--splits", default="val,test")
    parser.add_argument("--horizons", default="90,365")
    parser.add_argument("--risk-endpoints", default="death,disengagement,primary_referral_relapse,service_escalation,clinical_deterioration,high_acuity_state")
    parser.add_argument("--max-windows", type=int, default=0)
    parser.add_argument("--max-risk-train", type=int, default=500000)
    parser.add_argument("--bootstrap-clusters", type=int, default=200)
    parser.add_argument("--seed", type=int, default=20260526)
    parser.add_argument("--require-lightgbm", action="store_true")
    args = parser.parse_args()

    args.out_dir.mkdir(parents=True, exist_ok=True)
    meta = load_json(args.tensor_dir / "tensor_metadata.json")
    vocab = load_json(args.tensor_dir / "cat_value_vocab.json")
    ordinal_direction = load_json(args.stage0_dir / "ordinal_direction_table.json")
    ontology = build_action_ontology(meta, vocab)
    train = load_npz_to_memory(resolve_split_path(args.tensor_dir, "train"))
    if args.max_windows > 0:
        train = {k: v[: args.max_windows] if getattr(v, "shape", (0,))[0] == train["valid_mask"].shape[0] else v for k, v in train.items()}
    train_actions = load_actions("train", train, args.action_dir, meta, ordinal_direction, vocab, ontology, args.stage0_dir, args.max_windows)

    print(json.dumps({"stage": "fit_semi_markov", "train_windows": int(train["valid_mask"].shape[0])}), flush=True)
    semi = fit_semi_markov(train, smoothing=1.0)
    print(json.dumps({"stage": "fit_empirical_grammar"}), flush=True)
    grammar = fit_empirical_grammar(train, meta, vocab, smoothing=1.0)
    print(json.dumps({"stage": "fit_baseline_models_start", "splits": args.splits}), flush=True)
    horizons = [float(x) for x in args.horizons.split(",") if x.strip()]
    endpoints = [x.strip() for x in args.risk_endpoints.split(",") if x.strip()]

    semi_rows = []
    grammar_rows = []
    risk_rows = []
    propensity_rows = []
    risk_model_cache: dict[tuple[str, float, str], tuple[Any, str]] = {}
    propensity_model_cache: dict[int, Any] = {}
    for split in [s.strip() for s in args.splits.split(",") if s.strip()]:
        print(json.dumps({"stage": "evaluate_split_start", "split": split}), flush=True)
        arrays = load_npz_to_memory(resolve_split_path(args.tensor_dir, split))
        if args.max_windows > 0:
            arrays = {k: v[: args.max_windows] if getattr(v, "shape", (0,))[0] == arrays["valid_mask"].shape[0] else v for k, v in arrays.items()}
        actions = load_actions(split, arrays, args.action_dir, meta, ordinal_direction, vocab, ontology, args.stage0_dir, args.max_windows)
        print(json.dumps({"stage": "semi_markov_metrics", "split": split}), flush=True)
        semi_rows.extend(semi_markov_metrics(semi, arrays, split))
        print(json.dumps({"stage": "empirical_grammar_metrics", "split": split}), flush=True)
        grammar_rows.extend(empirical_grammar_metrics(grammar, arrays, split, meta, vocab))
        print(json.dumps({"stage": "risk_baseline_metrics", "split": split}), flush=True)
        risk_rows.extend(
            risk_baseline_rows(
                train,
                arrays,
                split,
                endpoints,
                horizons,
                args.seed,
                args.max_risk_train,
                args.bootstrap_clusters,
                args.require_lightgbm,
                risk_model_cache,
            )
        )
        print(json.dumps({"stage": "propensity_baseline_metrics", "split": split}), flush=True)
        propensity_rows.extend(propensity_baseline_rows(train, arrays, train_actions, actions, ontology, split, args.seed, args.max_risk_train, args.require_lightgbm, propensity_model_cache))
        print(json.dumps({"stage": "evaluate_split_done", "split": split}), flush=True)

    paths = {
        "semi_markov": args.out_dir / "semi_markov_care_process_metrics.csv",
        "empirical_grammar": args.out_dir / "empirical_state_year_grammar_metrics.csv",
        "risk": args.out_dir / "clinical_lightgbm_risk_baseline_metrics.csv",
        "propensity": args.out_dir / "empirical_lightgbm_propensity_baseline_metrics.csv",
        "behrt_status": args.out_dir / "behrt_style_baseline_status.json",
    }
    pd.DataFrame(semi_rows).to_csv(paths["semi_markov"], index=False)
    pd.DataFrame(grammar_rows).to_csv(paths["empirical_grammar"], index=False)
    pd.DataFrame(risk_rows).to_csv(paths["risk"], index=False)
    pd.DataFrame(propensity_rows).to_csv(paths["propensity"], index=False)
    behrt_status = {
        "status": "not_run_in_this_script",
        "reason": "Exact BEHRT-style baseline requires separate event-token MLM pretraining and landmark fine-tuning; this non-neural matched baseline suite intentionally does not substitute a tabular transformer for that baseline.",
        "required_spec": {
            "architecture": "6-layer Transformer encoder, 12 heads, hidden size 768",
            "pretraining": "15% masked language modeling over m1_event_token_v2 tokens",
            "finetuning": "landmark risk heads for death, disengagement, event0 primary referral relapse, service escalation, clinical deterioration, and high acuity",
        },
    }
    paths["behrt_status"].write_text(json.dumps(behrt_status, ensure_ascii=False, indent=2), encoding="utf-8")
    summary = {
        "status": "completed",
        "script": "evaluate_v5_matched_baselines.py",
        "splits": args.splits,
        "horizons": horizons,
        "risk_endpoints": endpoints,
        "max_risk_train": int(args.max_risk_train),
        "bootstrap_clusters": int(args.bootstrap_clusters),
        "require_lightgbm": bool(args.require_lightgbm),
        "lightgbm_backend": "lightgbm" if LGBMClassifier is not None else "unavailable",
        "outputs": {k: str(v) for k, v in paths.items()},
        "notes": (
            "LightGBM rows use the lightgbm backend; use --require-lightgbm for fail-fast publication runs. "
            "Clinical Landmark uses unweighted L2 logistic regression so predict_proba remains on the natural prevalence scale. "
            "Event-horizon endpoints use the same teacher-forced observed-future-contact evaluable mask as evaluate_v5_topjournal.py; "
            "PWE endpoints use the same next-transition CIF target convention. "
            "Supported event endpoints include primary_referral_relapse, service_escalation, clinical_deterioration, and high_acuity_state."
        ),
    }
    (args.out_dir / "matched_baseline_summary.json").write_text(json.dumps(json_ready(summary), ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps({"status": "completed", "out_dir": str(args.out_dir)}, ensure_ascii=False), flush=True)


if __name__ == "__main__":
    main()