"""
Precompute the (device, origin) feature/label table used by every model.

One row per (alive device, scenario origin). Features use only data at or before
the origin; no fleet-wide statistics are computed here, so slicing at an origin
cannot leak the future. Labels use the full eol_times and are for
training/evaluation only - the planner never reads them.

Cached to dataset/train/features.parquet (rebuild with --force).
"""

from pathlib import Path
import os
import sys

import numpy as np
import pandas as pd

from batteryswap_public.utils import load_dataset

MIN_TEMP = 10
MAX_TEMP = 30
WINDOW_DAYS = 42


def daily_voltage(timeseries: pd.DataFrame) -> pd.DataFrame:
    """In-band daily median voltage per device: (device_id, day, v)."""
    df = timeseries.reset_index()
    # device_id is categorical in the parquet; observed=False groupbys would emit
    # the full device x day cartesian product with NaN voltages
    df['device_id'] = df['device_id'].astype(str)
    df = df[(df['temperature'] > MIN_TEMP) & (df['temperature'] < MAX_TEMP)]
    df['day'] = df['end_time'].dt.normalize()
    out = df.groupby(['device_id', 'day'])['voltage'].median().rename('v').reset_index()
    return out.dropna(subset=['v']).sort_values('day')


def slope(days: np.ndarray, values: np.ndarray) -> float:
    """Per-day slope; NaN when under 2 points."""
    if len(values) < 2:
        return np.nan
    x = (days - days[-1]) / np.timedelta64(1, 'D')
    return float(np.polyfit(x.astype(float), values, 1)[0])


def detect_swaps(dall: pd.DataFrame, beta: float) -> pd.Series:
    """
    Last UNDOCUMENTED battery-swap day per device, from daily corrected
    voltage <= whatever the caller masked to (leak-free if dall is masked to
    the origin). A swap = consecutive daily readings jumping up by >= 0.15 V
    from a degraded level (<= 2.85) onto the fresh plateau (>= 3.00). The
    landing gate is load-bearing: the detected set is identical for jump
    thresholds 0.10-0.25 (real swaps jump >= 0.28 V, nearest noise landing
    2.989), so the rule sits in a wide dead zone. Devices without a detected
    swap are absent from the result.
    """
    d = dall.sort_values(['device_id', 'day'])
    vc = d['v'] - beta * (d['t'] - 20.0)
    same = d['device_id'].eq(d['device_id'].shift())
    jump = (vc - vc.shift() >= 0.15) & (vc.shift() <= 2.85) & (vc >= 3.00) & same
    return d.loc[jump].groupby('device_id')['day'].max()


def features_at_origin(daily: pd.DataFrame, devices: pd.DataFrame,
                       origin: pd.Timestamp, swaps: pd.Series = None) -> pd.DataFrame:
    """Feature rows for every device with any reading at or before origin.

    swaps: optional device -> last-swap-day map (detect_swaps). History before
    a device's swap belongs to the PREVIOUS battery: it is dropped so slopes,
    baselines and windows describe the battery actually in the slot, and
    age_days restarts at the swap.
    """
    h = daily[daily['day'] <= origin]
    if swaps is not None and len(swaps):
        cut = h['device_id'].map(swaps)
        h = h[cut.isna() | (h['day'] >= cut)]

    g = h.groupby('device_id')
    last = g.tail(1).set_index('device_id')
    feats = pd.DataFrame({'v_last': last['v']})
    feats['staleness'] = (origin - last['day']).dt.days

    # per-device historical reference: median over first 90 observed days
    first_day = g['day'].min()
    early = h[h['day'] <= h['device_id'].map(first_day) + pd.Timedelta(days=90)]
    feats['v_ref'] = early.groupby('device_id')['v'].median()
    feats['v_drop'] = feats['v_ref'] - feats['v_last']

    for w in (30, 90):
        hw = h[h['day'] > origin - pd.Timedelta(days=w)]
        gw = hw.groupby('device_id')
        feats[f'slope_{w}'] = gw.apply(
            lambda u: slope(u['day'].values, u['v'].values), include_groups=False)
        feats[f'v_min_{w}'] = gw['v'].min()
        feats[f'n_{w}'] = gw['v'].size()

    # window-free slope over the last K observations: sparse traces (the worst
    # recall misses) have <2 points in any fixed window but still have history
    lastk = g.tail(10)
    feats['slope_lastk'] = lastk.groupby('device_id').apply(
        lambda u: slope(u['day'].values, u['v'].values), include_groups=False)

    # recent drop: median of days -28..-7 minus the latest value (trend, not level)
    prev = h[(h['day'] > origin - pd.Timedelta(days=28)) &
             (h['day'] <= origin - pd.Timedelta(days=7))]
    feats['drop_recent'] = prev.groupby('device_id')['v'].median() - feats['v_last']

    # fraction of the last 30 observations below the personal baseline - 0.15
    last30 = g.tail(30).set_index('device_id')
    below = last30['v'] < (feats['v_ref'].reindex(last30.index) - 0.15)
    feats['below_frac'] = below.groupby(level=0).mean()

    feats['accel'] = feats['slope_30'] - feats['slope_90']

    # ---- fleet-relative encoding (REJECTED - phase 41) ----------------------
    # Ported from the blind-harness champion, where vrel and vmin7_rank were
    # the two largest single-feature wins. Both are computed at PLAN time over
    # the fleet actually in front of the planner, so they absorb a fleet-level
    # voltage shift that the absolute features (v_last, v_min_30, ...) cannot.
    # The motivation was TRANSFER to a held-out fleet.
    #
    # NOT IN models_cv.FEATURES, and they must not be added. Re-ablated on
    # recall at FIXED N (the phase-39 directive), which is the only way to see
    # a feature's effect on ordering separately from its effect on the topsum
    # count. Deaths caught, paired per origin, bootstrap 95% CI:
    #
    #                  native all48            static gate (5 splits)
    #   vrel+vmin7  N=12  -7 [-15,+0]        N=12  +4 [-2,+10]
    #               N=20  +5 [ -1,+11]       N=20  -1 [-9, +7]
    #
    # The sign flips with N on BOTH instruments and every CI straddles zero.
    # The ordering is untouched at 36-41 of 48 native origins and 228 of 240
    # static ones - these features barely enter the ranking at all. The one CI
    # that excludes zero is vrel alone, native, N=12: -6 [-12,-1], negative.
    # Cost confirms it: all48 1800.4 vs 1661.7 (+138.7).
    #
    # Left computed rather than deleted only to avoid regenerating the
    # packaged parquet on submission day; the two columns cost ~1s inside
    # build_split_table's 360s budget. Delete them when there is slack.
    #
    # Deliberately NOT filled here. The classifier pipeline median-imputes and
    # then scales, which parks a stale/dark device at "unremarkable" so it is
    # never selected; three independent measurements found that making those
    # devices visible instead is a monotone loss. Do not add fillna() or
    # na_option= to the rank.
    #
    # Population note: the denominator is whatever has history at or before the
    # origin, dead-but-not-yet-dropped devices included. That is identical in
    # the training-table build and in the runtime path, which is what matters -
    # both sides see the same population definition.
    h7 = h[h['day'] > origin - pd.Timedelta(days=7)]
    feats['v_min_7'] = h7.groupby('device_id')['v'].min()
    feats['vrel'] = feats['v_last'] - feats['v_last'].median()
    feats['vmin7_rank'] = feats['v_min_7'].rank(pct=True)

    install = devices.set_index('battery')['start_time'].dt.tz_localize(None)
    eff_install = install.reindex(feats.index)
    if swaps is not None and len(swaps):
        sw = swaps.reindex(feats.index)
        eff_install = eff_install.where(sw.isna() | (eff_install >= sw), sw)
    feats['age_days'] = (origin - eff_install).dt.days

    return feats


# ------------------------------------------------------------- physics model
# Empirically validated on the train split: temperature moves voltage by
# ~+5 mV/degC within a device-week, and the 82 observed deaths all cross
# ~2.40 V, approaching it along a shared discharge curve (sd 0.034 V at EOL).

TAU_MAX = 300      # curve support, days before EOL
FIT_WINDOW = 120   # days of history used for trajectory matching
BETA_DEFAULT = 0.005


def daily_all(timeseries: pd.DataFrame) -> pd.DataFrame:
    """Daily median voltage AND temperature per device, all readings kept."""
    df = timeseries.reset_index()
    df['device_id'] = df['device_id'].astype(str)
    df['day'] = df['end_time'].dt.normalize()
    out = df.groupby(['device_id', 'day']).agg(
        v=('voltage', 'median'), t=('temperature', 'median')).reset_index()
    return out.dropna(subset=['v']).sort_values('day')


def weekly_betas(daily: pd.DataFrame) -> pd.DataFrame:
    """Per (device, week) voltage-vs-temperature slope; 'week_end' for cutoffs."""
    d = daily.copy()
    d['week'] = d['day'].dt.to_period('W')
    def beta_of(g):
        if len(g) < 5 or g['t'].std() < 1:
            return np.nan
        return np.polyfit(g['t'], g['v'], 1)[0]
    b = d.groupby(['device_id', 'week']).apply(beta_of, include_groups=False).dropna()
    b = b.rename('beta').reset_index()
    b['week_end'] = b['week'].dt.end_time.dt.normalize()
    return b


def master_curve(daily: pd.DataFrame, dead_eol: pd.Series, beta: float,
                 min_dead: int = 5) -> 'np.ndarray | None':
    """Median temp-corrected voltage at each days-before-EOL in 0..TAU_MAX."""
    if len(dead_eol) < min_dead:
        return None
    d = daily[daily['device_id'].isin(dead_eol.index)].copy()
    d['v_corr'] = d['v'] - beta * (d['t'] - 20.0)
    d['tau'] = (d['device_id'].map(dead_eol).dt.normalize() - d['day']).dt.days
    d = d[(d['tau'] >= 0) & (d['tau'] <= TAU_MAX)]
    med = d.groupby('tau')['v_corr'].median()
    grid = np.arange(0, TAU_MAX + 1)
    curve = np.interp(grid, med.index.values, med.values)
    return np.maximum.accumulate(curve)  # enforce monotone non-decreasing in tau


def physics_features(daily: pd.DataFrame, origin: pd.Timestamp,
                     beta: float, curve, match: bool = True,
                     swaps: pd.Series = None) -> pd.DataFrame:
    """
    Physics-model features.
      v_corr_last   - latest temp-corrected voltage (full history, no window)
      thresh_dist   - v_corr_last minus the curve's EOL voltage
      tau_projected - C^-1(v_corr_last) minus staleness: days-to-EOL projected
                      forward through the silence. The key feature for devices
                      that go dark and keep discharging - a stale healthy-ish
                      reading minus months of silence can mean "due now".
      tau_hat/fit_err - trajectory registration (match=True only; expensive)
    """
    all_h = daily[daily['day'] <= origin]
    if swaps is not None and len(swaps):
        cut = all_h['device_id'].map(swaps)
        all_h = all_h[cut.isna() | (all_h['day'] >= cut)]
    last = all_h.sort_values('day').groupby('device_id').tail(1).set_index('device_id')
    v_corr = last['v'] - beta * (last['t'] - 20.0)
    stale = (origin - last['day']).dt.days

    out = pd.DataFrame({'v_corr_last': v_corr})

    # ghost marker: days since the corrected voltage first crossed below 2.45
    # (post-swap history only). A real recorded death follows the crossing
    # within days; a crossing >42d old with no EOL on file means the recording
    # process abandoned this device - it will never be scored as a death.
    vc_all = all_h['v'] - beta * (all_h['t'] - 20.0)
    below = all_h.loc[vc_all < 2.45]
    first_below = below.groupby('device_id')['day'].min()
    out['days_below_245'] = (origin - first_below.reindex(out.index)).dt.days

    # plunge precursors: volatility of the corrected voltage rises BEFORE the
    # level falls (AUC 0.83/0.81 on the voltage-silent dying cohort, where the
    # level itself is blind). These are the fast-plunge early-warning features.
    h30 = all_h[all_h['day'] > origin - pd.Timedelta(days=30)].copy()
    h30['v_corr'] = h30['v'] - beta * (h30['t'] - 20.0)
    g30 = h30.groupby('device_id')['v_corr']
    out['vstd_30'] = g30.std()
    out['rng_30'] = g30.max() - g30.min()
    if curve is not None:
        grid = np.arange(0, TAU_MAX + 1)
        tau_at_last = np.interp(v_corr.values, curve, grid)
        out['thresh_dist'] = v_corr - curve[0]
        out['tau_projected'] = tau_at_last - stale.values
    else:
        out['thresh_dist'] = np.nan
        out['tau_projected'] = np.nan
    out['tau_hat'] = np.nan
    out['fit_err'] = np.nan

    if match and curve is not None:
        h = all_h[all_h['day'] > origin - pd.Timedelta(days=FIT_WINDOW)].copy()
        h['v_corr'] = h['v'] - beta * (h['t'] - 20.0)
        h['delta'] = (origin - h['day']).dt.days
        taus = np.arange(1, TAU_MAX - FIT_WINDOW)
        for dev, g in h.groupby('device_id'):
            if len(g) >= 5:
                idx = taus[:, None] + g['delta'].values[None, :]
                err = np.abs(g['v_corr'].values[None, :] - curve[idx]).mean(axis=1)
                best = err.argmin()
                out.at[dev, 'tau_hat'] = float(taus[best])
                out.at[dev, 'fit_err'] = float(err[best])
    return out


def cohort_features(devices: pd.DataFrame, eol_times: pd.Series,
                    origin: pd.Timestamp) -> pd.DataFrame:
    """
    Install-cohort calibrated clock. Batteries from the same install batch die at
    nearly the same age (batch explains ~90% of lifetime variance), so
    age - median(batch lifetimes observed so far) predicts death even when the
    voltage has not moved yet (AUC 0.91 on voltage-silent pairs).
    Uses only deaths <= origin.
    """
    install = devices.set_index('battery')['start_time'].dt.tz_localize(None)
    dead = eol_times[eol_times.notna() & (eol_times <= origin)]
    life = (dead - install.reindex(dead.index)).dt.days.rename('life')

    out = pd.DataFrame(index=install.index)
    age = (origin - install).dt.days
    for col, key in [('cohort_rel_age', install),
                     ('cohort_rel_age_m', install.dt.to_period('M'))]:
        batch_life = life.groupby(key.reindex(life.index)).agg(['median', 'size'])
        med = key.map(batch_life['median'].where(batch_life['size'] >= 2))
        out[col] = age - med
        # neutral-when-unknown encoding: fires only past the batch's observed
        # lifetime, exactly 0 otherwise - no imputation noise for the majority
        out[col.replace('rel_age', 'overdue')] = (age - med).clip(lower=0).fillna(0)
    out.index.name = 'device_id'
    return out


def build(dataset_path: Path, out_path: Path) -> pd.DataFrame:
    locations, timeseries, eol_times, scenarios = load_dataset(dataset_path / 'train')
    daily = daily_voltage(timeseries)
    dall = daily_all(timeseries)
    betas = weekly_betas(dall)

    # real scenario origins, plus synthetic training-only origins: a denser grid
    # (every 4 days) extended back before s_0, so every observed death lands in
    # training windows and positive pairs are ~3x more numerous. Synthetic rows are
    # never scored - they only feed training_pairs().
    real_starts = [pd.Timestamp(sc['start_time']) for sc in scenarios]
    synth = pd.date_range('2025-02-15', max(real_starts), freq='4D')
    synth = [t for t in synth if t not in set(real_starts)]
    origins = ([(idx, t, False) for idx, t in enumerate(real_starts)] +
               [(1000 + i, t, True) for i, t in enumerate(synth)])

    rows = []
    for idx, origin, is_synth in origins:
        alive = eol_times.isna() | (eol_times > origin)
        alive_ids = eol_times.index[alive]

        # NOTE (phase 15): detect_swaps truncation was tested here and NOT
        # adopted - panel 1684.5 vs 1679.2. The undocumented-swap age error
        # touches only age_days (weak, linear); the machinery stays available
        # via the swaps= parameter but production builds run without it.
        past = betas[betas['week_end'] <= origin]['beta']
        beta = float(past.median()) if len(past) >= 20 else BETA_DEFAULT

        f = features_at_origin(daily, locations, origin)
        f = f.reindex(alive_ids)  # NaN feature rows kept: planner must handle them

        # physics features: beta and curve from information <= origin only.
        # Synthetic origins get the light variant (no trajectory matching).
        dead = eol_times[eol_times.notna() & (eol_times <= origin)]
        curve = master_curve(dall[dall['day'] <= origin], dead, beta)
        f = f.join(physics_features(dall, origin, beta, curve, match=not is_synth))
        f = f.join(cohort_features(locations, eol_times, origin))
        f['synthetic'] = is_synth
        f['origin_idx'] = idx
        f['origin'] = origin

        eol = eol_times.reindex(alive_ids)
        f['eol'] = eol.values
        f['dies42'] = ((eol.notna()) & (eol <= origin + pd.Timedelta(days=WINDOW_DAYS))).astype(int).values
        f['days_to_eol'] = ((eol - origin).dt.total_seconds() / 86400).values
        f['window_closed'] = origin + pd.Timedelta(days=WINDOW_DAYS) <= pd.Timestamp('2026-07-31')

        rows.append(f.reset_index().rename(columns={'index': 'device_id'}))
        print(f'origin {idx} {origin.date()}: {len(f)} alive, '
              f'{f["dies42"].sum()} die in window', file=sys.stderr)

    table = pd.concat(rows, ignore_index=True)
    table.to_parquet(out_path)
    return table


def load(dataset_path: Path = Path('dataset'), force: bool = False) -> pd.DataFrame:
    out_path = dataset_path / 'train' / 'features.parquet'
    if out_path.exists() and not force:
        return pd.read_parquet(out_path)
    return build(dataset_path, out_path)


if __name__ == '__main__':
    force = '--force' in sys.argv
    dataset_path = Path(os.environ.get('BATTERYSWAP_DATASET_PATH', 'dataset'))
    t = load(dataset_path, force=force)
    print(t.groupby('origin_idx')['dies42'].agg(['size', 'sum']).T)
    print(t.describe().round(3).to_string())