batteryswap_example/train.py

import pickle
import os
from pathlib import Path
from typing import Optional, Sequence

import pandas
import pandas as pd
import numpy as np
from pydantic_settings import BaseSettings, SettingsConfigDict
import structlog
from ortools.sat.python import cp_model

from batteryswap_public.interfaces import Planner, RULModel
from batteryswap_public.utils import load_dataset, iterate_scenarios
from batteryswap_public.evaluate import evaluate_plan, check_plan_valid

log = structlog.get_logger()

COST_SCALE = 600
MINUTES_PER_HOUR = 60

DEVICE_COLUMN = "device_id"
TIME_COLUMN = "end_time"
VALUE_COLUMNS = ("voltage", "temperature")


def normalize_timeseries(timeseries):
    frame = timeseries.copy()
    missing_identity = {DEVICE_COLUMN, TIME_COLUMN} - set(frame.columns)
    if missing_identity:
        frame = frame.reset_index()

    required = {DEVICE_COLUMN, TIME_COLUMN, *VALUE_COLUMNS}
    missing = required - set(frame.columns)
    if missing:
        raise ValueError(f"Timeseries is missing required columns: {sorted(missing)}")

    frame = frame.loc[:, [DEVICE_COLUMN, TIME_COLUMN, *VALUE_COLUMNS]].copy()
    frame[TIME_COLUMN] = pandas.to_datetime(frame[TIME_COLUMN])
    frame[DEVICE_COLUMN] = frame[DEVICE_COLUMN].astype(str)
    for column in VALUE_COLUMNS:
        frame[column] = pandas.to_numeric(frame[column], errors="coerce")
    return frame.sort_values([DEVICE_COLUMN, TIME_COLUMN], kind="stable").reset_index(
        drop=True
    )


def _setting(settings, name, default):
    if isinstance(settings, dict):
        return settings.get(name, default)
    return getattr(settings, name, default)


def _normalize_locations(locations):
    frame = locations.copy().reset_index(drop=True)
    aliases = {"device_id": "battery", "building_id": "building", "room_id": "room"}
    frame = frame.rename(
        columns={old: new for old, new in aliases.items() if new not in frame}
    )
    required = {"battery", "building", "room"}
    missing = required - set(frame.columns)
    if missing:
        raise ValueError(f"Locations is missing required columns: {sorted(missing)}")
    if frame["battery"].duplicated().any():
        raise ValueError("Each battery must have exactly one location")
    return frame


def _travel_lookup(travel_costs):
    frame = travel_costs.copy()
    required = {"from", "to", "hours"}
    missing = required - set(frame.columns)
    if missing:
        raise ValueError(f"Travel costs is missing required columns: {sorted(missing)}")
    return {
        (str(row["from"]), str(row["to"])): float(row["hours"])
        for _, row in frame.iterrows()
    }


def order_daily_route(batteries, locations, travel_costs, base_building):
    selected = set(str(value) for value in batteries)
    if not selected:
        return []
    loc = _normalize_locations(locations).set_index("battery")
    travel = _travel_lookup(travel_costs)
    buildings = set(loc.loc[list(selected), "building"].astype(str))
    current = str(base_building)
    building_order = []
    while buildings:
        next_building = min(
            buildings,
            key=lambda building: (
                travel.get(
                    (current, building), 0.0 if current == building else float("inf")
                ),
                building,
            ),
        )
        building_order.append(next_building)
        buildings.remove(next_building)
        current = next_building

    ordered = []
    for building in building_order:
        subset = loc.loc[list(selected)]
        subset = subset.loc[subset["building"].astype(str) == building].copy()
        subset["battery_key"] = subset.index.astype(str)
        subset = subset.sort_values(["room", "battery_key"], kind="stable")
        ordered.extend(subset.index.astype(str).tolist())
    return ordered


class MilpPlanner(Planner):
    def __init__(
        self,
        rul_estimator,
        solver_time_limit_seconds=30.0,
        late_risk_multiplier=1.0,
        candidate_benefit_threshold=30.0,
        solver_workers=8,
    ):
        self.rul_estimator = rul_estimator
        self.solver_time_limit_seconds = float(solver_time_limit_seconds)
        self.late_risk_multiplier = float(late_risk_multiplier)
        # Only batteries where swapping beats skipping by more than this margin
        # get decision variables. With ~450 batteries but only ~2-4% actually
        # due in the window, this shrinks the model ~20x and lets CP-SAT reach
        # optimality instead of timing out at a 99.9% gap.
        self.candidate_benefit_threshold = float(candidate_benefit_threshold)
        self.solver_workers = int(solver_workers)

    def _expected_costs(self, timeseries, batteries, settings):
        horizon = int(round(float(_setting(settings, "planning_window_days", 42))))
        normalized = normalize_timeseries(timeseries)
        scenario_start = normalized["end_time"].max().normalize()
        horizon_end = scenario_start + pandas.Timedelta(days=horizon)
        emergency_delay = 6 - horizon_end.weekday()
        costs = self.rul_estimator.expected_replacement_costs(
            timeseries,
            horizon_days=horizon,
            early_penalty=float(
                _setting(settings, "early_replacement_penalty_daily", 0.5)
            ),
            late_penalty=float(
                _setting(settings, "late_replacement_penalty_daily", 10.0)
            )
            * float(getattr(self, "late_risk_multiplier", 1.0)),
            no_swap_extension_days=emergency_delay,
        )
        expected_columns = list(range(horizon + 1)) + ["no_swap"]
        costs = costs.reindex(index=batteries, columns=expected_columns)
        finite = costs.to_numpy(dtype=float)
        fallback = (
            float(np.nanmax(finite[np.isfinite(finite)]))
            if np.isfinite(finite).any()
            else 1e6
        )
        return costs.replace([np.inf, -np.inf], np.nan).fillna(fallback + 1e3)

    def _solve_assignments(self, expected_costs, locations, travel_costs, settings):
        loc = _normalize_locations(locations).set_index("battery")
        batteries = list(expected_costs.index.astype(str))
        horizon = int(round(float(_setting(settings, "planning_window_days", 42))))
        real_days = list(range(horizon + 1))
        actions = real_days + ["no_swap"]
        base = str(_setting(settings, "base_location", ""))
        base_room = str(_setting(settings, "base_room", ""))
        travel = _travel_lookup(travel_costs)

        model = cp_model.CpModel()
        assignment = {
            (battery, action): model.new_bool_var(f"x_{index}_{action}")
            for index, battery in enumerate(batteries)
            for action in actions
        }
        for battery in batteries:
            model.add_exactly_one(assignment[battery, action] for action in actions)

        rooms = sorted(loc.loc[batteries, "room"].astype(str).unique())
        buildings = sorted(loc.loc[batteries, "building"].astype(str).unique())
        battery_rooms = loc.loc[batteries, "room"].astype(str).to_dict()
        battery_buildings = loc.loc[batteries, "building"].astype(str).to_dict()
        room_members = {
            room: [battery for battery in batteries if battery_rooms[battery] == room]
            for room in rooms
        }
        building_members = {
            building: [
                battery
                for battery in batteries
                if battery_buildings[battery] == building
            ]
            for building in buildings
        }
        room_visit = {
            (room, day): model.new_bool_var(f"room_{room}_{day}")
            for room in rooms
            for day in real_days
        }
        building_visit = {
            (building, day): model.new_bool_var(f"building_{building}_{day}")
            for building in buildings
            for day in real_days
        }

        for day in real_days:
            for room in rooms:
                members = room_members[room]
                for battery in members:
                    model.add(assignment[battery, day] <= room_visit[room, day])
                model.add(
                    room_visit[room, day]
                    <= sum(assignment[battery, day] for battery in members)
                )
            for building in buildings:
                members = building_members[building]
                for battery in members:
                    model.add(assignment[battery, day] <= building_visit[building, day])
                model.add(
                    building_visit[building, day]
                    <= sum(assignment[battery, day] for battery in members)
                )

        battery_minutes = round(
            float(_setting(settings, "time_per_battery_hours", 0.25)) * 60
        )
        room_minutes = round(
            float(_setting(settings, "time_per_room_change_hours", 0.5)) * 60
        )
        building_minutes = round(
            float(_setting(settings, "time_per_building_change_hours", 1.0)) * 60
        )
        building_work = {}
        for building in buildings:
            round_trip = travel.get((base, building), 0.0 if base == building else 24.0)
            round_trip += travel.get(
                (building, base), 0.0 if base == building else 24.0
            )
            building_work[building] = round(round_trip * 60) + (
                0 if building == base else building_minutes
            )

        maximum_daily = (
            len(batteries) * battery_minutes
            + len(rooms) * room_minutes
            + sum(building_work.values())
        )
        daily_work = {}
        daily_overtime = {}
        daily_limit_hit = {}
        overtime_start = round(float(_setting(settings, "overtime_start", 8.0)) * 60)
        daily_limit = round(
            float(_setting(settings, "worker_limit_daily_hours", 24.0)) * 60
        )

        for day in real_days:
            work = model.new_int_var(0, maximum_daily, f"work_{day}")
            expression = (
                battery_minutes * sum(assignment[battery, day] for battery in batteries)
                + room_minutes
                * sum(room_visit[room, day] for room in rooms if room != base_room)
                + sum(
                    building_work[building] * building_visit[building, day]
                    for building in buildings
                )
            )
            model.add(work == expression)
            daily_work[day] = work

            overtime = model.new_int_var(0, maximum_daily, f"overtime_{day}")
            model.add(overtime >= work - overtime_start)
            daily_overtime[day] = overtime

            hit = model.new_bool_var(f"daily_limit_hit_{day}")
            model.add(work <= daily_limit + maximum_daily * hit)
            daily_limit_hit[day] = hit

        weekly_limit_hit = {}
        weekly_limit = round(
            float(_setting(settings, "worker_limit_weekly_hours", 24.0)) * 60
        )
        for week_start in range(0, len(real_days), 7):
            week_days = real_days[week_start : week_start + 7]
            hit = model.new_bool_var(f"weekly_limit_hit_{week_start // 7}")
            weekly_maximum = maximum_daily * len(week_days)
            model.add(
                sum(daily_work[day] for day in week_days)
                <= max(weekly_limit - 1, -1) + weekly_maximum * hit
            )
            weekly_limit_hit[week_start] = hit

        objective_terms = []
        for battery_index, battery in enumerate(batteries):
            for action_index, action in enumerate(actions):
                coefficient = int(
                    round(float(expected_costs.loc[battery, action]) * COST_SCALE)
                )
                coefficient += action_index + battery_index % 3
                objective_terms.append(coefficient * assignment[battery, action])

        minute_cost = COST_SCALE // MINUTES_PER_HOUR
        objective_terms.extend(minute_cost * daily_work[day] for day in real_days)
        overtime_factor = float(_setting(settings, "overtime_penalty_factor", 2.0))
        overtime_minute_cost = int(round(overtime_factor * minute_cost))
        objective_terms.extend(
            overtime_minute_cost * daily_overtime[day] for day in real_days
        )
        daily_penalty = int(
            round(
                float(_setting(settings, "worker_limit_daily_penalty", 100.0))
                * COST_SCALE
            )
        )
        weekly_penalty = int(
            round(
                float(_setting(settings, "worker_limit_weekly_penalty", 100.0))
                * COST_SCALE
            )
        )
        objective_terms.extend(
            daily_penalty * value for value in daily_limit_hit.values()
        )
        objective_terms.extend(
            weekly_penalty * value for value in weekly_limit_hit.values()
        )
        model.minimize(sum(objective_terms))

        solver = cp_model.CpSolver()
        solver.parameters.max_time_in_seconds = self.solver_time_limit_seconds
        solver.parameters.num_search_workers = max(1, int(getattr(self, "solver_workers", 8)))
        solver.parameters.random_seed = 0
        status = solver.solve(model)
        # Never silently accept a non-solve: an UNKNOWN status here previously
        # fell through to the greedy fallback and cost us most of the late_swap.
        log.info(
            "cpsat-solve",
            status=solver.status_name(status),
            batteries=len(batteries),
            wall_seconds=round(solver.wall_time, 2),
        )
        if status not in (cp_model.OPTIMAL, cp_model.FEASIBLE):
            log.warning("cpsat-no-solution-using-greedy-fallback", status=solver.status_name(status))
            return {
                battery: min(
                    actions,
                    key=lambda action: (
                        expected_costs.loc[battery, action],
                        str(action),
                    ),
                )
                for battery in batteries
            }
        return {
            battery: next(
                action
                for action in actions
                if solver.value(assignment[battery, action])
            )
            for battery in batteries
        }

    def plan(self, timeseries, locations, travel_costs, settings):
        loc = _normalize_locations(locations)
        batteries = sorted(loc["battery"].astype(str).tolist())
        normalized_timeseries = normalize_timeseries(timeseries)
        if normalized_timeseries.empty:
            if "end_time" not in loc:
                raise ValueError("Cannot determine scenario start time")
            start_time = pandas.to_datetime(loc["end_time"]).max().normalize()
        else:
            start_time = normalized_timeseries["end_time"].max().normalize()

        expected_costs = self._expected_costs(timeseries, batteries, settings)

        # Restrict the solve to batteries where swapping actually beats skipping.
        # Only ~2-4% of batteries are due in any window; giving all ~450 of them
        # decision variables made the model too big to solve (CP-SAT returned
        # UNKNOWN at a 99.9% gap). Everything below the threshold is pinned to
        # no_swap, which its own expected costs already say is optimal.
        day_costs = expected_costs.drop(columns=["no_swap"])
        benefit = expected_costs["no_swap"] - day_costs.min(axis=1)
        # getattr keeps planners pickled before this attribute existed loadable
        threshold = float(getattr(self, "candidate_benefit_threshold", 30.0))
        candidates = sorted(benefit[benefit > threshold].index.astype(str))
        log.info("candidate-filter", total=len(batteries), candidates=len(candidates))

        assignments = {battery: "no_swap" for battery in batteries}
        if candidates:
            solved = self._solve_assignments(
                expected_costs.loc[candidates], loc, travel_costs, settings
            )
            assignments.update(solved)

        horizon = int(round(float(_setting(settings, "planning_window_days", 42))))
        base = str(_setting(settings, "base_location", ""))
        records = []
        for day in range(horizon + 1):
            selected = [battery for battery in batteries if assignments[battery] == day]
            for battery in order_daily_route(selected, loc, travel_costs, base):
                records.append(
                    {"day": start_time + pandas.Timedelta(days=day), "battery": battery}
                )

        no_swap_day = start_time + pandas.Timedelta(days=horizon + 1)
        for battery in sorted(
            battery for battery in batteries if assignments[battery] == "no_swap"
        ):
            records.append({"day": no_swap_day, "battery": battery})

        plan = pandas.DataFrame.from_records(
            records, columns=["day", "battery"]
        ).reset_index(drop=True)
        plan["day"] = pandas.to_datetime(plan["day"])
        check_plan_valid(plan, loc, start_time=start_time)
        return plan


class SearchPlanner(Planner):
    """Greedy construction + local search scored by the *real* evaluate_plan.

    The MILP approximates the official cost model; this scores candidate plans
    with the actual evaluator (using predicted EOL as surrogate truth), so it
    has zero modeling error -- it sees the emergency-visit mechanics, weekly
    limit accounting and end-of-day travel exactly as the scorer does.

    Only batteries that plausibly fail inside the window get scheduled; with
    ~2-4% due per scenario the search space is small enough to explore well.
    """

    def __init__(
        self,
        rul_estimator,
        candidate_benefit_threshold=5.0,
        iterations=400,
        random_seed=0,
        late_risk_multiplier=1.0,
    ):
        self.rul_estimator = rul_estimator
        self.candidate_benefit_threshold = float(candidate_benefit_threshold)
        self.iterations = int(iterations)
        self.random_seed = int(random_seed)
        self.late_risk_multiplier = float(late_risk_multiplier)

    def _expected_costs(self, timeseries, batteries, settings):
        return MilpPlanner._expected_costs(self, timeseries, batteries, settings)

    @staticmethod
    def _build_plan(assignment, all_batteries, start_time, park_day):
        """assignment: battery -> day offset (int). Others parked past window."""
        records = [
            {"day": start_time + pandas.Timedelta(days=int(d)), "battery": b}
            for b, d in assignment.items()
        ]
        assigned = set(assignment)
        records.extend(
            {"day": park_day, "battery": b} for b in all_batteries if b not in assigned
        )
        plan = pandas.DataFrame.from_records(records, columns=["day", "battery"])
        # stable ordering: by day, then grouped by building/room handled by caller
        plan = plan.sort_values(["day", "battery"], kind="stable").reset_index(drop=True)
        plan["day"] = pandas.to_datetime(plan["day"])
        return plan

    @staticmethod
    def _route_day(day_batteries, building_of, room_of, travel, base):
        """Nearest-building route for one day, using precomputed lookups.

        Same logic as order_daily_route but without rebuilding the location
        frame and travel dict on every call -- this runs inside the search loop.
        """
        remaining = set(building_of[b] for b in day_batteries)
        current = str(base)
        building_order = []
        while remaining:
            nxt = min(
                remaining,
                key=lambda bl: (
                    travel.get((current, bl), 0.0 if current == bl else float("inf")),
                    bl,
                ),
            )
            building_order.append(nxt)
            remaining.discard(nxt)
            current = nxt
        ordered = []
        for building in building_order:
            here = [b for b in day_batteries if building_of[b] == building]
            here.sort(key=lambda b: (room_of[b], b))
            ordered.extend(here)
        return ordered

    def plan(self, timeseries, locations, travel_costs, settings):
        import random

        loc = _normalize_locations(locations)
        batteries = sorted(loc["battery"].astype(str).tolist())
        normalized = normalize_timeseries(timeseries)
        if normalized.empty:
            start_time = pandas.to_datetime(loc["end_time"]).max().normalize()
        else:
            start_time = normalized["end_time"].max().normalize()

        horizon = int(round(float(_setting(settings, "planning_window_days", 42))))
        horizon_end = start_time + pandas.Timedelta(days=horizon)
        base = str(_setting(settings, "base_location", ""))
        park_day = start_time + pandas.Timedelta(days=horizon + 1)

        expected_costs = self._expected_costs(timeseries, batteries, settings)
        day_costs = expected_costs.drop(columns=["no_swap"])
        benefit = expected_costs["no_swap"] - day_costs.min(axis=1)
        candidates = sorted(
            benefit[benefit > self.candidate_benefit_threshold].index.astype(str)
        )
        log.info("candidate-filter", total=len(batteries), candidates=len(candidates))
        if not candidates:
            plan = self._build_plan({}, batteries, start_time, park_day)
            check_plan_valid(plan, loc, start_time=start_time)
            return plan

        # Surrogate EOL for scoring: the cost-minimising swap day per battery
        # already encodes the newsvendor trade-off (early 0.5/day vs late 10/day).
        target_day = {b: int(day_costs.loc[b].idxmin()) for b in candidates}
        surrogate_eol = pandas.Series(
            {
                b: (start_time + pandas.Timedelta(days=int(target_day[b])))
                if b in target_day
                else park_day + pandas.Timedelta(days=365)
                for b in batteries
            }
        )

        # Precomputed lookups so the search loop never rebuilds these.
        indexed = loc.set_index("battery")
        building_of = indexed["building"].astype(str).to_dict()
        room_of = indexed["room"].astype(str).to_dict()
        travel = _travel_lookup(travel_costs)
        parked = [b for b in batteries]

        def make_plan(assignment):
            by_day = {}
            for b, d in assignment.items():
                by_day.setdefault(int(d), []).append(b)
            rows_day, rows_bat = [], []
            for d in sorted(by_day):
                ordered = self._route_day(
                    sorted(by_day[d]), building_of, room_of, travel, base
                )
                stamp = start_time + pandas.Timedelta(days=d)
                rows_day.extend([stamp] * len(ordered))
                rows_bat.extend(ordered)
            assigned = set(assignment)
            rest = [b for b in parked if b not in assigned]
            rows_day.extend([park_day] * len(rest))
            rows_bat.extend(rest)
            plan = pandas.DataFrame({"day": rows_day, "battery": rows_bat})
            plan["day"] = pandas.to_datetime(plan["day"])
            return plan

        def score(assignment):
            try:
                _, _, overall = evaluate_plan(
                    make_plan(assignment),
                    loc,
                    travel_costs,
                    settings,
                    eol_times=surrogate_eol,
                    start_time=start_time,
                    verbose=0,
                )
                return float(overall["total_cost"])
            except Exception:
                return float("inf")

        # Greedy start: everyone at their own cost-minimising day.
        current = dict(target_day)
        current_cost = score(current)

        # Local search. Moves: shift a day, batch onto another candidate's day,
        # drop a battery, or restore a dropped one.
        rng = random.Random(self.random_seed)
        best, best_cost = dict(current), current_cost
        used_days = sorted(set(target_day.values()))
        for _ in range(self.iterations):
            trial = dict(current)
            battery = rng.choice(candidates)
            move = rng.random()
            if move < 0.4 and trial:
                # batch: move onto a day already being worked (saves a trip)
                if used_days:
                    trial[battery] = rng.choice(
                        sorted(set(trial.values())) or used_days
                    )
            elif move < 0.75:
                # shift, biased earlier (late costs 20x more than early)
                base_day = trial.get(battery, target_day[battery])
                shift = -rng.randint(1, 7) if rng.random() < 0.7 else rng.randint(1, 4)
                trial[battery] = int(min(max(base_day + shift, 0), horizon))
            elif move < 0.9:
                trial.pop(battery, None)  # drop
            else:
                trial[battery] = target_day[battery]  # restore

            trial_cost = score(trial)
            if trial_cost <= current_cost:
                current, current_cost = trial, trial_cost
                if trial_cost < best_cost:
                    best, best_cost = dict(trial), trial_cost

        log.info(
            "search-planner",
            candidates=len(candidates),
            scheduled=len(best),
            surrogate_cost=round(best_cost, 1),
        )
        plan = make_plan(best)
        check_plan_valid(plan, loc, start_time=start_time)
        return plan


RUL_DEVICE_COLUMN = "device_id"
RUL_TIME_COLUMN = "end_time"
RUL_VALUE_COLUMNS = ("voltage", "temperature")
RUL_WINDOW_DAYS = (7, 30, 90)


def _rul_slope(values, timestamps):
    y_all = values.to_numpy(dtype=float)
    time_all = timestamps.to_numpy(dtype="datetime64[ns]")
    valid = np.isfinite(y_all) & ~np.isnat(time_all)
    if np.count_nonzero(valid) < 2:
        return 0.0
    y = y_all[valid]
    selected_times = time_all[valid]
    x = (selected_times - selected_times.min()) / np.timedelta64(1, "D")
    x = x.astype(float)
    centered_x = x - x.mean()
    denominator = float(centered_x @ centered_x)
    if denominator == 0.0:
        return 0.0
    return float(centered_x @ (y - y.mean()) / denominator)


def _rul_finite(value):
    return float(value) if np.isfinite(value) else 0.0


def extract_snapshot_features(timeseries, reference_time=None, windows=RUL_WINDOW_DAYS):
    frame = normalize_timeseries(timeseries)
    if frame.empty:
        return pd.DataFrame(index=pd.Index([], name=RUL_DEVICE_COLUMN))

    reference = (
        pd.Timestamp(reference_time)
        if reference_time is not None
        else frame[RUL_TIME_COLUMN].max()
    )
    frame = frame.loc[frame[RUL_TIME_COLUMN] <= reference].copy()
    rows = []

    for device_id, group in frame.groupby(RUL_DEVICE_COLUMN, sort=True):
        group = group.sort_values(RUL_TIME_COLUMN, kind="stable")
        first_time = group[RUL_TIME_COLUMN].iloc[0]
        last_time = group[RUL_TIME_COLUMN].iloc[-1]
        row = {
            RUL_DEVICE_COLUMN: device_id,
            "device_age_days": max(
                (reference - first_time).total_seconds() / 86400.0, 0.0
            ),
            "history_span_days": max(
                (last_time - first_time).total_seconds() / 86400.0, 0.0
            ),
            "days_since_last_observation": max(
                (reference - last_time).total_seconds() / 86400.0, 0.0
            ),
            "observation_count": float(len(group)),
        }

        for value_column in RUL_VALUE_COLUMNS:
            values = group[value_column]
            row[f"{value_column}_latest"] = _rul_finite(values.iloc[-1])
            row[f"{value_column}_mean"] = _rul_finite(values.mean())
            row[f"{value_column}_std"] = _rul_finite(values.std(ddof=0))
            row[f"{value_column}_min"] = _rul_finite(values.min())
            row[f"{value_column}_max"] = _rul_finite(values.max())
            row[f"{value_column}_slope"] = _rul_finite(
                _rul_slope(values, group[RUL_TIME_COLUMN])
            )

        for days in windows:
            window = group.loc[
                group[RUL_TIME_COLUMN] >= reference - pd.Timedelta(days=int(days))
            ]
            row[f"observation_count_{days}d"] = float(len(window))
            for value_column in RUL_VALUE_COLUMNS:
                values = window[value_column]
                prefix = f"{value_column}_{days}d"
                row[f"{prefix}_mean"] = _rul_finite(values.mean())
                row[f"{prefix}_std"] = _rul_finite(values.std(ddof=0))
                row[f"{prefix}_min"] = _rul_finite(values.min())
                row[f"{prefix}_max"] = _rul_finite(values.max())
                row[f"{value_column}_slope_{days}d"] = _rul_finite(
                    _rul_slope(values, window[RUL_TIME_COLUMN])
                )

        rows.append(row)

    features = pd.DataFrame.from_records(rows).set_index(RUL_DEVICE_COLUMN)
    return features.astype(float)


def expected_costs_from_failure_distribution(
    failure_probability, replacement_days, early_penalty, late_penalty
):
    probabilities = np.asarray(failure_probability, dtype=float)
    probabilities = np.clip(probabilities, 0.0, None)
    total_probability = probabilities.sum()
    if total_probability <= 0:
        raise ValueError("Failure probabilities must contain positive mass")
    probabilities = probabilities / total_probability

    failure_days = np.arange(len(probabilities), dtype=float)
    replacement = np.asarray(replacement_days, dtype=float)[:, None]
    early_days = np.maximum(failure_days[None, :] - replacement, 0.0)
    late_days = np.maximum(replacement - failure_days[None, :], 0.0)
    costs = early_penalty * early_days + late_penalty * late_days
    return costs @ probabilities


def expected_no_swap_cost(
    failure_probability, horizon_day, emergency_day, late_penalty
):
    probabilities = np.asarray(failure_probability, dtype=float)
    probabilities = np.clip(probabilities, 0.0, None)
    total_probability = probabilities.sum()
    if total_probability <= 0:
        raise ValueError("Failure probabilities must contain positive mass")
    probabilities = probabilities / total_probability
    failure_days = np.arange(len(probabilities), dtype=float)
    due_inside_window = failure_days <= int(horizon_day)
    late_days = np.maximum(float(emergency_day) - failure_days, 0.0)
    return float(
        np.sum(probabilities[due_inside_window] * late_days[due_inside_window])
        * late_penalty
    )


class DiscreteHazardRULModel(RULModel):
    quantile_cols = ["p10", "p50", "p90"]
    period_days = 7
    horizon_cap_days = 126  # ~18 weekly periods; covers the 42-day planning window plus emergency margin

    def __init__(self, random_state=0):
        self.random_state = int(random_state)
        self.model_ = None
        self.feature_columns_ = []
        self.feature_medians_ = pd.Series(dtype=float)
        self.n_periods_ = self.horizon_cap_days // self.period_days
        self.fallback_hazard_ = 0.01

    def _prepare_features(self, features, fitting=False):
        numeric = features.apply(pd.to_numeric, errors="coerce").replace(
            [np.inf, -np.inf], np.nan
        )
        if fitting:
            self.feature_columns_ = list(numeric.columns)
            self.feature_medians_ = numeric.median().fillna(0.0)
        else:
            numeric = numeric.reindex(columns=self.feature_columns_)
        return numeric.fillna(self.feature_medians_).astype(float)

    def _expand_person_periods(self, features, durations, events):
        n_periods = self.n_periods_
        row_index = []
        row_periods = []
        labels = []
        for idx in features.index:
            duration = float(durations[idx])
            event = bool(events[idx])
            capped_duration = min(duration, float(self.horizon_cap_days))
            failure_period = None
            if event and duration <= self.horizon_cap_days:
                failure_period = min(
                    int(capped_duration // self.period_days), n_periods - 1
                )
            max_period = int(np.ceil(capped_duration / self.period_days))
            if failure_period is not None:
                max_period = max(max_period, failure_period + 1)
            max_period = min(max_period, n_periods)
            for period in range(max_period):
                row_index.append(idx)
                row_periods.append(period)
                is_failure = failure_period is not None and period == failure_period
                labels.append(1 if is_failure else 0)
                if is_failure:
                    break
        period_features = features.loc[row_index].copy()
        period_features["period"] = row_periods
        return period_features, np.array(labels, dtype=int)

    def fit_snapshots(self, snapshot_features, durations, events):
        common = snapshot_features.index.intersection(durations.index).intersection(
            events.index
        )
        if common.empty:
            raise ValueError("No aligned snapshot labels were provided")
        features = self._prepare_features(snapshot_features.loc[common], fitting=True)
        duration = pd.to_numeric(durations.loc[common], errors="coerce").clip(
            lower=0.25
        )
        event = events.loc[common].fillna(False).astype(bool)

        period_features, labels = self._expand_person_periods(features, duration, event)
        self.fallback_hazard_ = (
            float(np.clip(labels.mean(), 1e-3, 0.5)) if len(labels) else 0.01
        )

        self.model_ = None
        if labels.sum() >= 2 and len(labels) >= 10:
            from sklearn.ensemble import HistGradientBoostingClassifier

            model = HistGradientBoostingClassifier(random_state=self.random_state)
            model.fit(period_features.to_numpy(dtype=float), labels)
            self.model_ = model
        return self

    def fit(self, timeseries, rul):
        features = extract_snapshot_features(timeseries)
        labels = pd.to_numeric(rul, errors="coerce").reindex(features.index)
        events = pd.Series(True, index=features.index)
        return self.fit_snapshots(features, labels, events)

    def _period_hazards(self, features):
        prepared = self._prepare_features(features, fitting=False)
        n = len(prepared)
        n_periods = self.n_periods_
        hazards = np.full((n, n_periods), self.fallback_hazard_, dtype=float)
        if self.model_ is not None:
            base = prepared.to_numpy(dtype=float)
            for period in range(n_periods):
                period_col = np.full((n, 1), float(period), dtype=float)
                x = np.hstack([base, period_col])
                try:
                    hazards[:, period] = self.model_.predict_proba(x)[:, 1]
                except (ArithmeticError, ValueError):
                    pass
        return np.clip(hazards, 1e-4, 1.0 - 1e-4)

    def _survival(self, features, times):
        hazards = self._period_hazards(features)
        n = hazards.shape[0]
        period_survival = np.cumprod(1.0 - hazards, axis=1)
        period_survival = np.hstack(
            [np.ones((n, 1)), period_survival]
        )  # prepend day-0 survival = 1

        times = np.asarray(times, dtype=float)
        result = np.ones((len(times), n), dtype=float)
        for i, t in enumerate(times):
            period_idx = min(int(t // self.period_days), self.n_periods_)
            result[i, :] = period_survival[:, period_idx]
        return result

    def predict(self, timeseries):
        features = extract_snapshot_features(timeseries)
        times = np.arange(
            0, self.horizon_cap_days + self.period_days, self.period_days, dtype=float
        )
        survival = self._survival(features, times)
        quantile_days = {}
        for q_col, target in zip(self.quantile_cols, (0.9, 0.5, 0.1)):
            days = []
            for j in range(survival.shape[1]):
                below = np.where(survival[:, j] <= target)[0]
                days.append(
                    float(times[below[0]])
                    if len(below)
                    else float(self.horizon_cap_days)
                )
            quantile_days[q_col] = days
        return pd.DataFrame(quantile_days, index=features.index)[self.quantile_cols]

    def failure_probabilities(self, timeseries, max_day):
        features = extract_snapshot_features(timeseries)
        times = np.arange(max_day + 2, dtype=float)
        survival = self._survival(features, times)
        interval_mass = np.maximum(survival[:-1] - survival[1:], 0.0)
        probabilities = np.vstack([interval_mass, survival[-1:]]).T
        row_sums = probabilities.sum(axis=1, keepdims=True)
        probabilities = np.divide(
            probabilities,
            row_sums,
            out=np.zeros_like(probabilities),
            where=row_sums > 0,
        )
        return pd.DataFrame(
            probabilities, index=features.index, columns=range(max_day + 2)
        )

    def expected_replacement_costs(
        self,
        timeseries,
        horizon_days,
        early_penalty,
        late_penalty,
        no_swap_extension_days,
    ):
        emergency_day = int(horizon_days + no_swap_extension_days)
        max_failure_day = emergency_day + max(int(horizon_days), 30)
        probabilities = self.failure_probabilities(timeseries, max_day=max_failure_day)
        replacement_days = np.arange(int(horizon_days) + 1)
        rows = [
            np.append(
                expected_costs_from_failure_distribution(
                    row,
                    replacement_days,
                    early_penalty=float(early_penalty),
                    late_penalty=float(late_penalty),
                ),
                expected_no_swap_cost(
                    row,
                    horizon_day=int(horizon_days),
                    emergency_day=emergency_day,
                    late_penalty=float(late_penalty),
                ),
            )
            for row in probabilities.to_numpy(dtype=float)
        ]
        columns = list(replacement_days) + ["no_swap"]
        return pd.DataFrame(rows, index=probabilities.index, columns=columns)


def split_scenarios(scenarios, val_fraction=0.25, seed=0):
    """Deterministic shuffled train/val split across scenarios (not a
    prefix-limit) so evaluation isn't done on the same data used to fit."""
    import random

    rng = random.Random(seed)
    shuffled = list(scenarios)
    rng.shuffle(shuffled)
    n_val = max(1, round(len(shuffled) * val_fraction))
    return shuffled[n_val:], shuffled[:n_val]


def build_training_snapshots(
    locations, timeseries, eol_times, scenarios, limit_scenarios=None
):
    feature_parts = []
    duration_parts = []
    event_parts = []

    gen = iterate_scenarios(locations, timeseries, eol_times, scenarios)
    for scenario_number, (scenario, locs, cut, scenario_eol) in enumerate(gen):
        if limit_scenarios is not None and scenario_number >= limit_scenarios:
            break
        scenario_name = str(scenario["name"])
        scenario_start = pd.Timestamp(scenario["start_time"])
        features = extract_snapshot_features(cut, reference_time=scenario_start)
        batteries = features.index.astype(str)

        loc_by_battery = locs.set_index("battery")
        observed_eol = pd.to_datetime(scenario_eol.reindex(batteries))
        censor_end = pd.to_datetime(loc_by_battery.loc[batteries, "end_time"])
        event = observed_eol.notna()
        endpoint = observed_eol.where(event, censor_end)
        duration = ((endpoint - scenario_start) / pd.Timedelta(days=1)).astype(float)
        duration = duration.clip(lower=0.25)

        snapshot_index = pd.Index(
            [f"{battery}::{scenario_name}" for battery in batteries], name="snapshot_id"
        )
        features = features.copy()
        features.index = snapshot_index
        duration.index = snapshot_index
        event.index = snapshot_index
        feature_parts.append(features)
        duration_parts.append(duration.rename("duration"))
        event_parts.append(event.astype(bool).rename("event"))

    if not feature_parts:
        raise ValueError("No training scenarios produced snapshot features")
    return (
        pd.concat(feature_parts, axis=0),
        pd.concat(duration_parts, axis=0),
        pd.concat(event_parts, axis=0),
    )


def train_rul_model(locations, timeseries, eol_times, scenarios, limit_scenarios=None):
    features, durations, events = build_training_snapshots(
        locations, timeseries, eol_times, scenarios, limit_scenarios=limit_scenarios
    )
    log.info(
        "training-snapshots",
        rows=len(features),
        features=len(features.columns),
        observed_events=int(events.sum()),
        censored=int((~events).sum()),
    )
    return DiscreteHazardRULModel().fit_snapshots(features, durations, events)


class Config(BaseSettings):
    """
    Automatically provides command-line argument support for specified fields
    """

    model_config = SettingsConfigDict(
        env_prefix="",
        cli_parse_args=True,
        cli_ignore_unknown_args=True,
    )

    dataset_path: Optional[Path] = None
    split: str = "train"
    solver_time_limit_seconds: float = 20.0
    late_risk_multiplier: float = 1.0
    val_fraction: float = 0.25
    split_seed: int = 0


def main():
    cfg = Config()

    if cfg.dataset_path is None:
        dataset_path = os.environ.get("BATTERYSWAP_DATASET_PATH", None)
        assert dataset_path
        dataset_path = Path(dataset_path)
    else:
        dataset_path = cfg.dataset_path

    split_path = dataset_path / cfg.split
    locations, timeseries, eol_times, scenarios = load_dataset(split_path)

    log.info("evaluate-load-data", path=dataset_path)

    train_scenarios, val_scenarios = split_scenarios(
        scenarios, val_fraction=cfg.val_fraction, seed=cfg.split_seed
    )
    log.info(
        "scenario-split",
        total=len(scenarios),
        train=len(train_scenarios),
        val=len(val_scenarios),
    )
    rul_model = train_rul_model(locations, timeseries, eol_times, train_scenarios)
    log.info("train-done")

    log.info("evaluate-held-out")
    gen = iterate_scenarios(locations, timeseries, eol_times, val_scenarios)
    for scenario, locs, cut, eol in gen:
        scenario_name = scenario["name"]
        travel_costs = scenario["travel_costs"]
        settings = scenario["settings"]

        planner = MilpPlanner(
            rul_model,
            solver_time_limit_seconds=cfg.solver_time_limit_seconds,
            late_risk_multiplier=cfg.late_risk_multiplier,
        )
        plan = planner.plan(cut, locs, travel_costs, settings)

        start_time = pandas.Timestamp(scenario["start_time"])

        transitions, daily, overall = evaluate_plan(
            plan, locs, travel_costs, settings, eol_times=eol, start_time=start_time
        )

        print("scores", scenario_name, overall)

    log.info("refit-on-full-data-for-submission")
    rul_model_full = train_rul_model(locations, timeseries, eol_times, scenarios)

    # Save best planner
    planner = MilpPlanner(
        rul_model_full,
        solver_time_limit_seconds=cfg.solver_time_limit_seconds,
        late_risk_multiplier=cfg.late_risk_multiplier,
    )

    planner_path = 'batteryswap_example/planners/best.pickle' 
    with open(planner_path, "wb") as f:
        pickle.dump(planner, f)
        print('planner-save', planner_path)


if __name__ == '__main__':
    main()