"""
Model comparison under walk-forward CV.

Protocol
--------
- Training pairs at origin T: all (device, origin') rows whose 42-day window closed
  before T (origin' <= T - 42d). Window closed => label fully observed, and there is
  a built-in embargo: no training window overlaps the scored window.
- Walk-forward: the model is refit at every scored origin on exactly the pairs that
  would have been observable then.
- DEV block s_8..s_23: hyperparameters (p_star, day quantile, planner) are tuned here.
- TEST block s_24..s_41: scored once with dev-chosen settings.

Models predict, per alive device:
  p42   - probability of EOL within the 42-day window
  day_q - quantile of days-to-EOL conditional on dying in the window

Planners turn (p42, day_q) into a plan:
  onepd  - worst-first, one per day (control, shape of the shipped planner)
  sched  - swap each included battery on its predicted day
  batch  - sched + move each building's swaps to the group's earliest day and
           cap building-groups per day (travel is expensive, early days are cheap)
"""

from pathlib import Path
import os
import sys
import warnings

import numpy as np
import pandas as pd

# Benign and expected: tau_projected (and occasionally other physics columns)
# is all-NaN at the earliest origins, before the master curve exists; the
# imputer skips the column, which has been the accepted behavior since the
# feature was added (phase 6). One fit per origin per fold prints it, so an
# all48oof run emits it ~290 times - silence THIS message only, so real
# warnings stay visible.
warnings.filterwarnings(
    'ignore', message='Skipping features without any observed values')

from sklearn.impute import SimpleImputer
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import make_pipeline
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import HistGradientBoostingClassifier, GradientBoostingRegressor
from sklearn.metrics import roc_auc_score

from batteryswap_public.utils import load_dataset, iterate_scenarios
from batteryswap_public.evaluate import evaluate_plan, check_plan_valid

from batteryswap_example.features import load as load_features, WINDOW_DAYS

FEATURES = ['v_last', 'v_drop', 'slope_30', 'slope_90', 'v_min_30',
            'staleness', 'age_days',
            'slope_lastk', 'drop_recent', 'below_frac', 'accel', 'n_30',
            'tau_projected']
# NOT in FEATURES, deliberately: 'vrel', 'vmin7_rank'. features.py computes
# both (and v_min_7), and the training table carries them, so enabling them is
# a one-line change here. Held back from the pooling submission so that
# submission isolates ONE change and stays attributable - phase 28 stacked four
# and needed two further phases to untangle the damage. Evidence as of
# 2026-08-21 is also unsettled: vrel survived a leave-one-out ablation on the
# blind harness, vmin7_rank was ablated OUT (-12.8), and a static-fit gate
# showed the whole feature-set swap swinging +55 to -109 depending on fit pool.
# vstd_30/rng_30 (plunge-precursor volatility, AUC 0.83/0.81 on the silent
# cohort) rejected as GLOBAL features: dev +17, test -0.7 - the signal is
# diluted across the healthy majority. Used instead by the SilentFix override.
# Rejected by CV despite AUC 0.91 on voltage-silent pairs: cohort_rel_age /
# cohort_overdue (install-batch calibrated clock, computed in features.py).
# Coverage is too low at dev-era origins (14% exact / 54% monthly), so the
# no-information majority degrades the fit in every encoding tried
# (dev 2074 -> 2279 raw, -> 2433 overdue-neutral). Revisit only if far more
# batch deaths become observable.
# tau_projected = C^-1(v_corr_last) - staleness: targets devices that go dark and
# keep discharging (the model's main blind spot - 25 of 60 test-block misses).
# Other physics features (tau_hat, fit_err, v_corr_last, thresh_dist) remain
# computed in features.py but excluded: walk-forward CV showed no cost improvement
# over the empirical voltage features, which already saturate the same signal.
DEV_ORIGINS = range(8, 24)
TEST_ORIGINS = range(24, 42)
DAY_QUANTILES = (0.05, 0.20, 0.50, 0.80)


# --------------------------------------------------------------------------- models

class DayQuantiles:
    """Quantiles of days-to-EOL, fit on training pairs that died in-window."""

    def __init__(self, n_estimators=100, max_depth=2, quantiles=DAY_QUANTILES):
        self.models = {}
        self.fallback = {}
        self.n_estimators = n_estimators
        self.max_depth = max_depth
        self.quantiles = quantiles

    def fit(self, X, days):
        for q in self.quantiles:
            self.fallback[q] = float(np.quantile(days, q))
            if len(days) >= 40:
                m = make_pipeline(
                    SimpleImputer(strategy='median'),
                    GradientBoostingRegressor(loss='quantile', alpha=q,
                                              n_estimators=self.n_estimators,
                                              max_depth=self.max_depth,
                                              random_state=0))
                m.fit(X, days)
                self.models[q] = m
        return self

    def predict(self, X):
        out = {}
        for q in self.quantiles:
            if q in self.models:
                out[f'day_q{q}'] = np.clip(self.models[q].predict(X), 0, WINDOW_DAYS)
            else:
                out[f'day_q{q}'] = np.full(len(X), self.fallback[q])
        return pd.DataFrame(out, index=X.index)


Q6 = (0.05, 0.10, 0.20, 0.35, 0.50, 0.80)


class SklearnRUL:
    """p42 from a sklearn classifier + DayQuantiles for timing."""

    def __init__(self, kind):
        self.opts = {}
        if ':' in kind:
            kind, spec = kind.split(':', 1)
            self.opts = {k: float(v) for k, v in
                         (kv.split('=') for kv in spec.split(','))}
        self.kind = kind

    def fit(self, train):
        X, y = train[FEATURES], train['dies42']
        if self.kind == 'logistic':
            self.clf = make_pipeline(SimpleImputer(strategy='median'),
                                     StandardScaler(),
                                     LogisticRegression(max_iter=2000, C=1.0))
        elif self.kind == 'gbm':
            self.clf = HistGradientBoostingClassifier(
                max_iter=200, max_depth=3, learning_rate=0.1,
                min_samples_leaf=20, random_state=0)
        elif self.kind == 'logcal':
            # calibrated probabilities: makes sum(p42) a true expected count and
            # expected-cost planning meaningful
            from sklearn.calibration import CalibratedClassifierCV
            base = make_pipeline(SimpleImputer(strategy='median'),
                                 StandardScaler(),
                                 LogisticRegression(max_iter=2000, C=1.0))
            self.clf = CalibratedClassifierCV(base, method='sigmoid', cv=5)
        elif self.kind == 'logbag':
            # bagged logistic: averages away fit variance near the inclusion
            # cutoff, where a p42 flip decides catch-vs-emergency
            from sklearn.ensemble import BaggingClassifier
            base = make_pipeline(SimpleImputer(strategy='median'),
                                 StandardScaler(),
                                 LogisticRegression(max_iter=2000,
                                                    C=self.opts.get('C', 1.0)))
            self.clf = BaggingClassifier(
                base, n_estimators=int(self.opts.get('bags', 15)), random_state=0)
        elif self.kind == 'rf':
            from sklearn.ensemble import RandomForestClassifier
            cw = 'balanced_subsample' if self.opts.get('bal') else None
            self.clf = make_pipeline(
                SimpleImputer(strategy='median'),
                RandomForestClassifier(
                    n_estimators=int(self.opts.get('trees', 500)),
                    min_samples_leaf=int(self.opts.get('leaf', 5)),
                    max_features='sqrt', class_weight=cw,
                    random_state=0, n_jobs=-1))
        elif self.kind == 'xtrees':
            from sklearn.ensemble import ExtraTreesClassifier
            self.clf = make_pipeline(
                SimpleImputer(strategy='median'),
                ExtraTreesClassifier(n_estimators=500, min_samples_leaf=5,
                                     max_features='sqrt', random_state=0,
                                     n_jobs=-1))
        elif self.kind == 'mlp':
            from sklearn.neural_network import MLPClassifier
            self.clf = make_pipeline(
                SimpleImputer(strategy='median'),
                StandardScaler(),
                MLPClassifier(hidden_layer_sizes=(32, 16), alpha=1e-2,
                              max_iter=1000, early_stopping=True,
                              random_state=0))
        self.clf.fit(X, y)

        dying = train[train['dies42'] == 1]
        self.days = DayQuantiles(
            n_estimators=int(self.opts.get('dtrees', 100)),
            max_depth=int(self.opts.get('ddepth', 2)),
            quantiles=Q6 if self.opts.get('q6') else DAY_QUANTILES,
        ).fit(dying[FEATURES], dying['days_to_eol'])
        return self

    def predict(self, rows):
        X = rows[FEATURES]
        out = self.days.predict(X)
        out['p42'] = self.clf.predict_proba(X)[:, 1]
        if GHOST_ZERO and 'days_below_245' in rows.columns:
            # a device below the 2.45 death threshold for >42d with no EOL on
            # file will never be scored as a death (26 such ghosts in train) -
            # swapping it is a pure false-positive cost. Require it to STILL
            # be in the death zone: days_below_245 is sticky from the first
            # crossing, and a dip-then-recovery must stay eligible.
            ghost = ((rows['days_below_245'] > WINDOW_DAYS)
                     & (rows['v_corr_last'] < 2.55)).values
            out.loc[ghost, 'p42'] = 0.0
        return out


GHOST_ZERO = False   # phase-15 candidate: exclude label-process ghosts from plans
QUARANTINE = False   # phase-15 candidate: drop looks-dead-labeled-alive training rows


class LogbagLAH:
    """
    Two-regime scorer for the label process (phase 17). The metric charges
    RECORDED deaths, and EOL recording lags the physical 2.45-crossing by a
    median ~34d (a third of recorded deaths lag >42d; ghosts never get one).
    For devices already in the death zone the uncertainty is not "will it
    die" but "will the label arrive in this window" - a pure function of
    time-since-crossing. So:
      not crossed  -> champion logbag p42 (its real domain)
      crossed and still low -> p42 = A(a) = P(EOL recorded in (a, a+42] | none
             by a), a Kaplan-Meier on observed crossing->EOL lags (uncensored:
             recorded deaths; censored: crossed devices with no EOL yet).
    Unlike the phase-15 binary zero rule this keeps mid-lag devices included
    (A ~ 0.3-0.6) and lets true ghosts fade gradually as their censored mass
    flattens the curve. domain_split=True additionally trains the classifier
    only on not-crossed rows, since it is never queried in the death zone.
    """

    BUCKET = 14.0  # KM lag buckets (days); 82 events cannot support finer

    def __init__(self, domain_split=False):
        self.domain_split = domain_split

    def fit(self, train):
        crossed = train['days_below_245'].notna() & (train['days_below_245'] >= 0)
        self.who = SklearnRUL('logbag').fit(train[~crossed] if self.domain_split
                                            else train)

        # one lag observation per device: crossing -> EOL (or censored at the
        # last origin it appears). days_below_245 is leak-free per row.
        sub = train[crossed]
        per_dev = sub.sort_values('origin').groupby('device_id').tail(1)
        cross_date = per_dev['origin'] - pd.to_timedelta(per_dev['days_below_245'], unit='D')
        has_eol = per_dev['eol'].notna()
        lag = np.where(has_eol,
                       (per_dev['eol'] - cross_date).dt.days,
                       per_dev['days_below_245'])
        lag = np.clip(lag, 0, None)
        event = has_eol.values

        # discrete KM over lag buckets
        k = (lag // self.BUCKET).astype(int)
        kmax = int(k.max()) + 1 if len(k) else 1
        surv, s = [], 1.0
        for b in range(kmax):
            at_risk = (k >= b).sum()
            died = ((k == b) & event).sum()
            s *= 1.0 - (died / at_risk if at_risk > 0 else 0.0)
            surv.append(s)
        self.surv = np.array(surv)
        return self

    def _arrival(self, a):
        """P(EOL recorded within the next 42d | crossed a days ago, none yet)."""
        b0 = np.clip((a // self.BUCKET).astype(int), 0, len(self.surv) - 1)
        b1 = np.clip(((a + WINDOW_DAYS) // self.BUCKET).astype(int), 0, len(self.surv) - 1)
        s0 = self.surv[b0]
        s1 = self.surv[b1]
        return np.where(s0 > 1e-9, 1.0 - s1 / s0, 0.0)

    def predict(self, rows):
        out = self.who.predict(rows)
        crossed = (rows['days_below_245'].notna()
                   & (rows['days_below_245'] >= 0)
                   & (rows['v_corr_last'] < 2.50)).values
        if crossed.any():
            a = rows.loc[crossed, 'days_below_245'].values
            out.loc[crossed, 'p42'] = self._arrival(a)
        return out

SILENT_V = 2.55
SILENT_FEATURES = ['rng_30', 'vstd_30', 'tau_projected', 'cohort_rel_age_m',
                   'age_days', 'v_corr_last']


class SilentFix:
    """
    Base logbag plus a dedicated blind-zone model. Batteries whose voltage still
    looks healthy (v_corr_last > SILENT_V) get a rank boost from a second
    logistic trained ONLY on healthy-voltage pairs, using the plunge-precursor
    and cohort signals that carry AUC there while the level is blind.

    Scale-free: emits rank_score for ORDERING only; p42 (and hence the topsum
    count) is untouched.
    """

    def __init__(self, w=0.3):
        self.w = w

    def fit(self, train):
        self.base = SklearnRUL('logbag').fit(train)
        sil = train[train['v_corr_last'] > SILENT_V]
        self.silent = None
        if sil['dies42'].sum() >= 20:
            self.silent = make_pipeline(SimpleImputer(strategy='median'),
                                        StandardScaler(),
                                        LogisticRegression(max_iter=2000, C=1.0))
            self.silent.fit(sil[SILENT_FEATURES], sil['dies42'])
        return self

    def predict(self, rows):
        out = self.base.predict(rows)
        out['rank_score'] = out['p42']
        if self.silent is not None:
            mask = rows['v_corr_last'] > SILENT_V
            if mask.any():
                p_sil = self.silent.predict_proba(rows.loc[mask, SILENT_FEATURES])[:, 1]
                out.loc[mask, 'rank_score'] = out.loc[mask, 'p42'] + self.w * p_sil
        return out


class RankBlendRF:
    """logbag's probability scale with (logbag+RF) mean-rank ordering: keeps
    sum(p42) exactly (topsum count unchanged) while letting the forest's ranking
    reorder WHO is included - RF ranks better at some origins but its raw
    probabilities are too compressed for count-tracking."""

    def fit(self, train):
        self.a = SklearnRUL('logbag').fit(train)
        self.b = SklearnRUL('rf').fit(train)
        return self

    def predict(self, rows):
        out = self.a.predict(rows)
        mean_rank = (out['p42'].rank() + self.b.predict(rows)['p42'].rank()) / 2
        vals = np.sort(out['p42'].values)                 # same multiset of values
        out['p42'] = pd.Series(vals, index=mean_rank.sort_values().index).reindex(out.index)
        return out


class RankVoltage:
    """Control: score by v_last only; p42 is a rank in [0,1], day from DayQuantiles."""

    def fit(self, train):
        dying = train[train['dies42'] == 1]
        self.days = DayQuantiles().fit(dying[FEATURES], dying['days_to_eol'])
        return self

    def predict(self, rows):
        out = self.days.predict(rows[FEATURES])
        v = rows['v_last'].fillna(np.inf)
        out['p42'] = 1.0 - v.rank(pct=True).values  # lowest voltage -> highest score
        return out


class WeibullAFT:
    """lifelines Weibull AFT on censored pairs; survival curve gives p42 and timing."""

    def fit(self, train_censored):
        from lifelines import WeibullAFTFitter
        df = train_censored[FEATURES + ['duration', 'event']].copy()
        self.imputer = SimpleImputer(strategy='median').fit(df[FEATURES])
        df[FEATURES] = self.imputer.transform(df[FEATURES])
        self.aft = WeibullAFTFitter(penalizer=0.01)
        with warnings.catch_warnings():
            warnings.simplefilter('ignore')
            self.aft.fit(df, duration_col='duration', event_col='event')
        return self

    def predict(self, rows):
        X = rows[FEATURES].copy()
        X[FEATURES] = self.imputer.transform(X[FEATURES])
        times = np.arange(1, WINDOW_DAYS + 1)
        surv = self.aft.predict_survival_function(X, times=times)  # times x devices
        cdf = 1.0 - surv.T  # devices x times
        p42 = cdf.iloc[:, -1]
        out = pd.DataFrame({'p42': p42.values}, index=rows.index)

        # per-day death mass and expected residual life given surviving the window,
        # for the expected-cost planner
        pmf = np.diff(np.column_stack([np.zeros(len(cdf)), cdf.values]), axis=1)
        for t in times:
            out[f'pd_{t}'] = pmf[:, t - 1]
        far = np.arange(WINDOW_DAYS, 400)
        surv_far = self.aft.predict_survival_function(X, times=far).T.values
        s42 = np.maximum(surv_far[:, 0], 1e-12)
        out['e_far'] = WINDOW_DAYS + np.trapezoid(surv_far, far, axis=1) / s42
        # quantiles of death time conditional on dying within the window
        cond = cdf.div(p42.replace(0, np.nan), axis=0).clip(upper=1.0)
        for q in DAY_QUANTILES:
            day = (cond.values >= q).argmax(axis=1) + 1.0
            day[~(cond.values[:, -1] >= q)] = WINDOW_DAYS
            out[f'day_q{q}'] = np.where(p42.values > 0, day, WINDOW_DAYS)
        return out


class DiscreteHazard:
    """
    Discrete-time hazard logistic on person-period (pair, week-bucket) rows.

    Trains on ALL prior pairs, not just closed windows: a pair whose window is
    still open at the scored origin contributes its observed buckets (event if
    the death is already visible, censored otherwise). Recovers the freshest
    training data - the six origins immediately before the scored one - which
    training_pairs() discards entirely.

    p42 = 1 - prod_k (1 - h_k); conditional day quantiles from the bucket CDF.
    """

    BUCKET = 7.0
    K = 6  # buckets cover days (7k, 7(k+1)], k=0..5

    def __init__(self, extra=None):
        self.feats = FEATURES + (extra or [])

    def fit(self, train_censored):
        dur = train_censored['duration'].clip(upper=WINDOW_DAYS).values
        event = ((train_censored['event'] == 1) &
                 (train_censored['duration'] <= WINDOW_DAYS)).values
        X = train_censored[self.feats].values
        rows, ys, ks = [], [], []
        for k in range(self.K):
            lo, hi = self.BUCKET * k, self.BUCKET * (k + 1)
            died_here = event & (dur > lo) & (dur <= hi)
            # censored-mid-bucket exposure is dropped: only fully observed
            # buckets contribute negatives
            in_bucket = (dur >= hi) | died_here
            rows.append(X[in_bucket])
            ys.append(died_here[in_bucket])
            ks.append(np.full(in_bucket.sum(), k))
        Xpp = np.vstack(rows)
        kpp = np.concatenate(ks)
        Xpp = np.column_stack([Xpp] + [(kpp == k).astype(float)
                                       for k in range(1, self.K)])
        self.clf = make_pipeline(SimpleImputer(strategy='median'),
                                 StandardScaler(),
                                 LogisticRegression(max_iter=2000, C=1.0))
        self.clf.fit(Xpp, np.concatenate(ys))

        dead = train_censored[(train_censored['event'] == 1) &
                              (train_censored['duration'] <= WINDOW_DAYS)]
        self.days = DayQuantiles().fit(dead[FEATURES], dead['duration'])
        return self

    def predict(self, rows):
        X = rows[self.feats].values
        n = len(rows)
        haz = np.empty((n, self.K))
        for k in range(self.K):
            dummies = [(np.full(n, k) == j).astype(float) for j in range(1, self.K)]
            haz[:, k] = self.clf.predict_proba(np.column_stack([X] + dummies))[:, 1]
        surv = np.cumprod(1.0 - haz, axis=1)
        p42 = 1.0 - surv[:, -1]
        out = self.days.predict(rows[FEATURES])
        out['p42'] = p42
        # override day quantiles with the hazard model's own conditional CDF
        cdf = (1.0 - surv) / np.maximum(p42[:, None], 1e-12)
        ends = self.BUCKET * np.arange(1, self.K + 1)
        for q in DAY_QUANTILES:
            day = np.array([np.interp(q, np.concatenate([[0.0], cdf[i]]),
                                      np.concatenate([[0.0], ends]))
                            for i in range(n)])
            out[f'day_q{q}'] = np.clip(day, 0, WINDOW_DAYS)
        return out



def determined_labels(train_censored):
    """Pairs whose dies42 label is already fully determined at the cut:
    survived the whole window (duration >= 42) or died inside it. Recovers
    open-window positives that training_pairs() throws away."""
    neg = train_censored['duration'] >= WINDOW_DAYS
    pos = (train_censored['event'] == 1) & (train_censored['duration'] <= WINDOW_DAYS)
    sub = train_censored[neg | pos].copy()
    sub['dies42'] = ((sub['event'] == 1) &
                     (sub['duration'] <= WINDOW_DAYS)).astype(int)
    sub['days_to_eol'] = sub['duration']
    return sub


class LogBagOpen:
    """logbag trained on the determined-label set (closed windows + open-window
    deaths observed before the origin)."""

    def fit(self, train_censored):
        from sklearn.ensemble import BaggingClassifier
        train = determined_labels(train_censored)
        X, y = train[FEATURES], train['dies42']
        base = make_pipeline(SimpleImputer(strategy='median'),
                             StandardScaler(),
                             LogisticRegression(max_iter=2000, C=1.0))
        self.clf = BaggingClassifier(base, n_estimators=15, random_state=0)
        self.clf.fit(X, y)
        dying = train[train['dies42'] == 1]
        self.days = DayQuantiles().fit(dying[FEATURES], dying['days_to_eol'])
        return self

    def predict(self, rows):
        X = rows[FEATURES]
        out = self.days.predict(X)
        out['p42'] = self.clf.predict_proba(X)[:, 1]
        return out


class BlendDH:
    """Mean of LogBagOpen and DiscreteHazard p42: same features, but different
    training sets and likelihoods, so their errors decorrelate."""

    def fit(self, train_censored):
        self.a = LogBagOpen().fit(train_censored)
        self.b = DiscreteHazard().fit(train_censored)
        return self

    def predict(self, rows):
        out = self.a.predict(rows)
        out['p42'] = 0.5 * (out['p42'] + self.b.predict(rows)['p42'])
        return out


class BlendDHAux:
    """Mean of logbag_aux (closed windows) and DiscreteHazard-with-aux
    (censored open windows): decorrelated training sets, shared feature set."""

    def fit(self, train_censored):
        # recover the scored origin T: censored durations end exactly at T
        tc = train_censored
        T = (tc['origin'] + pd.to_timedelta(tc['duration'], unit='D')).max()
        closed = tc[tc['origin'] <= T - pd.Timedelta(days=WINDOW_DAYS)]
        self.a = LogBagFeats(['rng_90', 'vstd_30', 'rng_ratio', 'gap_max_60']).fit(closed)
        self.b = DiscreteHazard(['rng_90', 'vstd_30', 'rng_ratio', 'gap_max_60']).fit(train_censored)
        return self

    def predict(self, rows):
        out = self.a.predict(rows)
        out['p42'] = 0.5 * (out['p42'] + self.b.predict(rows)['p42'])
        return out


class CostWeightBag:
    """
    Bagged logistic with DECISION-STAKE sample weights (phase 17): a missed
    positive costs ~emergency (~200); a false inclusion costs its era's
    false-swap price (0.5/day against the end_time+30 imputation, so early
    origins' negatives matter ~10x more than near-data-end ones). Weights
    normalized to mean 1 so the L2 scale (C=1) stays comparable to logbag.
    """

    MISS_W = 200.0

    def __init__(self, n_estimators=15):
        self.n_estimators = n_estimators

    def fit(self, train):
        X, y = train[FEATURES], train['dies42']
        data_end = train['origin'].max()
        fp_price = (0.5 * ((data_end - train['origin']).dt.days + 30 - 21)
                    ).clip(lower=1.0).values
        w = np.where(y.values == 1, self.MISS_W, fp_price)
        w = w / w.mean()
        rng = np.random.RandomState(0)
        self.models = []
        n = len(X)
        for _ in range(self.n_estimators):
            idx = rng.randint(0, n, n)
            m = make_pipeline(SimpleImputer(strategy='median'),
                              StandardScaler(),
                              LogisticRegression(max_iter=2000, C=1.0))
            m.fit(X.iloc[idx], y.iloc[idx],
                  logisticregression__sample_weight=w[idx])
            self.models.append(m)
        dying = train[train['dies42'] == 1]
        self.days = DayQuantiles().fit(dying[FEATURES], dying['days_to_eol'])
        return self

    def predict(self, rows):
        X = rows[FEATURES]
        out = self.days.predict(X)
        out['p42'] = np.mean([m.predict_proba(X)[:, 1] for m in self.models], axis=0)
        return out


class WeightedBagLogistic:
    """Bagged logistic with recency-weighted samples: pairs from recent origins
    count more, on the theory that fleet composition and degradation mix drift."""

    def __init__(self, half_life=180.0, n_estimators=15):
        self.half_life = half_life
        self.n_estimators = n_estimators

    def fit(self, train):
        X, y = train[FEATURES], train['dies42']
        age = (train['origin'].max() - train['origin']).dt.days.values
        w = 0.5 ** (age / self.half_life)
        rng = np.random.RandomState(0)
        self.models = []
        n = len(X)
        for _ in range(self.n_estimators):
            idx = rng.randint(0, n, n)
            m = make_pipeline(SimpleImputer(strategy='median'),
                              StandardScaler(),
                              LogisticRegression(max_iter=2000, C=1.0))
            m.fit(X.iloc[idx], y.iloc[idx],
                  logisticregression__sample_weight=w[idx])
            self.models.append(m)
        dying = train[train['dies42'] == 1]
        self.days = DayQuantiles().fit(dying[FEATURES], dying['days_to_eol'])
        return self

    def predict(self, rows):
        X = rows[FEATURES]
        out = self.days.predict(X)
        out['p42'] = np.mean([m.predict_proba(X)[:, 1] for m in self.models], axis=0)
        return out



class BlendRUL:
    """Mean of logbag and gbm probabilities: linear and tree models miss
    different deaths; averaging keeps sum(p42) meaningful for topsum."""

    def fit(self, train):
        self.a = SklearnRUL('logbag').fit(train)
        try:
            self.b = SklearnRUL('gbm').fit(train)
        except ValueError:  # hist-GBM needs more rows than the earliest origins have
            self.b = None
        return self

    def predict(self, rows):
        out = self.a.predict(rows)
        if self.b is not None:
            out['p42'] = 0.5 * (out['p42'] + self.b.predict(rows)['p42'])
        return out


class LogBagC:
    """logbag with a non-default inverse regularization C."""

    def __init__(self, C):
        self.C = C

    def fit(self, train):
        from sklearn.ensemble import BaggingClassifier
        X, y = train[FEATURES], train['dies42']
        base = make_pipeline(SimpleImputer(strategy='median'),
                             StandardScaler(),
                             LogisticRegression(max_iter=2000, C=self.C))
        self.clf = BaggingClassifier(base, n_estimators=15, random_state=0)
        self.clf.fit(X, y)
        dying = train[train['dies42'] == 1]
        self.days = DayQuantiles().fit(dying[FEATURES], dying['days_to_eol'])
        return self

    def predict(self, rows):
        X = rows[FEATURES]
        out = self.days.predict(X)
        out['p42'] = self.clf.predict_proba(X)[:, 1]
        return out



class LogBagFeats:
    """logbag with extra feature columns (aux usage-dynamics table)."""

    def __init__(self, extra, n_estimators=15):
        self.feats = FEATURES + extra
        self.n_estimators = n_estimators

    def fit(self, train):
        from sklearn.ensemble import BaggingClassifier
        X, y = train[self.feats], train['dies42']
        base = make_pipeline(SimpleImputer(strategy='median'),
                             StandardScaler(),
                             LogisticRegression(max_iter=2000, C=1.0))
        self.clf = BaggingClassifier(base, n_estimators=self.n_estimators, random_state=0)
        self.clf.fit(X, y)
        dying = train[train['dies42'] == 1]
        self.days = DayQuantiles().fit(dying[FEATURES], dying['days_to_eol'])
        return self

    def predict(self, rows):
        out = self.days.predict(rows[FEATURES])
        out['p42'] = self.clf.predict_proba(rows[self.feats])[:, 1]
        return out


class LogBagUCB(LogBagFeats):
    """Rank by mean + k*std across the bagged estimators: a battery the ensemble
    disagrees about gets an inclusion boost (missing it costs ~4x including it).
    N still comes from sum of mean p42, so the count stays calibrated."""

    def __init__(self, extra, k=0.5):
        super().__init__(extra)
        self.k = k

    def predict(self, rows):
        out = self.days.predict(rows[FEATURES])
        X = rows[self.feats]
        Xi = pd.DataFrame(X, columns=self.feats)
        probs = np.stack([est.predict_proba(Xi.iloc[:, feats])[:, 1]
                          for est, feats in zip(self.clf.estimators_,
                                                self.clf.estimators_features_)])
        mean, std = probs.mean(axis=0), probs.std(axis=0)
        out['p42'] = np.clip(mean + self.k * std, 0, 1)
        out['p42_sum'] = mean.sum()
        return out


class LogBagLog(LogBagFeats):
    """Same features, but log1p-compress the heavy-tailed ones so the linear
    model is not dominated by extreme staleness/gap values."""

    SKEWED = ['staleness', 'gap_max_60', 'n_30']

    def _tx(self, rows):
        X = rows[self.feats].copy()
        for c in self.SKEWED:
            if c in X:
                X[c] = np.log1p(X[c].clip(lower=0))
        return X

    def fit(self, train):
        from sklearn.ensemble import BaggingClassifier
        X, y = self._tx(train), train['dies42']
        base = make_pipeline(SimpleImputer(strategy='median'),
                             StandardScaler(),
                             LogisticRegression(max_iter=2000, C=1.0))
        self.clf = BaggingClassifier(base, n_estimators=self.n_estimators,
                                     random_state=0)
        self.clf.fit(X, y)
        dying = train[train['dies42'] == 1]
        self.days = DayQuantiles().fit(dying[FEATURES], dying['days_to_eol'])
        return self

    def predict(self, rows):
        out = self.days.predict(rows[FEATURES])
        out['p42'] = self.clf.predict_proba(self._tx(rows))[:, 1]
        return out




class WienerFHT:
    """
    First-hitting-time RUL: temp-corrected voltage is Brownian motion with
    per-device drift; EOL is the first passage of THRESH = 2.40 V (validated,
    sd 0.034 across all observed deaths).

    P(T <= t) for barrier distance d > 0, drift m (toward barrier), variance s2:
        F(t) = Phi((m t - d)/sqrt(s2 t)) + exp(2 m d / s2) Phi(-(m t + d)/sqrt(s2 t))
    valid for any sign of m. Devices dark for `stale` days are conditioned on
    having survived the silence: p42 = (F(stale+42) - F(stale)) / (1 - F(stale)) -
    uncertainty through the reporting gap is handled by the process itself.

    Per-device (m, s2) are shrunk toward population medians (hierarchical-lite);
    day quantiles come from the conditional first-passage CDF.
    """

    THRESH = 2.40
    K_SHRINK = 6.0

    def fit(self, train):
        ok = train['wnr_n'] >= 3
        self.m_pop = float(np.nanmedian(train.loc[ok, 'wnr_m'])) if ok.any() else 1e-4
        self.s2_pop = float(np.nanmedian(train.loc[ok, 'wnr_s2'])) if ok.any() else 1e-4
        dying = train[train['dies42'] == 1]
        self.days = DayQuantiles().fit(dying[FEATURES], dying['days_to_eol'])
        return self

    def predict(self, rows):
        from scipy.stats import norm
        from scipy.special import log_ndtr

        n = len(rows)
        w = rows['wnr_n'].fillna(0).values
        m = np.where(w > 0, rows['wnr_m'].fillna(self.m_pop).values, self.m_pop)
        s2 = np.where(w > 0, rows['wnr_s2'].fillna(self.s2_pop).values, self.s2_pop)
        m = (w * m + self.K_SHRINK * self.m_pop) / (w + self.K_SHRINK)
        s2 = np.maximum((w * s2 + self.K_SHRINK * self.s2_pop) / (w + self.K_SHRINK),
                        1e-8)
        d = rows['v_corr_last'].values - self.THRESH
        d = np.where(np.isnan(d), 0.35, d)      # never-measured device
        at_barrier = d <= 1e-6
        d = np.maximum(d, 1e-6)
        stale = rows['staleness'].fillna(0).clip(lower=0).values

        def F(times):
            """first-passage CDF, times: devices x k"""
            sig = np.sqrt(s2[:, None] * times)
            z1 = (m[:, None] * times - d[:, None]) / sig
            z2 = -(m[:, None] * times + d[:, None]) / sig
            with np.errstate(over='ignore'):
                term2 = np.exp(np.minimum(
                    2.0 * m[:, None] * d[:, None] / s2[:, None] + log_ndtr(z2), 0.0))
            return np.clip(norm.cdf(z1) + term2, 0.0, 1.0)

        t_grid = np.arange(1, WINDOW_DAYS + 1, dtype=float)
        Ft = F(stale[:, None] + t_grid[None, :])
        F0 = np.where(stale > 0, F(np.maximum(stale, 1e-6)[:, None])[:, 0], 0.0)
        cond = np.clip((Ft - F0[:, None]) / np.maximum(1.0 - F0, 1e-9)[:, None],
                       0.0, 1.0)
        p42 = np.where(at_barrier, 1.0, cond[:, -1])

        out = self.days.predict(rows[FEATURES])
        out['p42'] = p42
        qcdf = cond / np.maximum(p42[:, None], 1e-12)
        for q in DAY_QUANTILES:
            idx = (qcdf >= q).argmax(axis=1)
            hit = qcdf[:, -1] >= q
            day = np.where(hit, t_grid[idx], WINDOW_DAYS).astype(float)
            out[f'day_q{q}'] = np.where(at_barrier, 0.0, day)
        return out


class LogWienerBlend:
    """logbag p42 blended with the Wiener FHT p42."""

    def __init__(self, mode='mean'):
        self.mode = mode

    def fit(self, train):
        self.a = SklearnRUL('logbag').fit(train)
        self.b = WienerFHT().fit(train)
        return self

    def predict(self, rows):
        out = self.a.predict(rows)
        pw = self.b.predict(rows)['p42']
        if self.mode == 'mean':
            out['p42'] = 0.5 * (out['p42'] + pw)
        elif self.mode == 'stale':
            stale = rows['staleness'].fillna(0).values
            out['p42'] = np.where(stale > 14, np.maximum(out['p42'], pw), out['p42'])
        return out


class ConformalCut:
    """
    Conformal-risk inclusion: wraps logbag and, at each refit, sets the score
    cutoff to the risk_q quantile of the model's scores on *dying* calibration
    pairs from the most recent closed origins. Scale-free by construction: the
    cutoff moves with whatever scale the refit produces. mode 'split' holds the
    calibration block out of training; 'in' trains on everything (slightly
    optimistic scores, but twice the recency).
    """

    def __init__(self, risk_q=0.15, mode='in', calib_days=120):
        self.risk_q = risk_q
        self.mode = mode
        self.calib_days = calib_days

    def fit(self, train):
        t_max = train['origin'].max()
        calib = train[train['origin'] > t_max - pd.Timedelta(days=self.calib_days)]
        fit_on = train if self.mode == 'in' else \
            train[train['origin'] <= t_max - pd.Timedelta(days=self.calib_days)]
        if len(fit_on) < 500 or fit_on['dies42'].sum() < 10:
            fit_on = train
        self.model = SklearnRUL('logbag').fit(fit_on)
        dead = calib[calib['dies42'] == 1]
        if len(dead) >= 8:
            scores = self.model.predict(dead)['p42']
            self.cut = float(np.quantile(scores, self.risk_q))
        else:
            self.cut = 0.5
        return self

    def predict(self, rows):
        out = self.model.predict(rows)
        out['p_cut'] = self.cut
        return out



class CohortGate:
    """
    logbag everywhere; on devices whose install batch has a calibrated monthly
    clock (cohort_rel_age_m present), lift p42 to a cohort-hazard logistic's
    probability when it is higher. mode 'silent' lifts only devices the voltage
    model scores low (<0.2) - the miss bucket - leaving confident scores alone.
    """

    CFEATS = ['cohort_rel_age_m', 'cohort_overdue_m', 'age_days']

    def __init__(self, mode='all'):
        self.mode = mode

    def fit(self, train):
        self.base = SklearnRUL('logbag').fit(train)
        cov = train[train['cohort_rel_age_m'].notna()]
        self.cmodel = None
        if len(cov) >= 300 and cov['dies42'].sum() >= 10:
            self.cmodel = make_pipeline(SimpleImputer(strategy='median'),
                                        StandardScaler(),
                                        LogisticRegression(max_iter=2000, C=1.0))
            self.cmodel.fit(cov[self.CFEATS], cov['dies42'])
        return self

    def predict(self, rows):
        out = self.base.predict(rows)
        if self.cmodel is None:
            return out
        covered = rows['cohort_rel_age_m'].notna().values
        if not covered.any():
            return out
        pc = np.zeros(len(rows))
        pc[covered] = self.cmodel.predict_proba(
            rows.loc[covered, self.CFEATS])[:, 1]
        p = out['p42'].values
        lift = covered & (pc > p)
        if self.mode == 'silent':
            lift &= p < 0.2
        out['p42'] = np.where(lift, pc, p)
        return out


class LogbagDH:
    """
    Champion WHO (logbag p42, identical training pairs) + hazard WHEN: day
    quantiles from DiscreteHazard's conditional bucket CDF instead of the
    4-knot GBR quantile models. Isolates the value of a better death-day
    estimator without touching inclusion.
    """

    def fit(self, train):
        self.who = SklearnRUL('logbag').fit(train)
        aug = train.copy()
        aug['duration'] = train['days_to_eol'].fillna(10000.0).clip(lower=0.5)
        aug['event'] = train['dies42'].astype(int)
        self.when = DiscreteHazard().fit(aug)
        return self

    def predict(self, rows):
        out = self.when.predict(rows)          # day_q from hazard CDF
        out['p42'] = self.who.predict(rows)['p42'].values
        return out


MODELS = {
    'rank-v': RankVoltage,
    'logistic': lambda: SklearnRUL('logistic'),
    'logcal': lambda: SklearnRUL('logcal'),
    'logbag': lambda: SklearnRUL('logbag'),
    'gbm': lambda: SklearnRUL('gbm'),
    'aft': WeibullAFT,
    'logbag_rw120': lambda: WeightedBagLogistic(half_life=120.0),
    'logbag_rw240': lambda: WeightedBagLogistic(half_life=240.0),
    'dthaz': DiscreteHazard,
    'blend': BlendRUL,
    'logbag_c03': lambda: LogBagC(0.3),
    'logbag_c3': lambda: LogBagC(3.0),
    'logbag_c01': lambda: LogBagC(0.1),
    'logbag_open': LogBagOpen,
    'blend_dh': BlendDH,
    'logbag_aux': lambda: LogBagFeats(['rng_90', 'vstd_30', 'rng_ratio', 'gap_max_60']),
    'logbag_aux2': lambda: LogBagFeats(['rng_30', 'rng_90', 'vstd_30']),
    'rank_rf': RankBlendRF,
    'silentfix015': lambda: SilentFix(0.15),
    'silentfix03': lambda: SilentFix(0.3),
    'silentfix06': lambda: SilentFix(0.6),
    'logbag_proj': lambda: LogBagFeats(['v_proj']),
    'logbag_rank': lambda: LogBagFeats(['r_vlast', 'r_tau']),
    'logbag_rankproj': lambda: LogBagFeats(['v_proj', 'r_vlast', 'r_tau']),
    'aux_rv': lambda: LogBagFeats(['rng_90', 'vstd_30']),
    'aux_rvg': lambda: LogBagFeats(['rng_90', 'vstd_30', 'gap_max_60']),
    'aux_rvr': lambda: LogBagFeats(['rng_90', 'vstd_30', 'rng_ratio']),
    'aux_all': lambda: LogBagFeats(['rng_90', 'vstd_30', 'rng_ratio', 'gap_max_60',
                                    'hot_90', 'oob_90']),
    'aux_proj': lambda: LogBagFeats(['rng_90', 'vstd_30', 'rng_ratio', 'gap_max_60',
                                     'v_proj']),
    'aux_rank': lambda: LogBagFeats(['rng_90', 'vstd_30', 'rng_ratio', 'gap_max_60',
                                     'r_vlast', 'r_tau']),
    'aux_bag30': lambda: LogBagFeats(['rng_90', 'vstd_30', 'rng_ratio', 'gap_max_60'], n_estimators=30),
    'aux_bag50': lambda: LogBagFeats(['rng_90', 'vstd_30', 'rng_ratio', 'gap_max_60'], n_estimators=50),
    'dthaz_aux': lambda: DiscreteHazard(['rng_90', 'vstd_30', 'rng_ratio', 'gap_max_60']),
    'blend_dhaux': BlendDHAux,
    'aux_ucb05': lambda: LogBagUCB(['rng_90', 'vstd_30', 'rng_ratio', 'gap_max_60'], k=0.5),
    'aux_ucb10': lambda: LogBagUCB(['rng_90', 'vstd_30', 'rng_ratio', 'gap_max_60'], k=1.0),
    'aux_log': lambda: LogBagLog(['rng_90', 'vstd_30', 'rng_ratio', 'gap_max_60']),
    'wiener': WienerFHT,
    'lw_mean': lambda: LogWienerBlend('mean'),
    'lw_stale': lambda: LogWienerBlend('stale'),
    'crt_in15': lambda: ConformalCut(0.15, 'in'),
    'crt_in25': lambda: ConformalCut(0.25, 'in'),
    'crt_sp15': lambda: ConformalCut(0.15, 'split'),
    'crt_in40': lambda: ConformalCut(0.40, 'in'),
    'crt_in50': lambda: ConformalCut(0.50, 'in'),
    'logbag_fht': lambda: LogBagFeats(['fht_days', 'wnr_m', 'wnr_s2']),
    'logbag_fht2': lambda: LogBagFeats(['fht_days']),
    'cg_all': lambda: CohortGate('all'),
    'cg_silent': lambda: CohortGate('silent'),
    'logbag_dh': LogbagDH,
    'logbag_lah': LogbagLAH,
    'logbag_lah_ds': lambda: LogbagLAH(domain_split=True),
    'logbag_cw': CostWeightBag,
}


def resolve_model(name):
    """'logbag:bags=30,q6=1' -> SklearnRUL variant; plain names via MODELS."""
    if name in MODELS:
        return MODELS[name]
    return lambda: SklearnRUL(name)


# ------------------------------------------------------------------------- planners

def assemble(selected_days, locations, start_time, sort_key=None):
    parked_day = start_time + pd.Timedelta(days=WINDOW_DAYS + 1)
    scheduled = set(selected_days.index)
    parked = [b for b in locations['battery'] if b not in scheduled]
    plan = pd.concat([
        pd.DataFrame({'day': selected_days.values, 'battery': selected_days.index}),
        pd.DataFrame({'day': parked_day, 'battery': parked}),
    ], ignore_index=True)
    # within a day the evaluator visits in row order: group by building then room,
    # otherwise interleaved buildings pay repeated building-changes and travel.
    # sort_key overrides the building order (per-battery rank, e.g. NN route);
    # parked rows fall past end_time and are cut by the evaluator, so their
    # order is irrelevant.
    loc = locations.set_index('battery')
    plan['building'] = plan['battery'].map(loc['building'])
    plan['room'] = plan['battery'].map(loc['room'])
    if sort_key is not None:
        plan['building'] = plan['battery'].map(sort_key).fillna(10 ** 9)
    plan = (plan.sort_values(['day', 'building', 'room'], kind='stable')
                .drop(columns=['building', 'room']).reset_index(drop=True))
    check_plan_valid(plan, locations, start_time=start_time)
    return plan


def plan_onepd(pred, included, locations, start_time, day_q):
    order = pred.loc[included, 'p42'].sort_values(ascending=False)
    days = start_time + pd.to_timedelta(np.arange(len(order)), unit='D')
    return assemble(pd.Series(days, index=order.index), locations, start_time)


def plan_sched(pred, included, locations, start_time, day_q):
    offs = pred.loc[included, f'day_q{day_q}'].round().clip(0, WINDOW_DAYS)
    days = start_time + pd.to_timedelta(offs.values, unit='D')
    return assemble(pd.Series(days, index=offs.index), locations, start_time)


def plan_batch(pred, included, locations, start_time, day_q, max_groups_per_day=2):
    offs = pred.loc[included, f'day_q{day_q}'].round().clip(0, WINDOW_DAYS)
    bldg = locations.set_index('battery')['building'].reindex(offs.index)

    # one visit per building at that group's earliest predicted day
    group_day = offs.groupby(bldg).min().sort_values()

    # cap building-groups per day; push overflow to the next free day
    used = {}
    final = {}
    for building, day in group_day.items():
        d = int(day)
        while used.get(d, 0) >= max_groups_per_day and d < WINDOW_DAYS:
            d += 1
        used[d] = used.get(d, 0) + 1
        final[building] = d

    offs = bldg.map(final).astype(float)
    days = start_time + pd.to_timedelta(offs.values, unit='D')
    return assemble(pd.Series(days, index=offs.index), locations, start_time)


def _batch_days(offs, locations, start_time, max_groups_per_day=2):
    """Shared batching pass: one visit per building at the group's earliest day."""
    bldg = locations.set_index('battery')['building'].reindex(offs.index)
    group_day = offs.groupby(bldg).min().sort_values()
    used, final = {}, {}
    for building, day in group_day.items():
        d = int(day)
        while used.get(d, 0) >= max_groups_per_day and d < WINDOW_DAYS:
            d += 1
        used[d] = used.get(d, 0) + 1
        final[building] = d
    offs = bldg.map(final).astype(float)
    days = start_time + pd.to_timedelta(offs.values, unit='D')
    return assemble(pd.Series(days, index=offs.index), locations, start_time)


EARLY, LATE = 0.5, 10.0
# The evaluator's true emergency queue opens at the END OF THE WEEK containing
# end_time (day 48 for the all-Monday train scenarios), +1 day per queued miss.
# Exact pricing (48/50/52) was tested and REJECTED (phase 21): it only flips
# decisions at near-data-end origins and loses there - the underprice acts as
# a regularizer against over-including marginal survivors. Keep 45.
EMERGENCY_DAY = WINDOW_DAYS + 3
# The evaluator prices a swap on a device with no observed EOL against
# (location end_time + settings.unobserved_eol_days). That setting is
# per-scenario and READ FROM THE SCENARIO at runtime; it was hardcoded to the
# default 30.0 in two places here, and submit.py imports plan_tsroute from
# this module, so the hardcode shipped. Every train scenario happens to use
# 30.0, which is exactly why local CV could never catch it - same failure
# shape as the hardcoded 42-day park that repark_for_window exists to undo.
# plan_tsroute now overwrites this from the settings it is handed.
UNOBSERVED_EOL_DAYS = 30.0
CENSOR_AWARE = True     # price unobservable (post-data-end) deaths via imputation
PARK_MOBILITY = 0.0     # extra cost per expected emergency (travel/overtime)


def plan_expcost(pred, included, locations, start_time, day_q):
    """
    Include/park and swap-day both chosen by expected cost under the model's
    death distribution. `included` and `day_q` are ignored.
    """
    t = np.arange(1, WINDOW_DAYS + 1)
    pmf = pred[[f'pd_{i}' for i in t]].values          # devices x 42
    p_surv = np.clip(1.0 - pmf.sum(axis=1), 0.0, 1.0)
    e_far = pred['e_far'].values

    d = np.arange(0, WINDOW_DAYS + 1)
    delta = d[None, :, None] - t[None, None, :]        # 1 x 43 x 42
    swap_penalty = np.where(delta >= 0, LATE * delta, EARLY * (-delta))
    cost = (pmf[:, None, :] * swap_penalty).sum(axis=2)          # devices x 43
    cost += p_surv[:, None] * EARLY * np.maximum(e_far[:, None] - d[None, :], 0)

    cost_park = (pmf * (LATE * (EMERGENCY_DAY - t))[None, :]).sum(axis=1)

    best_day = cost.argmin(axis=1)
    include = cost[np.arange(len(cost)), best_day] < cost_park
    offs = pd.Series(best_day[include].astype(float), index=pred.index[include])
    if offs.empty:
        return assemble(pd.Series(dtype='datetime64[ns]'), locations, start_time)
    return _batch_days(offs, locations, start_time)


def plan_topsum(pred, included, locations, start_time, day_q):
    """
    Adaptive inclusion: take the top-N batteries by p42 with N = alpha * sum(p42).
    sum(p42) tracks the expected death count even when the probability scale
    drifts between refits, unlike a fixed threshold. alpha arrives via p_star.
    """
    alpha = plan_topsum.alpha
    n = int(round(alpha * pred['p42'].sum()))
    chosen = pred['p42'].nlargest(n).index
    offs = pred.loc[chosen, f'day_q{day_q}'].round().clip(0, WINDOW_DAYS)
    if offs.empty:
        return assemble(pd.Series(dtype='datetime64[ns]'), locations, start_time)
    return _batch_days(offs, locations, start_time)


plan_topsum.alpha = 1.2


# ---------------------------------------------------------- expected-cost planning

def cost_curves(pred, locations, start_time, blend=0.0):
    """
    Per-battery expected timing cost for swapping on each day d in 0..42, plus the
    cost of parking (emergency if it dies in-window).

    Death-day CDF within the window is piecewise-linear through the predicted
    conditional quantiles. The surviving mass pays early-swap at 0.5/day against the
    evaluator's imputed EOL (device end_time + unobserved_eol_days).
    """
    qs = list(DAY_QUANTILES)
    tq = pred[[f'day_q{q}' for q in qs]].values          # devices x quantiles
    tq = np.maximum.accumulate(np.clip(tq, 0, WINDOW_DAYS), axis=1)  # monotone
    p42 = pred['p42'].values

    grid = np.arange(0, WINDOW_DAYS + 1, dtype=float)
    n = len(pred)
    cdf = np.empty((n, len(grid)))
    xs = np.concatenate([[0.0], np.zeros(len(qs)), [WINDOW_DAYS]])
    ys = np.array([0.0] + qs + [1.0])
    for i in range(n):
        xs[1:-1] = tq[i]
        cdf[i] = np.interp(grid, xs, ys)

    if blend > 0:
        # the day models are trained on typical dying batteries and predict "dies
        # soon" too confidently for marginal ones; pull low-p42 batteries'
        # conditional CDF toward uniform, which defers their optimal swap day
        w = 1.0 - blend * (1.0 - p42[:, None])
        uniform = (grid / WINDOW_DAYS)[None, :]
        cdf = w * cdf + (1.0 - w) * uniform

    pmf = np.diff(cdf, axis=1)                            # devices x 42, mass at t=1..42
    t = np.arange(1, WINDOW_DAYS + 1, dtype=float)

    end = locations.set_index('battery')['end_time'].dt.tz_localize(None)
    # unclipped: for long-stale devices the imputed EOL (last data + 30d) can
    # precede the window, and the evaluator then charges a swap as LATE at
    # 10/day - the old clip(lower=43) made exactly those devices look free
    e_far = ((end.reindex(pred.index) - start_time).dt.total_seconds() / 86400
             + UNOBSERVED_EOL_DAYS).values

    # censoring-aware: deaths after the split's data end are never in the
    # shipped eol_times, so the evaluator cannot charge late/emergency for
    # them - it prices those devices via the end_time+30 imputation instead
    # (i.e. our survivor branch). Move unobservable death mass to the
    # survivor branch and drop it from the parking (emergency) estimate.
    # Inert whenever data extends >= 42d past the origin (all train panel
    # origins); bites on near-data-end scenarios (runtime splits, s_42-47).
    if CENSOR_AWARE:
        D = (end.max() - start_time).total_seconds() / 86400
        pmf = pmf * (t <= D)[None, :]
    obs_frac = pmf.sum(axis=1)

    delta = grid[None, :, None] - t[None, None, :]        # 1 x 43 x 42
    pen = np.where(delta >= 0, LATE * delta, EARLY * (-delta))
    curves = p42[:, None] * (pmf[:, None, :] * pen).sum(axis=2)
    surv_delta = e_far[:, None] - grid[None, :]
    curves += (1 - p42 * obs_frac)[:, None] * np.where(surv_delta >= 0,
                                                       EARLY * surv_delta,
                                                       LATE * -surv_delta)

    # an emergency also pays its own travel/overtime, not just late-days;
    # PARK_MOBILITY adds that per expected observable emergency
    park = (p42 * (pmf * (LATE * (EMERGENCY_DAY - t))[None, :]).sum(axis=1)
            + PARK_MOBILITY * p42 * obs_frac)
    return curves, park


def plan_ecday(pred, included, locations, start_time, day_q):
    """Include iff min-day expected cost beats parking; swap at the argmin day."""
    curves, park = cost_curves(pred, locations, start_time)
    best_day = curves.argmin(axis=1)
    keep = curves[np.arange(len(curves)), best_day] < park
    offs = pd.Series(best_day[keep].astype(float), index=pred.index[keep])
    if offs.empty:
        return assemble(pd.Series(dtype='datetime64[ns]'), locations, start_time)
    return _batch_days(offs, locations, start_time)


# --------------------------------------------------------------- mobility simulator

def mobility_cost(day_batteries, locations, travel_hours, settings, start_time):
    """
    Travel + overtime + daily/weekly limit cost of a visit schedule, replicating
    evaluate_plan's state machine exactly (verified against it with eol_times=None).
    day_batteries: {day_offset: [battery, ...]} with per-day visit order preserved.
    """
    base, base_room = settings.base_location, settings.base_room
    loc = locations.set_index('battery')
    total = 0.0
    room = base_room
    week_used = 0.0
    week_start = 0

    # Evaluator quirk, replicated deliberately: end_previous_day computes the old
    # day's overtime including the return-to-base leg, but do_action then adds that
    # same leg to the NEW day's clock - so each return leg counts toward both
    # adjacent visit-days' overtime, and lands in the week of the *next* visit day.
    carry = None  # return travel awaiting attribution; None until first transition

    end_day = WINDOW_DAYS + (6 - (start_time + pd.Timedelta(days=WINDOW_DAYS)).weekday())
    for d in range(0, end_day + 1):
        if d >= week_start + 7:
            if week_used >= settings.worker_limit_weekly_hours:
                total += settings.worker_limit_weekly_penalty
            week_used, week_start = 0.0, d
        batts = day_batteries.get(d, [])
        if not batts:
            continue
        if carry is None:
            # first visit day: if not day 0, the initial end_previous_day travels
            # base->base
            carry = 0.0 if d == 0 else float(travel_hours.at[base, base])
            total += carry
        hours = 0.0
        building = base
        for b in batts:
            tb, tr = loc.at[b, 'building'], loc.at[b, 'room']
            if tb != building:
                hours += settings.time_per_building_change_hours + travel_hours.at[building, tb]
                building = tb
            if tr != room:
                hours += settings.time_per_room_change_hours
                room = tr
            hours += settings.time_per_battery_hours
        ret = float(travel_hours.at[building, base])
        clock = carry + hours + ret
        total += hours + ret
        total += max(clock - settings.overtime_start, 0.0) * settings.overtime_penalty_factor
        if clock > settings.worker_limit_daily_hours:
            total += settings.worker_limit_daily_penalty
        week_used += carry + hours
        carry = ret
    if carry:
        week_used += carry
    if week_used >= settings.worker_limit_weekly_hours:
        total += settings.worker_limit_weekly_penalty
    return total



def _nn_building_order(bldgs, travel_hours, base):
    """Nearest-neighbor route over a day's buildings starting from base:
    {building: visit_rank}. Deterministic (ties broken by name)."""
    left = sorted(set(bldgs))
    order, cur = {}, base
    for i in range(len(left)):
        nxt = min(left, key=lambda b: (float(travel_hours.at[cur, b]), b))
        order[nxt] = i
        left.remove(nxt)
        cur = nxt
    return order


def optimize_days(groups, curves, locations, travel_hours, settings, start_time,
                  restarts=None, iters=None):
    restarts = restarts if restarts is not None else optimize_days.restarts
    iters = iters if iters is not None else optimize_days.iters
    """
    Local search over building-group visit days: shift moves, merge moves (jump a
    group onto another group's day), and deterministic random restarts.
    Objective = expected timing cost (model curves) + exact mobility cost.
    """
    loc_b = locations.set_index('battery')['building']
    gbldg = {g: loc_b.at[bs[0]] for g, bs in groups.items()}

    def timing(day_of):
        return sum(curves.at[b, day_of[g]] for g, bs in groups.items() for b in bs)

    def mobility(day_of):
        day_groups = {}
        for g in sorted(day_of, key=lambda g: (day_of[g], g)):
            day_groups.setdefault(day_of[g], []).append(g)
        if optimize_days.nn_order:
            # visit a day's buildings in nearest-neighbor route order instead
            # of name order (matters only on 3+ building days: with a
            # symmetric travel matrix a 2-building day is order-invariant)
            for gs in day_groups.values():
                rank = _nn_building_order([gbldg[g] for g in gs], travel_hours,
                                          settings.base_location)
                gs.sort(key=lambda g: (rank[gbldg[g]], g))
        day_batteries = {d: [b for g in gs for b in groups[g]]
                         for d, gs in day_groups.items()}
        return mobility_cost(day_batteries, locations, travel_hours,
                             settings, start_time)

    def climb(day_of):
        score = timing(day_of) + mobility(day_of)
        for _ in range(iters):
            improved = False
            for g in groups:
                others = {day_of[o] for o in groups if o != g}
                candidates = ({day_of[g] + s for s in
                               (-10, -7, -5, -3, -2, -1, 1, 2, 3, 5, 7, 10)} | others)
                for d in candidates:
                    if not 0 <= d <= WINDOW_DAYS or d == day_of[g]:
                        continue
                    trial = dict(day_of, **{g: int(d)})
                    s = timing(trial) + mobility(trial)
                    if s < score - 1e-9:
                        day_of, score, improved = trial, s, True
            if not improved:
                break
        return day_of, score

    base = {g: int(min(curves.loc[bs].values.argmin(axis=1)))
            for g, bs in groups.items()}
    best, best_score = climb(base)
    rng = np.random.RandomState(0)
    for _ in range(restarts):
        start = {g: int(np.clip(base[g] + rng.randint(-6, 7), 0, WINDOW_DAYS))
                 for g in groups}
        cand, s = climb(start)
        if s < best_score:
            best, best_score = cand, s

    if optimize_days.constructive:
        # limit-aware constructive start (phase 18): the argmin start piles
        # urgent groups into the same days/weeks and the +-10d shift moves may
        # not escape the resulting overtime/limit basin. Greedily place groups
        # in urgency order, each on the day minimizing the EXACT partial cost
        # (timing + real mobility incl. limit penalties), then climb from it.
        cons = {}
        for g in sorted(groups, key=base.get):
            best_d, best_s = None, np.inf
            for dcand in range(0, WINDOW_DAYS + 1):
                trial = dict(cons, **{g: dcand})
                s = (sum(curves.at[b, trial[gg]] for gg in trial for b in groups[gg])
                     + mobility({gg: trial[gg] for gg in trial}))
                if s < best_s:
                    best_d, best_s = dcand, s
            cons[g] = best_d
        cand, s = climb(cons)
        if s < best_score:
            best, best_score = cand, s
    return best


optimize_days.restarts = 2
optimize_days.iters = 60
optimize_days.constructive = False
optimize_days.nn_order = False   # visit a day's buildings in NN route order


def plan_ecroute(pred, included, locations, start_time, day_q,
                 travel_costs=None, settings=None, margin=None):
    """
    ecday selection, then local search over building-group days minimizing
    (exact mobility cost) + (expected timing cost from the model's curves).
    margin > 1 includes more marginal batteries (insurance); reuses p_star slot.
    """
    margin = margin if margin else (plan_ecroute.margin or 1.0)
    curves, park = cost_curves(pred, locations, start_time,
                               blend=plan_ecroute.blend)
    best_day = curves.argmin(axis=1)
    keep = curves[np.arange(len(curves)), best_day] < margin * park
    chosen = pred.index[keep]
    if len(chosen) == 0:
        return assemble(pd.Series(dtype='datetime64[ns]'), locations, start_time)

    curves = pd.DataFrame(curves[keep], index=chosen)
    bldg = locations.set_index('battery')['building'].reindex(chosen)
    room = locations.set_index('battery')['room']
    groups = {g: sorted(members.index, key=lambda b: room.get(b, ''))
              for g, members in curves.groupby(bldg)}
    travel_hours = travel_costs.set_index(['from', 'to'])['hours'].unstack()

    day_of = optimize_days(groups, curves, locations, travel_hours,
                           settings, start_time)

    offs = bldg.map(day_of).astype(float)
    days = start_time + pd.to_timedelta(offs.values, unit='D')
    return assemble(pd.Series(days, index=offs.index), locations, start_time)


plan_ecroute.margin = 1.0
plan_ecroute.blend = 0.0


def plan_tsroute(pred, included, locations, start_time, day_q,
                 travel_costs=None, settings=None):
    """
    Hybrid: topsum inclusion (count-based, robust to probability-scale drift)
    with expected-cost day placement and mobility routing.

    p_be:  break-even top-up - also include every battery with p42 >= p_be
           (a miss costs ~4x a false inclusion, so p_be ~ 0.2 is +EV) [0 = off]
    floor_day/p_conf: batteries with p42 < p_conf may not be scheduled before
           floor_day - defers uncertain insurance swaps that would otherwise
           pay early-replacement for weeks [floor_day 0 = off]
    """
    if settings is not None:
        global UNOBSERVED_EOL_DAYS
        UNOBSERVED_EOL_DAYS = float(settings.unobserved_eol_days)
    psum = pred['p42_sum'].iloc[0] if 'p42_sum' in pred else pred['p42'].sum()
    if plan_tsroute.floor_scale > 0:
        # cold-start inclusion floor: sum(p42) tracks prevalence only once the
        # training era has seen it - at warm-up origins it undercounts ~4x and
        # the count rule parks visibly-dying batteries (46% of all-48 late
        # cost). Deaths RECORDED in the last 42d are observable at plan time
        # (the harness itself drops eol <= origin devices) and estimate the
        # coming window's prevalence without trusting the probability scale.
        psum = max(psum, plan_tsroute.floor_scale * plan_tsroute.recent_deaths)
    # consume-once appetite boost: armed by the model's predict (T3 warmup
    # gate) for the plan that immediately follows, then self-resets so a
    # fallback path can never inherit a stale boost
    boost = plan_tsroute.alpha_boost
    plan_tsroute.alpha_boost = 1.0
    n = int(round(plan_tsroute.alpha * boost * psum))
    order_col = 'rank_score' if 'rank_score' in pred.columns else 'p42'
    chosen = pred[order_col].nlargest(n).index
    if plan_tsroute.p_be > 0:
        topups = pred.index[pred['p42'] >= plan_tsroute.p_be]
        chosen = chosen.union(topups)
    curves, park = cost_curves(pred, locations, start_time)
    if plan_tsroute.ec_margin > 0:
        # break-even top-up: include every battery whose best swap day is
        # cheaper than parking it. The false-swap price is 0.5/day against the
        # imputed EOL (end_time+30), so this set is tiny when data extends far
        # past the window (train dev block) and widens automatically when the
        # scenario sits near the split's data end (runtime), where false swaps
        # cost almost nothing but misses still cost 10/day.
        # ec_pmax guards the p42/curve estimation error: the top-up only fires
        # for batteries whose worst-case cost of a wrong inclusion (false swap
        # at day 42) is at most ec_pmax - EV says include, price caps the risk.
        # Real-dev CV rejected the ungated rule (+24, worst origin +461); the
        # counterfactual cheap-price CV shows -76 where wrong costs ~19.
        end = locations.set_index('battery')['end_time'].dt.tz_localize(None)
        e_far = ((end.reindex(pred.index) - start_time).dt.total_seconds() / 86400
                 + UNOBSERVED_EOL_DAYS).values
        # top-up only fresh devices (imputed EOL beyond the window): for stale
        # ones the survivor branch is a bookkeeping fiction - swapping "on the
        # imputed day" scores ~0 but buys nothing and costs real mobility
        # (panel: admitting them turned s_23/s_41 from ties into losses)
        fp_price = EARLY * (e_far - WINDOW_DAYS)
        be = pred.index[(e_far >= WINDOW_DAYS)
                        & (curves.min(axis=1) < plan_tsroute.ec_margin * park)
                        & (fp_price <= plan_tsroute.ec_pmax)]
        chosen = chosen.union(be)
    if len(chosen) == 0:
        return assemble(pd.Series(dtype='datetime64[ns]'), locations, start_time)

    if plan_tsroute.floor_day > 0:
        uncertain = (pred['p42'] < plan_tsroute.p_conf).values
        curves[np.ix_(uncertain, np.arange(plan_tsroute.floor_day))] = np.inf
    curves = pd.DataFrame(curves, index=pred.index).loc[chosen]
    bldg = locations.set_index('battery')['building'].reindex(chosen)
    room = locations.set_index('battery')['room']
    key = bldg
    if plan_tsroute.split_days > 0:
        # one visit day per building forces urgent and insurance swaps in the
        # same building onto a compromise day; sub-group by best-day bucket so
        # the optimizer MAY split a building across days (merge moves can
        # always rejoin them at zero travel cost - same-day groups share the
        # mobility simulation)
        bucket = curves.values.argmin(axis=1) // int(plan_tsroute.split_days)
        key = bldg.astype(str) + '~' + pd.Series(bucket, index=curves.index).astype(str)
    groups = {g: sorted(m.index, key=lambda b: room.get(b, ''))
              for g, m in curves.groupby(key)}
    travel_hours = travel_costs.set_index(['from', 'to'])['hours'].unstack()

    day_of = optimize_days(groups, curves * plan_tsroute.timing_weight,
                           locations, travel_hours, settings, start_time)

    offs = key.map(day_of).astype(float)
    days = start_time + pd.to_timedelta(offs.values, unit='D')
    sort_key = None
    if optimize_days.nn_order:
        # emit rows in the same NN building order the mobility sim priced
        base_b = settings.base_location
        gb = {g: bldg.loc[bs[0]] for g, bs in groups.items()}
        ranks = {}
        for d in set(day_of.values()):
            gs = [g for g in day_of if day_of[g] == d]
            ranks[d] = _nn_building_order([gb[g] for g in gs], travel_hours, base_b)
        sort_key = pd.Series({b: ranks[day_of[g]][gb[g]]
                              for g, bs in groups.items() for b in bs})
    return assemble(pd.Series(days, index=offs.index), locations, start_time,
                    sort_key=sort_key)


plan_tsroute.alpha = 0.8
plan_tsroute.split_days = 0
plan_tsroute.timing_weight = 1.0   # scales timing curves in optimize_days' objective
plan_tsroute.alpha_boost = 1.0     # consume-once; see the count line above
plan_tsroute.p_be = 0.0
plan_tsroute.floor_day = 0
plan_tsroute.p_conf = 0.5
plan_tsroute.ec_margin = 0.0
plan_tsroute.ec_pmax = 25.0
plan_tsroute.floor_scale = 0.0   # 0 = off; k in max(sum_p42, k * recent_deaths)
plan_tsroute.recent_deaths = 0.0  # deaths recorded in (origin-42d, origin]; set per origin

def plan_tspig(pred, included, locations, start_time, day_q,
               travel_costs=None, settings=None):
    """
    tsroute + exact-marginal membership repair. After the day search:
      drop pass - remove any included battery whose expected timing benefit at its
                  assigned day is worse than parking, counting the exact mobility
                  saving (a group that empties refunds its whole trip);
      add pass  - piggyback parked batteries onto buildings already being visited
                  when benefit at that day exceeds the exact mobility delta
                  (usually just battery-handling time), then whole new building
                  groups when their pooled benefit pays for the trip.
    Inclusion thereby becomes routing-aware instead of count-only.
    """
    n = int(round(plan_tspig.alpha * pred['p42'].sum()))
    chosen = pred['p42'].nlargest(n).index
    curves_np, park_np = cost_curves(pred, locations, start_time)
    curves_all = pd.DataFrame(curves_np, index=pred.index)
    park = pd.Series(park_np, index=pred.index)
    loc = locations.set_index('battery')
    bldg_all, room_all = loc['building'], loc['room']
    travel_hours = travel_costs.set_index(['from', 'to'])['hours'].unstack()

    if len(chosen) == 0:
        return assemble(pd.Series(dtype='datetime64[ns]'), locations, start_time)

    curves = curves_all.loc[chosen]
    bldg = bldg_all.reindex(chosen)
    groups = {g: sorted(m.index, key=lambda b: room_all.get(b, ''))
              for g, m in curves.groupby(bldg)}
    day_of = optimize_days(groups, curves, locations, travel_hours,
                           settings, start_time)

    def mobility(gs, days):
        day_batteries = {}
        for g in sorted(days, key=lambda g: (days[g], g)):
            if gs.get(g):
                day_batteries.setdefault(days[g], []).extend(gs[g])
        return mobility_cost(day_batteries, locations, travel_hours,
                             settings, start_time)

    cur_mob = mobility(groups, day_of)
    in_plan = set().union(*groups.values())

    for _ in range(4):
        improved = False
        # drop pass
        for g in list(groups):
            for b in list(groups[g]):
                trial = dict(groups, **{g: [x for x in groups[g] if x != b]})
                dm = mobility(trial, day_of) - cur_mob
                if park.at[b] - curves_all.at[b, day_of[g]] + dm < -1e-9:
                    groups, cur_mob = trial, cur_mob + dm
                    in_plan.discard(b)
                    improved = True
        # piggyback adds within visited buildings
        active = {g for g in groups if groups[g]}
        cands = [b for b in pred.index
                 if b not in in_plan and bldg_all.get(b) in active]
        for b in sorted(cands, key=lambda b: -pred.at[b, 'p42']):
            g = bldg_all.get(b)
            if not groups.get(g):
                continue
            trial = dict(groups, **{g: sorted(groups[g] + [b],
                                              key=lambda x: room_all.get(x, ''))})
            dm = mobility(trial, day_of) - cur_mob
            if curves_all.at[b, day_of[g]] - park.at[b] + dm < -1e-9:
                groups, cur_mob = trial, cur_mob + dm
                in_plan.add(b)
                improved = True
        # whole new building groups whose pooled benefit pays for the trip
        outside = pred.index.difference(list(in_plan))
        for g, members in bldg_all.reindex(outside).groupby(bldg_all.reindex(outside)):
            if groups.get(g):
                continue
            gain = park.reindex(members.index) - curves_all.loc[members.index].min(axis=1)
            keep = gain[gain > 0].index
            if len(keep) == 0:
                continue
            day = int(curves_all.loc[keep].sum(axis=0).values.argmin())
            trial = dict(groups, **{g: sorted(keep, key=lambda x: room_all.get(x, ''))})
            tdays = dict(day_of, **{g: day})
            dm = mobility(trial, tdays) - cur_mob
            delta = (curves_all.loc[keep, day] - park.reindex(keep)).sum() + dm
            if delta < -1e-9:
                groups, day_of, cur_mob = trial, tdays, cur_mob + dm
                in_plan.update(keep)
                improved = True
        if not improved:
            break

    groups = {g: bs for g, bs in groups.items() if bs}
    if not groups:
        return assemble(pd.Series(dtype='datetime64[ns]'), locations, start_time)
    # final day re-optimization with the repaired membership
    day_of = optimize_days(groups, curves_all.loc[sorted(in_plan)], locations,
                           travel_hours, settings, start_time)
    offs, idx = [], []
    for g, bs in groups.items():
        for b in bs:
            idx.append(b)
            offs.append(float(day_of[g]))
    days = start_time + pd.to_timedelta(offs, unit='D')
    return assemble(pd.Series(days, index=idx), locations, start_time)


plan_tspig.alpha = 0.8


def plan_tspig2(pred, included, locations, start_time, day_q,
                travel_costs=None, settings=None):
    """
    Rank-gated piggyback: run tsroute on the usual top-N (alpha), then extend
    inclusion to batteries ranked within alpha2 * sum(p42) whose building is
    already visited, scheduling them on that building's existing day. Extra
    inclusions are chosen by rank (scale-free), not expected value; only the
    mobility context decides whether the insurance is cheap enough.
    """
    n = int(round(plan_tspig2.alpha * pred['p42'].sum()))
    chosen = pred['p42'].nlargest(n).index
    if len(chosen) == 0:
        return assemble(pd.Series(dtype='datetime64[ns]'), locations, start_time)

    curves_np, _ = cost_curves(pred, locations, start_time)
    curves_all = pd.DataFrame(curves_np, index=pred.index)
    loc = locations.set_index('battery')
    bldg_all, room_all = loc['building'], loc['room']
    travel_hours = travel_costs.set_index(['from', 'to'])['hours'].unstack()

    curves = curves_all.loc[chosen]
    bldg = bldg_all.reindex(chosen)
    groups = {g: sorted(m.index, key=lambda b: room_all.get(b, ''))
              for g, m in curves.groupby(bldg)}
    day_of = optimize_days(groups, curves, locations, travel_hours,
                           settings, start_time)

    n2 = int(round(plan_tspig2.alpha2 * pred['p42'].sum()))
    extras = pred['p42'].nlargest(n2).index.difference(chosen)
    offs, idx = [], []
    for g, bs in groups.items():
        for b in bs:
            idx.append(b); offs.append(float(day_of[g]))
    for b in extras:
        g = bldg_all.get(b)
        if g in day_of:
            idx.append(b); offs.append(float(day_of[g]))
    days = start_time + pd.to_timedelta(offs, unit='D')
    return assemble(pd.Series(days, index=idx), locations, start_time)


plan_tspig2.alpha = 0.8
plan_tspig2.alpha2 = 1.3

def plan_tscut(pred, included, locations, start_time, day_q,
               travel_costs=None, settings=None):
    """tsroute scheduling, but inclusion = p42 >= the model's conformal cutoff."""
    chosen = pred.index[pred['p42'] >= pred['p_cut'].iloc[0]]
    if len(chosen) == 0:
        return assemble(pd.Series(dtype='datetime64[ns]'), locations, start_time)
    curves, _ = cost_curves(pred, locations, start_time)
    curves = pd.DataFrame(curves, index=pred.index).loc[chosen]
    bldg = locations.set_index('battery')['building'].reindex(chosen)
    room = locations.set_index('battery')['room']
    groups = {g: sorted(m.index, key=lambda b: room.get(b, ''))
              for g, m in curves.groupby(bldg)}
    travel_hours = travel_costs.set_index(['from', 'to'])['hours'].unstack()
    day_of = optimize_days(groups, curves, locations, travel_hours,
                           settings, start_time)
    offs = bldg.map(day_of).astype(float)
    days = start_time + pd.to_timedelta(offs.values, unit='D')
    return assemble(pd.Series(days, index=offs.index), locations, start_time)


PLANNERS ={'onepd': plan_onepd, 'sched': plan_sched, 'batch': plan_batch,
            'expcost': plan_expcost, 'topsum': plan_topsum,
            'ecday': plan_ecday, 'ecroute': plan_ecroute, 'tsroute': plan_tsroute,
            'tspig': plan_tspig, 'tspig2': plan_tspig2,
            'tscut': plan_tscut}


# ------------------------------------------------------------------------ harness

def training_pairs(table, origin_time):
    closed = table['origin'] <= origin_time - pd.Timedelta(days=WINDOW_DAYS)
    sub = table[closed]
    if QUARANTINE and 'v_corr_last' in sub.columns:
        # drop "looks dead, labeled survives" rows: devices already below the
        # death threshold whose death never made it into eol_times (26 ghost
        # trajectories in train). Feeding them as negatives teaches the model
        # that textbook plunges survive - label noise on the hardest pattern.
        poison = (sub['dies42'] == 0) & (sub['v_corr_last'] < 2.48)
        sub = sub[~poison]
    return sub


def censored_pairs(table, origin_time):
    """All prior pairs with duration censored at origin_time (for AFT)."""
    prior = table[table['origin'] < origin_time].copy()
    eol = prior['eol']
    end = eol.where(eol.notna() & (eol <= origin_time), origin_time)
    prior['duration'] = ((end - prior['origin']).dt.total_seconds() / 86400).clip(lower=0.5)
    prior['event'] = (eol.notna() & (eol <= origin_time)).astype(int)
    return prior


def walk_forward(table, scenario_ctx, origins, configs):
    """configs: list of dicts(model, planner, p_star, day_q, name)."""
    # device-level EOL ledger for the cold-start floor: every death is
    # observable from the moment it happens (eol <= origin), unlike training
    # labels, which need the full 42d window closed
    dev_eol = table[['device_id', 'eol']].dropna().drop_duplicates('device_id')
    fitted = {}
    records = []
    for idx in origins:
        ctx = scenario_ctx[idx]
        origin_time = ctx['start_time']
        recent_deaths = int(((dev_eol['eol'] > origin_time - pd.Timedelta(days=WINDOW_DAYS))
                             & (dev_eol['eol'] <= origin_time)).sum())
        # unique dead devices visible to training: every death <= origin lands
        # as a positive in some closed window (the 4d synthetic grid reaches
        # back >42d), so this equals the training set's unique positives
        train_deaths = int((dev_eol['eol'] <= origin_time).sum())
        rows = table[table['origin_idx'] == idx].set_index('device_id')
        # label-free prevalence signal for the earliest origins, where even
        # recent_deaths lags (recording delay): fresh devices whose corrected
        # voltage crossed the 2.45 death threshold within the last 42d
        zone_young = int(((rows['days_below_245'] >= 0)
                          & (rows['days_below_245'] <= WINDOW_DAYS)
                          & (rows['staleness'] <= 7)).sum()) \
            if 'days_below_245' in rows.columns else 0

        preds = {}
        for model_name in {c['model'] for c in configs}:
            if (model_name, idx) not in fitted:
                factory = resolve_model(model_name)
                if model_name in ('aft', 'dthaz', 'dthaz_aux', 'logbag_open', 'blend_dh', 'blend_dhaux'):
                    train = censored_pairs(table, origin_time)
                else:
                    train = training_pairs(table, origin_time)
                fitted[(model_name, idx)] = factory().fit(train)
            preds[model_name] = fitted[(model_name, idx)].predict(rows)

        for cfg in configs:
            pred = preds[cfg['model']]
            if cfg['planner'] == 'topsum':
                PLANNERS['topsum'].alpha = cfg['p_star']  # p_star carries alpha
            if cfg['planner'] == 'ecroute':
                PLANNERS['ecroute'].margin = cfg['p_star'] or 1.0  # p_star carries margin
                PLANNERS['ecroute'].blend = cfg.get('blend', 0.0)
            if cfg['planner'] == 'tspig':
                PLANNERS['tspig'].alpha = cfg['p_star'] or 0.8
            if cfg['planner'] == 'tspig2':
                PLANNERS['tspig2'].alpha = cfg['p_star'] or 0.8
                PLANNERS['tspig2'].alpha2 = cfg.get('alpha2', 1.3)
            if cfg['planner'] == 'tsroute':
                PLANNERS['tsroute'].alpha = cfg['p_star'] or 0.8   # p_star carries alpha
                PLANNERS['tsroute'].p_be = cfg.get('p_be', 0.0)
                PLANNERS['tsroute'].floor_day = cfg.get('floor_day', 0)
                PLANNERS['tsroute'].p_conf = cfg.get('p_conf', 0.5)
                PLANNERS['tsroute'].ec_margin = cfg.get('ec_margin', 0.0)
                PLANNERS['tsroute'].ec_pmax = cfg.get('ec_pmax', 25.0)
                PLANNERS['tsroute'].split_days = cfg.get('split_days', 0)
                PLANNERS['tsroute'].floor_scale = cfg.get('floor', 0.0)
                # floor_gate: apply the floor only while the training set has
                # too few observed deaths to calibrate sum(p42) (the cold-start
                # failure); 0 = ungated. floor_zone: strengthen the prevalence
                # signal with the label-free death-zone count.
                gate = cfg.get('floor_gate', 0)
                prev = max(recent_deaths, zone_young) if cfg.get('floor_zone') \
                    else recent_deaths
                PLANNERS['tsroute'].recent_deaths = (
                    float(prev) if (gate == 0 or train_deaths < gate) else 0.0)
            global PARK_MOBILITY
            PARK_MOBILITY = cfg.get('park_mobility', 0.0)
            global EMERGENCY_DAY
            EMERGENCY_DAY = cfg.get('emergency_day', WINDOW_DAYS + 3)
            optimize_days.restarts = int(cfg.get('restarts', 2))
            optimize_days.iters = int(cfg.get('iters', 60))
            optimize_days.constructive = bool(cfg.get('constructive', False))
            optimize_days.nn_order = bool(cfg.get('nn', False))
            included = pred.index[pred['p42'] >= cfg['p_star']]
            kwargs = {}
            if cfg['planner'] in ('ecroute', 'tsroute', 'tspig', 'tspig2', 'tscut'):
                kwargs = dict(travel_costs=ctx['travel_costs'],
                              settings=ctx['settings'])
            plan = PLANNERS[cfg['planner']](pred, included, ctx['locs'],
                                            origin_time, cfg['day_q'], **kwargs)
            _, _, overall = evaluate_plan(
                plan, ctx['locs'], ctx['travel_costs'], ctx['settings'],
                eol_times=ctx['eol'], start_time=origin_time, verbose=0)
            rs = rank_stats(pred, plan, ctx['eol'], origin_time)
            records.append(dict(overall, **rs, origin=idx, **cfg))

        auc = {m: roc_auc_score(rows['dies42'], p['p42'])
               for m, p in preds.items() if rows['dies42'].sum() > 0}
        print(f"s_{idx}: " + ' '.join(f'{m}={a:.3f}' for m, a in auc.items()),
              file=sys.stderr)
    return pd.DataFrame(records)


RANK_NS = (12, 20)


def rank_stats(pred, plan, eol, start_time, order_col=None):
    """Ranking precision at the operating point, per origin.

    Measurement law #1 from the research loop: mean cost and AUC both hide
    whether a change moved the ORDER or just the depth of the cut. A model
    can hold AUC 0.947 and still have recall 0.000 on the subgroup that
    actually drives late cost. So every protocol here reports, alongside the
    cost:

      rec12/rec20  ordering only - share of the window's deaths inside the
                   top 12/20 of the score, independent of how many the
                   planner chose to swap;
      n_sched      what the planner actually bought (in-window swaps);
      caught/prec  the operating point: deaths among those swaps, and the
                   share of swaps that were real deaths;
      miss_rank    median rank of the deaths that were missed - the number
                   that says whether a miss was near the cut or invisible.
    """
    order = pred[order_col] if order_col and order_col in pred.columns else (
        pred['rank_score'] if 'rank_score' in pred.columns else pred['p42'])
    order = order.sort_values(ascending=False)
    window_end = start_time + pd.Timedelta(days=WINDOW_DAYS)
    e = eol.reindex(order.index)
    deaths = set(e.index[(e > start_time) & (e <= window_end)])
    out = {'n_deaths': len(deaths)}
    if not deaths:
        return out
    ranks = {b: r for r, b in enumerate(order.index, start=1)}
    dr = sorted(ranks[b] for b in deaths if b in ranks)
    for n in RANK_NS:
        out[f'rec{n}'] = sum(1 for r in dr if r <= n) / len(dr)
    missed = [r for r in dr if r > max(RANK_NS)]
    out['miss_rank'] = float(np.median(missed)) if missed else np.nan
    sched = set(plan.loc[plan['day'] <= window_end, 'battery'])
    out['n_sched'] = len(sched)
    out['caught'] = len(sched & deaths)
    out['prec'] = out['caught'] / len(sched) if sched else np.nan
    out['recall'] = out['caught'] / len(deaths)
    return out

def load_context(dataset_path):
    locations, timeseries, eol_times, scenarios = load_dataset(dataset_path / 'train')
    ctx = {}
    gen = iterate_scenarios(locations, timeseries, eol_times, scenarios)
    for idx, (scenario, locs, cut, eol) in enumerate(gen):
        ctx[idx] = dict(start_time=pd.Timestamp(scenario['start_time']),
                        locs=locs, eol=eol,
                        travel_costs=scenario['travel_costs'],
                        settings=scenario['settings'])
    return ctx


def summarize(df, label):
    df['name'] = (df['model'] + '/' + df['planner'] + '/p' +
                  df['p_star'].astype(str) + '/q' + df['day_q'].astype(str))
    for key, tag in [('blend', 'b'), ('p_be', 'be'), ('floor_day', 'f'), ('p_conf', 'c'),
                     ('alpha2', 'a2'), ('emergency_day', 'e'), ('restarts', 'r'),
                     ('iters', 'i'), ('ec_margin', 'ec'), ('ec_pmax', 'pm'),
                     ('split_days', 'sd'), ('park_mobility', 'pk'), ('floor', 'fl'),
                     ('floor_gate', 'fg'), ('floor_zone', 'fz'), ('constructive', 'cn'),
                     ('nn', 'nn')]:
        if key in df.columns and df[key].nunique(dropna=False) > 1:
            df['name'] += '/' + tag + df[key].fillna(0).astype(str)
    mean = df.groupby('name')['total_cost'].agg(['mean', 'std']).sort_values('mean')
    print(f'\n=== {label} ===')
    print(mean.round(1).to_string())
    # component decomposition by default (phase 19): late/early/ops drove
    # every real insight since the leaderboard comparison - stop reburying it
    comps = [c for c in ('late_swap', 'early_swap') if c in df.columns]
    ops = [c for c in ('travel', 'overtime', 'daily_limit', 'weekly_limit',
                       'battery_swap', 'building_change', 'room_change')
           if c in df.columns]
    if comps:
        dec = df.groupby('name')[comps].mean()
        if ops:
            dec['ops'] = df.groupby('name')[ops].mean().sum(axis=1)
        print(dec.round(1).reindex(mean.index).to_string())
    # ranking precision at the operating N (measurement law #1): cost alone
    # cannot say whether a change reordered the list or just cut it deeper
    rk = [c for c in ('rec12', 'rec20', 'recall', 'prec', 'n_sched',
                      'miss_rank', 'n_deaths') if c in df.columns]
    if rk:
        g = df.groupby('name')
        r = g[rk].mean()
        # pool over origins rather than averaging per-origin ratios: a
        # 1-death origin must not outvote a 12-death one
        if {'caught', 'n_deaths', 'n_sched'} <= set(df.columns):
            r['recall'] = g['caught'].sum() / g['n_deaths'].sum()
            r['prec'] = g['caught'].sum() / g['n_sched'].sum()
        print(r.round(3).reindex(mean.index).to_string())
    return mean


ALL_ORIGINS = range(8, 42)  # merged dev+test block (protocol 2026-08-18)


def selection_stability(df, n_random=2000, n_blocks=6, seed=0):
    """
    Repeated-split champion selection over the merged origin block.

    Training is walk-forward and time-locked, so only the *selection* step is
    resampled: choose the best config on half the origins, measure it on the
    other half. Random iid splits estimate selection frequency; blocked splits
    (contiguous chunks) respect the correlation between origins whose 42-day
    windows overlap. Acceptance rule (EXPERIMENTS.md): a challenger should be
    selected in >= 60% of random splits AND be modal under blocked splits.
    """
    from itertools import combinations

    df = df.copy()
    df['name'] = (df['model'] + '/' + df['planner'] + '/p' +
                  df['p_star'].astype(str))
    # distinguish configs that differ only in planner knobs (same tags as summarize)
    for key, tag in [('blend', 'b'), ('p_be', 'be'), ('floor_day', 'f'), ('p_conf', 'c'),
                     ('alpha2', 'a2'), ('emergency_day', 'e'), ('restarts', 'r'),
                     ('iters', 'i'), ('ec_margin', 'ec'), ('ec_pmax', 'pm'),
                     ('split_days', 'sd'), ('park_mobility', 'pk'), ('floor', 'fl'),
                     ('floor_gate', 'fg'), ('floor_zone', 'fz'), ('constructive', 'cn'),
                     ('nn', 'nn')]:
        if key in df.columns and df[key].nunique(dropna=False) > 1:
            df['name'] += '/' + tag + df[key].fillna(0).astype(str)
    pv = df.pivot_table(index='origin', columns='name', values='total_cost')
    pv = pv.sort_index()
    names = pv.columns.tolist()
    X = pv.values
    n = len(pv)
    rng = np.random.RandomState(seed)

    def run(splits, label):
        # decided-splits-only (phase 19): a challenger byte-identical to the
        # incumbent on most origins used to lose every all-tie split to
        # alphabetical argmin order - selection% was mechanically capped
        # (phase 11 had to argue around this by hand). Ties are now excluded
        # from the denominator and reported separately.
        sel = {m: 0 for m in names}
        best = {m: 0 for m in names}
        regrets, tied = [], 0
        for A in splits:
            B = np.setdiff1d(np.arange(n), A)
            ma, mb = X[A].mean(0), X[B].mean(0)
            if ma.max() - ma.min() < 1e-9:
                tied += 1
                continue
            i = ma.argmin()
            j = mb.argmin()
            sel[names[i]] += 1
            best[names[j]] += 1
            regrets.append(mb[i] - mb[j])
        k = max(1, len(splits) - tied)
        print(f'\n--- selection stability: {label} ({len(splits)} splits, '
              f'{tied} all-tie excluded) ---')
        print(f"{'config':40s} {'selected%':>9s} {'best-holdout%':>13s} {'mean':>8s}")
        for m in sorted(names, key=lambda m: -sel[m]):
            print(f'{m:40s} {100*sel[m]/k:9.1f} {100*best[m]/k:13.1f} '
                  f'{X[:, names.index(m)].mean():8.1f}')
        if regrets:
            print(f'mean selection regret {np.mean(regrets):.1f} '
                  f'(90th pct {np.percentile(regrets, 90):.1f})')

    half = n // 2
    run([rng.choice(n, half, replace=False) for _ in range(n_random)],
        f'random {half}/{n - half}')
    blocks = np.array_split(np.arange(n), n_blocks)
    run([np.concatenate([blocks[i] for i in c])
         for c in combinations(range(n_blocks), n_blocks // 2)],
        f'blocked ({n_blocks} chunks choose {n_blocks // 2})')


def main():
    dataset_path = Path(os.environ.get('BATTERYSWAP_DATASET_PATH', 'dataset'))
    table = load_features(dataset_path)
    # derived: voltage projected through the reporting gap by the device's own
    # slope, and leak-free within-origin fleet ranks (data <= origin only)
    table['v_proj'] = table['v_last'] + table['slope_lastk'] * table['staleness']

    table['r_vlast'] = table.groupby('origin_idx')['v_last'].rank(pct=True)
    table['r_tau'] = table.groupby('origin_idx')['tau_projected'].rank(pct=True)
    try:
        from batteryswap_example.aux_features import load as load_aux
        aux = load_aux(dataset_path)
        dup = [c for c in aux.columns if c in table.columns
               and c not in ('device_id', 'origin_idx')]
        table = table.merge(aux.drop(columns=dup),
                            on=['device_id', 'origin_idx'], how='left')
    except FileNotFoundError:
        pass
    try:
        from batteryswap_example.fht_features import load as load_fht
        table = table.merge(load_fht(dataset_path), on=['device_id', 'origin_idx'],
                            how='left')
        table['fht_days'] = ((table['v_corr_last'] - 2.40)
                             / table['wnr_m'].clip(lower=1e-4)
                             - table['staleness']).clip(-300, 600)
    except FileNotFoundError:
        pass
    ctx = load_context(dataset_path)

    phase = sys.argv[1] if len(sys.argv) > 1 else 'dev'

    if phase == 'dev':
        configs = []
        for model in MODELS:
            for planner in PLANNERS:
                for p_star in (0.1, 0.2, 0.35):
                    for day_q in (0.05, 0.20):
                        configs.append(dict(model=model, planner=planner,
                                            p_star=p_star, day_q=day_q))
        df = walk_forward(table, ctx, DEV_ORIGINS, configs)
        df.to_csv('/tmp/dev_results.csv', index=False)
        summarize(df, f'DEV s_8..s_23 ({len(configs)} configs)')
    elif phase == 'panel':
        # merged-block evaluation (protocol of 2026-08-18): score configs on all
        # 34 usable origins, then repeated-split selection stability
        import json
        configs = json.loads(Path('/tmp/panel_configs.json').read_text())
        df = walk_forward(table, ctx, ALL_ORIGINS, configs)
        df.to_csv('/tmp/panel_results.csv', index=False)
        summarize(df, f'PANEL s_{ALL_ORIGINS[0]}..s_{ALL_ORIGINS[-1]}')
        selection_stability(df)
    else:
        # phase 'dev2' or 'test': configs from /tmp/<phase>_configs.json
        import json
        configs = json.loads(Path(f'/tmp/{phase}_configs.json').read_text())
        origins = DEV_ORIGINS if phase.startswith('dev') else TEST_ORIGINS
        df = walk_forward(table, ctx, origins, configs)
        df.to_csv(f'/tmp/{phase}_results.csv', index=False)
        summarize(df, f'{phase.upper()} origins {list(origins)[0]}..{list(origins)[-1]}')


if __name__ == '__main__':
    main()