Blame for src/batteryswapai/competition_features.py

from __future__ import annotations

from collections.abc import Iterable

import numpy as np
import pandas as pd


ROLLING_WINDOWS_DAYS = (3, 7, 14, 30, 60, 90, 180, 365)

CONTEXT_LOW_VOLTAGE = 2.45
TRAJECTORY_LAGS = tuple(range(0, 365, 14))
TRAJECTORY_COLUMNS = tuple(f"traj_lag{lag}" for lag in TRAJECTORY_LAGS)
TRAJECTORY_COVERAGE_COLUMN = "traj_coverage"
MIN_TRAJECTORY_COVERAGE = 0.5
TRAJECTORY_REFERENCE_OFFSET = 0
CONTEXT_COLUMNS = (
    "room_voltage_mean",
    "room_voltage_min",
    "building_voltage_mean",
    "voltage_minus_room",
    "voltage_minus_building",
    "room_low_count",
    "building_low_count",
    "scenario_low_fraction",
)

NON_FEATURE_COLUMNS = {
    "scenario",
    "scenario_time",
    "device_id",
    "battery",
    "building",
    "room",
    "start_time",
    "end_time",
    "effective_eol",
    "target_rul_days",
    "event_observed",
    # Dataset/scenario boundary fields are useful to the planner and reports,
    # but they are not battery-health measurements. Learning from them makes
    # same-scenario validation optimistic and breaks on a shifted hidden split.
    "censor_proxy_rul_days",
    "scenario_month_sin",
    "scenario_month_cos",
    # Scenario-level context, used only by the auxiliary context model.
    *CONTEXT_COLUMNS,
    # Trajectory history, used only by the degradation models.
    *TRAJECTORY_COLUMNS,
    TRAJECTORY_COVERAGE_COLUMN,
}


def _as_columns(timeseries: pd.DataFrame) -> pd.DataFrame:
    frame = timeseries
    if "device_id" not in frame.columns or "end_time" not in frame.columns:
        frame = frame.reset_index()
    required = {"device_id", "end_time", "voltage", "temperature"}
    missing = required - set(frame.columns)
    if missing:
        raise ValueError(f"Battery time series is missing columns: {sorted(missing)}")
    return frame[["device_id", "end_time", "voltage", "temperature"]].copy()


def build_daily_features(
    timeseries: pd.DataFrame,
    windows: Iterable[int] = ROLLING_WINDOWS_DAYS,
) -> pd.DataFrame:
    """Build past-only daily features for all scenarios."""

    frame = _as_columns(timeseries)
    frame["end_time"] = pd.to_datetime(frame["end_time"], errors="coerce")
    frame["voltage"] = pd.to_numeric(frame["voltage"], errors="coerce")
    frame["temperature"] = pd.to_numeric(frame["temperature"], errors="coerce")
    frame = frame.dropna(subset=["device_id", "end_time", "voltage"])
    frame["day"] = frame["end_time"].dt.floor("D")

    daily = (
        frame.groupby(["device_id", "day"], observed=True, sort=True)
        .agg(
            voltage_median=("voltage", "median"),
            voltage_mean=("voltage", "mean"),
            voltage_min=("voltage", "min"),
            voltage_max=("voltage", "max"),
            voltage_std=("voltage", "std"),
            temperature_median=("temperature", "median"),
            temperature_std=("temperature", "std"),
            observations_day=("voltage", "size"),
        )
        .reset_index()
        .sort_values(["device_id", "day"])
        .reset_index(drop=True)
    )

    stable = frame[frame["temperature"].between(10.0, 30.0, inclusive="both")]
    stable_daily = (
        stable.groupby(["device_id", "day"], observed=True, sort=True)["voltage"]
        .median()
        .rename("stable_voltage_median")
        .reset_index()
    )
    daily = daily.merge(stable_daily, on=["device_id", "day"], how="left")
    daily["stable_voltage_median"] = daily["stable_voltage_median"].fillna(daily["voltage_median"])

    grouped = daily.groupby("device_id", observed=True, sort=False)
    first_day = grouped["day"].transform("min")
    daily["history_days"] = (daily["day"] - first_day).dt.total_seconds() / 86400.0
    daily["days_since_previous"] = grouped["day"].diff().dt.total_seconds() / 86400.0

    for window in tuple(int(value) for value in windows):
        min_periods = max(2, window // 4)
        rolled = (
            grouped["stable_voltage_median"]
            .rolling(window, min_periods=min_periods)
            .agg(["mean", "min", "max", "std"])
            .reset_index(level=0, drop=True)
            .reindex(daily.index)
        )
        for statistic in ("mean", "min", "max", "std"):
            daily[f"voltage_{statistic}_{window}d"] = rolled[statistic]

        daily[f"temperature_mean_{window}d"] = (
            grouped["temperature_median"]
            .rolling(window, min_periods=min_periods)
            .mean()
            .reset_index(level=0, drop=True)
            .reindex(daily.index)
        )
        lag_voltage = grouped["stable_voltage_median"].shift(window - 1)
        lag_day = grouped["day"].shift(window - 1)
        span = (daily["day"] - lag_day).dt.total_seconds() / 86400.0
        daily[f"voltage_slope_{window}d"] = (
            daily["stable_voltage_median"] - lag_voltage
        ) / span.replace(0.0, np.nan)
        daily[f"span_{window}d"] = span

    daily["voltage_recent_vs_30d"] = daily["voltage_mean_7d"] - daily["voltage_mean_30d"]
    daily["voltage_recent_vs_90d"] = daily["voltage_mean_14d"] - daily["voltage_mean_90d"]
    daily["voltage_range_day"] = daily["voltage_max"] - daily["voltage_min"]
    return daily


def scenario_history(history: pd.DataFrame, scenario_time) -> pd.DataFrame:
    """Restrict raw observations to what the scenario can see, exactly as iterate_scenarios does."""

    start = pd.Timestamp(scenario_time)
    frame = history
    if "end_time" not in frame.columns:
        frame = frame.reset_index()
    visible = frame[pd.to_datetime(frame["end_time"]) <= start]
    return visible.set_index(["device_id", "end_time"])


def scenario_history_snapshot(
    history: pd.DataFrame,
    locations: pd.DataFrame,
    scenario_name: str,
    scenario_time,
    windows: Iterable[int] = ROLLING_WINDOWS_DAYS,
) -> pd.DataFrame:
    """Features for one scenario, built only from that scenario's visible history."""

    visible = scenario_history(history, scenario_time)
    return scenario_snapshot(
        build_daily_features(visible, windows=windows), locations, scenario_name, scenario_time
    )


def build_trajectory_matrix(daily_features: pd.DataFrame) -> dict:
    """Dense device-by-day smoothed voltage grid, built once per split and cached."""

    cached = daily_features.attrs.get("_trajectory_matrix")
    if cached is not None:
        return cached

    frame = daily_features[["device_id", "day", "stable_voltage_median"]].dropna()
    frame = frame.sort_values(["device_id", "day"])
    smooth = frame.groupby("device_id", observed=True)["stable_voltage_median"].transform(
        lambda values: values.rolling(7, min_periods=3).median()
    )
    frame = frame.assign(smooth=smooth).dropna(subset=["smooth"])
    if frame.empty:
        empty = {"index": {}, "origin": pd.Timestamp("2000-01-01"),
                 "grid": np.zeros((0, 1), dtype=np.float32),
                 "first": np.zeros(0, dtype=np.int64),
                 "observed_last": np.zeros(0, dtype=np.int64)}
        daily_features.attrs["_trajectory_matrix"] = empty
        return empty
    devices = sorted(frame["device_id"].astype(str).unique())
    index = {device: position for position, device in enumerate(devices)}
    origin = frame["day"].min().normalize()
    span = int((frame["day"].max().normalize() - origin).days) + 1
    grid = np.full((len(devices), span), np.nan, dtype=np.float32)
    rows = frame["device_id"].astype(str).map(index).to_numpy(dtype=np.int64)
    columns = (frame["day"] - origin).dt.days.to_numpy(dtype=np.int64)
    grid[rows, columns] = frame["smooth"].to_numpy()

    # Carry the last reading forward to the end of the grid. Stopping at the device's
    # final observation would make availability depend on whether it reported again later.
    first = np.full(len(devices), np.iinfo(np.int64).max, dtype=np.int64)
    observed_last = np.full(len(devices), -1, dtype=np.int64)
    for position in range(len(devices)):
        seen = np.flatnonzero(np.isfinite(grid[position]))
        if seen.size == 0:
            continue
        low = int(seen[0])
        segment = grid[position, low:]
        carry = np.maximum.accumulate(
            np.where(np.isfinite(segment), np.arange(segment.size), 0)
        )
        grid[position, low:] = segment[carry]
        first[position] = low
        observed_last[position] = int(seen[-1])

    built = {"index": index, "origin": origin, "grid": grid, "first": first,
             "observed_last": observed_last}
    daily_features.attrs["_trajectory_matrix"] = built
    return built


def trajectory_bins(matrix: dict, batteries, cutoff) -> np.ndarray:
    """Voltage at the cutoff and every 14 days before it; never reads past the cutoff."""

    grid = matrix["grid"]
    first = matrix["first"]
    rows = np.array([matrix["index"].get(str(name), -1) for name in batteries], dtype=np.int64)
    cut = int((pd.Timestamp(cutoff).normalize() - matrix["origin"]).days)
    lags = np.asarray(TRAJECTORY_LAGS, dtype=np.int64) + TRAJECTORY_REFERENCE_OFFSET
    columns = cut - lags[None, :]
    out = np.full((len(rows), lags.size), np.nan, dtype=float)
    known = rows >= 0
    if known.any():
        safe = np.clip(columns, 0, grid.shape[1] - 1)
        values = grid[rows[known][:, None], safe]
        usable = (columns >= 0) & (columns >= first[rows[known]][:, None])
        out[known] = np.where(usable, values, np.nan)
    return out


def scenario_snapshot(
    daily_features: pd.DataFrame,
    locations: pd.DataFrame,
    scenario_name: str,
    scenario_time: pd.Timestamp | str,
) -> pd.DataFrame:
    """Get the last feature row visible at the scenario time."""

    start = pd.Timestamp(scenario_time).normalize()
    alive = locations.copy()
    if "battery" not in alive.columns:
        raise ValueError("locations must contain a battery column")

    available = daily_features[daily_features["day"] <= start]
    latest = available.groupby("device_id", observed=True, sort=False).tail(1)
    snapshot = alive.merge(latest, left_on="battery", right_on="device_id", how="left")
    snapshot["scenario"] = scenario_name
    snapshot["scenario_time"] = start

    snapshot["location_age_days"] = (
        start - pd.to_datetime(snapshot["start_time"], errors="coerce").dt.normalize()
    ).dt.total_seconds() / 86400.0
    snapshot["data_gap_days"] = (
        start - pd.to_datetime(snapshot["day"], errors="coerce")
    ).dt.total_seconds() / 86400.0
    snapshot["censor_proxy_rul_days"] = (
        pd.to_datetime(snapshot["end_time"], errors="coerce").dt.normalize()
        + pd.Timedelta(days=30)
        - start
    ).dt.total_seconds() / 86400.0
    snapshot["scenario_month_sin"] = np.sin(2.0 * np.pi * start.month / 12.0)
    snapshot["scenario_month_cos"] = np.cos(2.0 * np.pi * start.month / 12.0)

    snapshot["voltage_rank_global"] = snapshot["stable_voltage_median"].rank(pct=True)
    snapshot["voltage_rank_building"] = snapshot.groupby("building", observed=True)[
        "stable_voltage_median"
    ].rank(pct=True)
    snapshot["voltage_rank_room"] = snapshot.groupby("room", observed=True)[
        "stable_voltage_median"
    ].rank(pct=True)
    snapshot["building_battery_count"] = snapshot.groupby("building", observed=True)[
        "battery"
    ].transform("size")
    snapshot["room_battery_count"] = snapshot.groupby("room", observed=True)["battery"].transform(
        "size"
    )
    snapshot = add_context_columns(snapshot)
    bins = trajectory_bins(
        build_trajectory_matrix(daily_features), snapshot["battery"].astype(str), start
    )
    for position, column in enumerate(TRAJECTORY_COLUMNS):
        snapshot[column] = bins[:, position]
    snapshot[TRAJECTORY_COVERAGE_COLUMN] = np.isfinite(bins).mean(axis=1)
    return snapshot


def add_context_columns(snapshot: pd.DataFrame) -> pd.DataFrame:
    """Room and building degradation context for one scenario."""

    voltage = pd.to_numeric(snapshot["stable_voltage_median"], errors="coerce")
    room = snapshot.groupby("room", observed=True)["stable_voltage_median"]
    building = snapshot.groupby("building", observed=True)["stable_voltage_median"]
    snapshot["room_voltage_mean"] = room.transform("mean")
    snapshot["room_voltage_min"] = room.transform("min")
    snapshot["building_voltage_mean"] = building.transform("mean")
    snapshot["voltage_minus_room"] = voltage - snapshot["room_voltage_mean"]
    snapshot["voltage_minus_building"] = voltage - snapshot["building_voltage_mean"]
    low = (voltage < CONTEXT_LOW_VOLTAGE).astype(float)
    snapshot["room_low_count"] = low.groupby(snapshot["room"]).transform("sum")
    snapshot["building_low_count"] = low.groupby(snapshot["building"]).transform("sum")
    snapshot["scenario_low_fraction"] = float(low.mean())
    return snapshot


def attach_training_targets(
    snapshot: pd.DataFrame,
    eol_times: pd.Series,
    unobserved_eol_days: float = 30.0,
) -> pd.DataFrame:
    """Attach public training targets."""

    result = snapshot.copy()
    observed = pd.to_datetime(result["battery"].map(eol_times), errors="coerce")
    assumed = (
        pd.to_datetime(result["end_time"], errors="coerce").dt.normalize()
        + pd.to_timedelta(unobserved_eol_days, unit="D")
    )
    result["event_observed"] = observed.notna().astype("int8")
    result["effective_eol"] = observed.fillna(assumed)
    result["target_rul_days"] = (
        result["effective_eol"] - pd.to_datetime(result["scenario_time"])
    ).dt.total_seconds() / 86400.0
    return result


def numeric_feature_columns(snapshot: pd.DataFrame) -> list[str]:
    return [
        column
        for column in snapshot.columns
        if column not in NON_FEATURE_COLUMNS
        and not column.lower().startswith("unnamed")
        and pd.api.types.is_numeric_dtype(snapshot[column])
    ]