"""Scaling evaluation: run an arm on a synthetic fleet at scale K.

Deployment-faithful: ONE static fit on the real packaged train table
(submit.py's held-out branch), then the 48 real scenario times planned on
the synthetic fleet and costed by the real evaluator with derived EOLs.

Arms: champ (plain logbag, alpha .8) | ft2 (frontier, alpha .8)
      | ft2a10 (frontier, alpha 1.0)
Usage: scale_run.py K ARM
"""
import sys, time, pickle
from pathlib import Path
import numpy as np
import pandas as pd

OURS = Path('/Users/THIF/repos/BatterySwapAI2026-Example')
sys.path.insert(0, str(OURS))
from batteryswap_public.utils import load_dataset, iterate_scenarios
from batteryswap_public.evaluate import evaluate_plan
from batteryswap_example.features import (
    daily_voltage, daily_all, weekly_betas, master_curve, physics_features,
    features_at_origin, WINDOW_DAYS, BETA_DEFAULT)
from batteryswap_example.fht_features import daily_corrected, wiener_at_origin
import batteryswap_example.models_cv as mc
import batteryswap_example.submit as sub

K = int(sys.argv[1])
ARM = sys.argv[2]
ALPHA = 1.0 if ARM.endswith('a10') else 0.8
if ARM == 'ft2fix':
    ALPHA = 0.8

with open(f'/tmp/c5/fleet_k{K}.pkl', 'rb') as f:
    FL = pickle.load(f)
locations = FL['locations']
for _c in ('start_time', 'end_time'):
    if getattr(locations[_c].dt, 'tz', None) is not None:
        locations[_c] = locations[_c].dt.tz_localize(None)
eol = FL['eol']
if getattr(eol.dt, 'tz', None) is not None:
    eol = eol.dt.tz_localize(None)
travel = FL['travel']
ts = FL['ts'].set_index(['device_id', 'end_time']).sort_index()

_, _, eol_real, scenarios = load_dataset(OURS / 'dataset/train')
# fleet-wide EOL series shaped like load_dataset's (NaT for survivors);
# clone-0 devices take the OFFICIAL values (exact), shifted clones the
# telemetry-derived ones (0 mismatches on the checkable set). The 6 real
# deaths the derivation misses become permanent low-voltage "ghosts" in
# shifted clones - noted, and shared by every arm equally.
eol_syn = eol.reindex(locations['battery'])
eol_syn.index = locations['battery']
er = eol_real.copy()
if getattr(er.dt, 'tz', None) is not None:
    er = er.dt.tz_localize(None)
c0 = eol_syn.index.intersection(er.index)
eol_syn.loc[c0] = er.reindex(c0)
# scenarios with the synthetic travel matrix
scenarios = [dict(sc, travel_costs=travel) for sc in scenarios]
GEN = list(iterate_scenarios(locations, ts, eol_syn, scenarios))

# ---- one static foreign fit (the deployment path) ----
base = pd.read_parquet(OURS / 'batteryswap_example/data/train_features.parquet')
base = base[base['window_closed']]
model = (mc.SklearnRUL('logbag') if ARM == 'champ' else sub.Ft2RUL()).fit(base)
if ARM == 'ft2fix':
    # scale hardening: vmin7_rank as a FIXED-reference percentile (train
    # ECDF of v_min_7) instead of a live-fleet rank - fleet-size invariant
    _q = np.sort(base['v_min_7'].dropna().to_numpy())
    _inner = model
    class _Fix:
        def predict(self, rows):
            r = rows.copy()
            v = r['v_min_7'].to_numpy(dtype=float)
            pct = np.searchsorted(_q, v, side='right') / len(_q)
            r['vmin7_rank'] = np.where(np.isfinite(v), pct, np.nan)
            return _inner.predict(r)
    model = _Fix()

# ---- per-fleet caches ----
t0 = time.time()
daily = daily_voltage(ts)
dall = daily_all(ts)
betas = weekly_betas(dall)
fd = daily_corrected(ts)
print(f'fleet caches {time.time()-t0:.0f}s', file=sys.stderr, flush=True)

p = mc.plan_tsroute
p.alpha = ALPHA
p.ec_margin, p.ec_pmax = 0.8, 25.0
p.split_days = 7
p.floor_scale, p.recent_deaths = 2.5, 0.0

W = pd.Timedelta(days=WINDOW_DAYS)
recs = []
for i, (sc, fleet, cut, eol_f) in enumerate(GEN):
    T = pd.Timestamp(sc['start_time'])
    tt = time.time()
    feats = features_at_origin(daily, fleet, T)
    past = betas[betas['week_end'] <= T]['beta']
    beta = float(past.median()) if len(past) >= 20 else BETA_DEFAULT
    dead = eol[eol.notna() & (eol <= T)]
    curve = master_curve(dall[dall['day'] <= T], dead, beta)
    feats = feats.join(physics_features(dall, T, beta, curve, match=False))
    if ARM != 'champ':
        feats = feats.join(wiener_at_origin(
            fd, T.tz_localize(None) if T.tz is not None else T))
    feats = feats.reindex(fleet['battery'])
    t_feat = time.time() - tt

    pred = model.predict(feats)
    t_pred = time.time() - tt - t_feat
    plan = mc.plan_tsroute(pred, None, fleet, T, 0.05,
                           travel_costs=travel, settings=sc['settings'])
    t_plan = time.time() - tt - t_feat - t_pred
    _, _, overall = evaluate_plan(plan, fleet, travel, sc['settings'],
                                  eol_times=eol_f, start_time=T)
    pday = plan['day']
    if getattr(pday.dt, 'tz', None) is not None:
        pday = pday.dt.tz_localize(None)
    Tn = T.tz_localize(None) if T.tz is not None else T
    inwin = plan[(pday <= Tn + W).values]
    swap_day = inwin.groupby('battery')['day'].min()
    dies = eol[(eol > Tn) & (eol <= Tn + W)
               & eol.index.isin(set(fleet['battery']))]
    caught = int(dies.index.isin(swap_day.index).sum())
    recs.append(dict(overall, origin=i, n_swaps=int(inwin['battery'].nunique()),
                     deaths=len(dies), caught=caught,
                     t_feat=t_feat, t_pred=t_pred, t_plan=t_plan))
    print(f's_{i}: {overall["total_cost"]:8.1f} deaths {len(dies):3d} '
          f'caught {caught:3d} swaps {recs[-1]["n_swaps"]:3d} '
          f'[f{t_feat:.1f} p{t_pred:.1f} pl{t_plan:.1f}s]',
          file=sys.stderr, flush=True)

df = pd.DataFrame(recs)
df.to_csv(f'/tmp/c5/scale_k{K}_{ARM}.csv', index=False)
n_dev = locations['battery'].nunique()
print(f"\nK={K} ARM={ARM}  devices={n_dev}")
print(f"  mean total/scenario {df['total_cost'].mean():9.1f}   "
      f"per-DEVICE {df['total_cost'].mean()/n_dev:7.3f}")
print(f"  swaps/scen {df['n_swaps'].mean():6.1f}  deaths/scen "
      f"{df['deaths'].mean():5.1f}  recall {df['caught'].sum()/max(df['deaths'].sum(),1):.3f}  "
      f"swaps/death {df['n_swaps'].sum()/max(df['deaths'].sum(),1):.2f}")
print(f"  late/scen {df['late_swap'].mean():7.1f}  early/scen "
      f"{df['early_swap'].mean():7.1f}")
print(f"  sec/scenario: feat {df['t_feat'].mean():.1f} pred "
      f"{df['t_pred'].mean():.1f} plan {df['t_plan'].mean():.1f}")