from __future__ import annotations from dataclasses import dataclass import numpy as np import pandas as pd from batteryswap_public.interfaces import Planner from .competition_features import build_daily_features, scenario_snapshot from .competition_model import EventTimeModel @dataclass(frozen=True) class PlannerPolicy: event_risk_threshold: float = 0.10 prediction_offset_days: float = -25.0 building_batch_window_days: int = 4 unscheduled_margin_days: int = 1 capacity_lookback_days: int = 14 capacity_lookahead_days: int = 7 capacity_daily_limit_fraction: float = 1.0 capacity_weekly_limit_fraction: float = 1.0 capacity_limit_penalty_multiplier: float = 1.0 capacity_operational_cost_weight: float = 1.5 use_expected_cost: bool = False risk_calibration_scale: float = 1.50 expected_event_day_shift: float = 0.0 expected_service_cost_hours: float = 2.0 expected_gain_margin: float = 10.0 expected_emergency_buffer_days: float = 6.0 emergency_operational_scale: float = 0.0 stale_risk_cutoff_days: float = 7.0 recent_gap_risk_factor: float = 0.75 stale_risk_factor: float = 0.10 minimum_scheduled_batteries: int = 0 maximum_scheduled_batteries: int | None = None route_last_leg_factor: float = 1.0 survivor_horizon_scale: float = 1.0 scheduled_fraction: float | None = None risk_quantile_spread_days: float = 0.0 selection_key: str = "expected_gain" class CompetitionPlanner(Planner): def __init__(self, model: EventTimeModel, policy: PlannerPolicy | None = None): self.model = model self.policy = policy or PlannerPolicy() def plan(self, battery_data, locations, travel_costs, settings): flat = battery_data.reset_index() start_time = pd.to_datetime(flat["end_time"], errors="coerce").max().normalize() daily = build_daily_features(battery_data) snapshot = scenario_snapshot( daily, locations, scenario_name="direct", scenario_time=start_time, ) return self.plan_snapshot(snapshot, locations, travel_costs, settings, start_time) def plan_snapshot( self, snapshot: pd.DataFrame, locations: pd.DataFrame, travel_costs: pd.DataFrame, settings, start_time: pd.Timestamp | str, predicted_rul: np.ndarray | None = None, predicted_risk: np.ndarray | None = None, predicted_survivor_rul: np.ndarray | None = None, ) -> pd.DataFrame: start = pd.Timestamp(start_time).normalize() horizon = int(settings.planning_window_days) predictions = ( self.model.predict_rul(snapshot) if predicted_rul is None else np.asarray(predicted_rul, dtype=float) ) risks = ( self.model.predict_event_risk(snapshot) if predicted_risk is None else np.asarray(predicted_risk, dtype=float) ) if "data_gap_days" in snapshot.columns: gaps = pd.to_numeric(snapshot["data_gap_days"], errors="coerce").to_numpy( dtype=float ) freshness_factor = np.ones(len(snapshot), dtype=float) recent_gap = (gaps > 0.0) & (gaps <= self.policy.stale_risk_cutoff_days) stale = (gaps > self.policy.stale_risk_cutoff_days) | ~np.isfinite(gaps) freshness_factor[recent_gap] = self.policy.recent_gap_risk_factor freshness_factor[stale] = self.policy.stale_risk_factor risks = risks * freshness_factor raw_predictions = predictions.copy() predictions = raw_predictions + self.policy.prediction_offset_days # Only a battery certain to fail is worth replacing at the 5% quantile. spread = self.policy.risk_quantile_spread_days if spread > 0.0: late = float(settings.late_replacement_penalty_daily) early = float(settings.early_replacement_penalty_daily) ratio = (late + early) / early certainty = np.clip(risks * self.policy.risk_calibration_scale, 1e-6, 1.0) optimal = np.clip(1.0 / (ratio * certainty), 1.0 / ratio, 1.0) predictions = predictions + spread * (optimal - 1.0 / ratio) work = snapshot[["battery", "building", "room"]].copy() work["predicted_rul"] = predictions work["predicted_risk"] = risks work["target_day"] = np.floor(np.clip(predictions, 0.0, horizon + 1.0)).astype(int) if self.policy.use_expected_cost: planned_day = np.clip(predictions, 0.0, float(horizon)) expected_event_day = np.clip( raw_predictions + self.policy.expected_event_day_shift, 0.0, float(horizon), ) calibrated_risk = np.clip( risks * self.policy.risk_calibration_scale, 0.0, 1.0, ) scheduled_event_cost = ( float(settings.early_replacement_penalty_daily) * np.maximum(expected_event_day - planned_day, 0.0) + float(settings.late_replacement_penalty_daily) * np.maximum(planned_day - expected_event_day, 0.0) ) emergency_day = float(horizon) + self.policy.expected_emergency_buffer_days no_schedule_event_cost = ( float(settings.late_replacement_penalty_daily) * np.maximum(emergency_day - expected_event_day, 0.0) + self.policy.expected_service_cost_hours + self._emergency_operational_cost(work, travel_costs, settings) ) survivor_rul = self._survivor_rul( snapshot, horizon, predicted_survivor_rul ) healthy_early_cost = ( float(settings.early_replacement_penalty_daily) * np.maximum(survivor_rul - planned_day, 0.0) ) expected_gain = ( calibrated_risk * (no_schedule_event_cost - scheduled_event_cost) - (1.0 - calibrated_risk) * healthy_early_cost - self.policy.expected_service_cost_hours ) work["expected_gain"] = expected_gain work["survivor_rul"] = survivor_rul work["scheduled"] = ( (expected_gain > self.policy.expected_gain_margin) & (predictions <= horizon) ) else: work["scheduled"] = ( (risks >= self.policy.event_risk_threshold) & (predictions <= horizon) ) # Keeps workload stable under calibration shift. priority_column = "expected_gain" if self.policy.use_expected_cost else "predicted_risk" key = self.policy.selection_key if key == "predicted_rul": ranked = work.sort_values(["predicted_rul", "predicted_risk", "battery"], ascending=[True, False, True]) elif key == "survivor_rul" and "survivor_rul" in work.columns: ranked = work.sort_values(["survivor_rul", "predicted_risk", "battery"], ascending=[True, False, True]) elif key == "risk_per_day" and "survivor_rul" in work.columns: # Value per unit of asset life given up, so near-misses outrank distant deaths. score = work["predicted_risk"] / np.maximum(work["survivor_rul"], 1.0) ranked = work.assign(_s=score).sort_values(["_s", "predicted_risk", "battery"], ascending=[False, False, True]) else: ranked = work.sort_values( [priority_column, "predicted_risk", "predicted_rul", "battery"], ascending=[False, False, True, True], ) minimum = min(max(int(self.policy.minimum_scheduled_batteries), 0), len(work)) maximum = self.policy.maximum_scheduled_batteries # Workload tracks fleet size rather than an absolute count. if self.policy.scheduled_fraction is not None: quota = int(round(self.policy.scheduled_fraction * len(work))) quota = max(quota, minimum) if maximum is not None: quota = min(quota, int(maximum)) minimum = maximum = quota selected = set(work.index[work["scheduled"]]) if len(selected) < minimum: selected.update(ranked.index[:minimum]) if maximum is not None and len(selected) > int(maximum): selected = set( ranked.loc[ranked.index.isin(selected)].index[: int(maximum)] ) work["scheduled"] = work.index.isin(selected) candidates = work[work["scheduled"]].copy() if not candidates.empty: candidates = self._schedule_candidates( candidates, travel_costs, settings, horizon, ) unscheduled_day = start + pd.Timedelta( days=horizon + self.policy.unscheduled_margin_days ) planned_rows: list[dict] = [] if not candidates.empty: for assigned_day, day_frame in candidates.groupby("assigned_day", sort=True): ordered = self._route_day( day_frame, travel_costs, settings.base_location, float(settings.overtime_penalty_factor) * self.policy.route_last_leg_factor, ) date = start + pd.Timedelta(days=int(assigned_day)) planned_rows.extend({"day": date, "battery": battery} for battery in ordered) scheduled_ids = set(candidates["battery"]) if not candidates.empty else set() remaining = work[~work["battery"].isin(scheduled_ids)].sort_values( ["building", "room", "battery"] ) planned_rows.extend( {"day": unscheduled_day, "battery": battery} for battery in remaining["battery"] ) return pd.DataFrame(planned_rows, columns=["day", "battery"]).reset_index(drop=True) def _emergency_operational_cost(self, work, travel_costs, settings) -> np.ndarray: """What the evaluator actually charges for a battery left to the emergency queue. A missed required battery becomes its own working day: a dedicated round trip from base, the swap itself, then straight home. Measured over 144 train cases that costs 65.56 beyond the late penalty, against the flat 2.0 the selection assumed - a 33x underestimate that hid the geography of a miss entirely. Scaling by distance is what lets a remote candidate outrank an equally risky one next door. """ scale = float(self.policy.emergency_operational_scale) if scale <= 0.0: return np.zeros(len(work), dtype=float) distances = travel_costs.set_index(["from", "to"])["hours"] base = settings.base_location buildings = work["building"].astype(str).to_numpy() out = np.zeros(len(work), dtype=float) cache: dict[str, float] = {} for position, building in enumerate(buildings): if building not in cache: try: leg = float(distances.loc[(base, building)]) except KeyError: leg = 0.0 hours = ( 2.0 * leg + float(settings.time_per_building_change_hours) + float(settings.time_per_room_change_hours) + float(settings.time_per_battery_hours) ) overtime = float(settings.overtime_penalty_factor) * max( hours - float(settings.overtime_start), 0.0 ) limits = 0.0 if hours > float(settings.worker_limit_daily_hours): limits += float(settings.worker_limit_daily_penalty) # each emergency day consumes this share of a week's budget limits += float(settings.worker_limit_weekly_penalty) * min( hours / float(settings.worker_limit_weekly_hours), 1.0 ) cache[building] = hours + overtime + limits out[position] = cache[building] return scale * out def _survivor_rul( self, snapshot: pd.DataFrame, horizon: int, override: np.ndarray | None, ) -> np.ndarray: """Remaining life if the battery outlives the window; drives the cost of swapping early.""" values = override if values is None and self.model is not None: # An artifact without the head must degrade, not crash the evaluator. predict = getattr(self.model, "predict_survivor_rul", None) values = predict(snapshot) if predict is not None else None if values is None: return np.full(len(snapshot), float(horizon), dtype=float) values = np.asarray(values, dtype=float) * self.policy.survivor_horizon_scale values = np.where(np.isfinite(values), values, float(horizon)) return np.maximum(values, float(horizon)) # Assigns days using route and capacity costs. def _schedule_candidates(self, candidates, travel_costs, settings, horizon): candidates = candidates.copy() distances = travel_costs.set_index(["from", "to"])["hours"] base = str(settings.base_location) base_room = str(settings.base_room) route_cache: dict[tuple[str, ...], tuple[float, float]] = {} last_leg_weight = ( float(settings.overtime_penalty_factor) * self.policy.route_last_leg_factor ) buildings = {index: str(row["building"]) for index, row in candidates.iterrows()} rooms = {index: str(row["room"]) for index, row in candidates.iterrows()} targets = {index: int(row["target_day"]) for index, row in candidates.iterrows()} priorities: dict[int, float] = {} for index, row in candidates.iterrows(): if ( "expected_gain" in candidates.columns and np.isfinite(float(row["expected_gain"])) ): priorities[index] = float(row["expected_gain"]) else: priorities[index] = float(row["predicted_risk"]) daily_limit = ( float(settings.worker_limit_daily_hours) * self.policy.capacity_daily_limit_fraction ) weekly_limit = ( float(settings.worker_limit_weekly_hours) * self.policy.capacity_weekly_limit_fraction ) hard_scale = self.policy.capacity_limit_penalty_multiplier operational_weight = self.policy.capacity_operational_cost_weight building_time = float(settings.time_per_building_change_hours) room_time = float(settings.time_per_room_change_hours) battery_time = float(settings.time_per_battery_hours) def route_legs(remote_buildings: frozenset[str]) -> tuple[float, float]: key = tuple(sorted(remote_buildings)) if not key: return 0.0, 0.0 cached = route_cache.get(key) if cached is None: _, cycle, last = self._order_buildings( list(key), distances, base, last_leg_weight ) cached = (cycle, last) route_cache[key] = cached return cached def hours_from_parts( remote_buildings: frozenset[str], remote_rooms: frozenset[tuple[str, str]], base_rooms: frozenset[str], count: int, ) -> tuple[float, float]: base_room_changes = len(base_rooms) - int(base_room in base_rooms) cycle, last = route_legs(remote_buildings) if not remote_buildings and count: cycle = last = float(distances.loc[(base, base)]) hours = ( cycle + len(remote_buildings) * building_time + (len(remote_rooms) + base_room_changes) * room_time + count * battery_time ) return hours, last def empty_state() -> dict: return { "indices": [], "remote": frozenset(), "remote_rooms": frozenset(), "base_rooms": frozenset(), "count": 0, "hours": 0.0, "ret": 0.0, } def build_state(indices: list[int]) -> dict: remote_buildings: set[str] = set() remote_rooms: set[tuple[str, str]] = set() base_rooms: set[str] = set() for index in indices: building = buildings[index] room = rooms[index] if building == base: base_rooms.add(room) else: remote_buildings.add(building) remote_rooms.add((building, room)) remote = frozenset(remote_buildings) remote_room_set = frozenset(remote_rooms) base_room_set = frozenset(base_rooms) hours, last = hours_from_parts( remote, remote_room_set, base_room_set, len(indices), ) return { "indices": list(indices), "remote": remote, "remote_rooms": remote_room_set, "base_rooms": base_room_set, "count": len(indices), "hours": hours, "ret": last, } def add_state(state: dict, index: int) -> dict: building = buildings[index] room = rooms[index] remote = state["remote"] remote_rooms = state["remote_rooms"] base_rooms = state["base_rooms"] if building == base: next_remote = remote next_remote_rooms = remote_rooms next_base_rooms = base_rooms | {room} else: next_remote = remote | {building} next_remote_rooms = remote_rooms | {(building, room)} next_base_rooms = base_rooms count = state["count"] + 1 hours, last = hours_from_parts( next_remote, next_remote_rooms, next_base_rooms, count, ) return { "indices": [*state["indices"], index], "remote": next_remote, "remote_rooms": next_remote_rooms, "base_rooms": next_base_rooms, "count": count, "hours": hours, "ret": last, } def day_cost(carry: float, hours: float) -> float: total = carry + hours overtime = max(total - float(settings.overtime_start), 0.0) daily_penalty = ( float(settings.worker_limit_daily_penalty) * hard_scale if total > daily_limit else 0.0 ) return ( operational_weight * hours + overtime * float(settings.overtime_penalty_factor) + daily_penalty ) def week_cost(hours: float) -> float: return ( float(settings.worker_limit_weekly_penalty) * hard_scale if hours > weekly_limit else 0.0 ) def timing_cost(indices: list[int], day: int) -> float: early_days = 0.0 late_days = 0.0 for index in indices: target = targets[index] if day < target: early_days += target - day else: late_days += day - target return ( early_days * float(settings.early_replacement_penalty_daily) + late_days * float(settings.late_replacement_penalty_daily) ) states: dict[int, dict] = {} week_hours: dict[int, float] = {} # The evaluator charges each day's return leg to the next working day too. def labor_with(overrides: dict[int, dict]) -> float: merged = {**states, **overrides} if overrides else states working = sorted(day for day, state in merged.items() if state["count"]) total = 0.0 carry = 0.0 for day in working: state = merged[day] total += day_cost(carry, state["hours"]) carry = state["ret"] return total def day_total(day: int) -> float: working = sorted(other for other, state in states.items() if state["count"]) carry = 0.0 for other in working: if other == day: return carry + states[other]["hours"] carry = states[other]["ret"] return states[day]["hours"] if day in states else 0.0 base_labor = 0.0 windows: dict[int, range] = {} for index in candidates.index: target = targets[index] windows[index] = range( max(0, target - self.policy.capacity_lookback_days), min(horizon, target + self.policy.capacity_lookahead_days) + 1, ) def insertion_option(index: int, day: int) -> tuple[float, dict]: old_state = states.get(day, empty_state()) new_state = add_state(old_state, index) week = day // 7 old_week_hours = week_hours.get(week, 0.0) new_week_hours = ( old_week_hours - old_state["hours"] + new_state["hours"] ) cost = ( timing_cost([index], day) + labor_with({day: new_state}) - base_labor + week_cost(new_week_hours) - week_cost(old_week_hours) ) return cost, new_state # Inserts urgent jobs by regret. pending = sorted( candidates.index, key=lambda index: ( targets[index], -priorities[index], str(candidates.at[index, "battery"]), ), ) while pending: batch = pending[: min(8, len(pending))] chosen: tuple | None = None for index in batch: options = [] for day in windows[index]: cost, new_state = insertion_option(index, day) options.append( ( cost, abs(day - targets[index]), day, new_state, ) ) options.sort(key=lambda option: (option[0], option[1], option[2])) best = options[0] second_cost = options[1][0] if len(options) > 1 else best[0] key = ( second_cost - best[0], priorities[index], -targets[index], str(candidates.at[index, "battery"]), ) if chosen is None or key > chosen[0]: chosen = (key, index, best) assert chosen is not None _, index, best = chosen _, _, day, new_state = best old_state = states.get(day, empty_state()) week = day // 7 week_hours[week] = ( week_hours.get(week, 0.0) - old_state["hours"] + new_state["hours"] ) states[day] = new_state base_labor = labor_with({}) candidates.at[index, "assigned_day"] = day pending.remove(index) def relocate_delta( indices: list[int], source: int, target: int, ) -> tuple[float, dict | None, dict | None]: if source == target: return 0.0, None, None source_state = states.get(source) if source_state is None: return float("inf"), None, None moving = set(indices) if not moving.issubset(source_state["indices"]): return float("inf"), None, None source_after = build_state( [index for index in source_state["indices"] if index not in moving] ) target_state = states.get(target, empty_state()) target_after = build_state([*target_state["indices"], *indices]) affected_weeks = {source // 7, target // 7} old_week_cost = sum( week_cost(week_hours.get(week, 0.0)) for week in affected_weeks ) new_week_cost = 0.0 for week in affected_weeks: hours = week_hours.get(week, 0.0) if source // 7 == week: hours += source_after["hours"] - source_state["hours"] if target // 7 == week: hours += target_after["hours"] - target_state["hours"] new_week_cost += week_cost(hours) delta = ( labor_with({source: source_after, target: target_after}) - base_labor + new_week_cost - old_week_cost + timing_cost(indices, target) - timing_cost(indices, source) ) return delta, source_after, target_after def apply_move( indices: list[int], source: int, target: int, source_after: dict, target_after: dict, ) -> None: source_state = states[source] target_state = states.get(target, empty_state()) source_week = source // 7 target_week = target // 7 week_hours[source_week] = ( week_hours.get(source_week, 0.0) + source_after["hours"] - source_state["hours"] ) if target_week == source_week: week_hours[source_week] += ( target_after["hours"] - target_state["hours"] ) else: week_hours[target_week] = ( week_hours.get(target_week, 0.0) + target_after["hours"] - target_state["hours"] ) if source_after["count"]: states[source] = source_after else: states.pop(source, None) states[target] = target_after candidates.loc[indices, "assigned_day"] = target nonlocal base_labor base_labor = labor_with({}) # Repairs overloaded or fragmented days. for _ in range(2): improved = False days = sorted( states, key=lambda day: ( day_total(day) <= daily_limit, day, ), ) units: list[tuple[list[int], int]] = [] for day in days: indices = states[day]["indices"] groups: dict[tuple[str, str], list[int]] = {} for index in indices: groups.setdefault((buildings[index], rooms[index]), []).append(index) units.extend( (group, day) for group in groups.values() if len(group) > 1 ) units.extend(([index], day) for index in indices) for indices, source in units: if source not in states: continue if not set(indices).issubset(states[source]["indices"]): continue first = max( 0, min( targets[index] - self.policy.capacity_lookback_days for index in indices ), ) last = min( horizon, max( targets[index] + self.policy.capacity_lookahead_days for index in indices ), ) best_delta = 0.0 best_day = source best_source_state = None best_target_state = None for target in range(first, last + 1): if target == source: continue delta, source_after, target_after = relocate_delta( indices, source, target, ) option = (delta, abs(target - source), target) current = ( best_delta, abs(best_day - source), best_day, ) if option < current: best_delta = delta best_day = target best_source_state = source_after best_target_state = target_after if ( best_day != source and best_delta < -1e-9 and best_source_state is not None and best_target_state is not None ): apply_move( indices, source, best_day, best_source_state, best_target_state, ) improved = True if not improved: break candidates["assigned_day"] = candidates["assigned_day"].astype(int) return candidates @staticmethod def _order_buildings( remote: list[str], distances: pd.Series, base: str, last_leg_weight: float = 0.0, ) -> tuple[list[str], float, float]: """Returns (order, cycle_hours, return_leg).""" remote = sorted(remote) if not remote: return [], 0.0, 0.0 def legs(order: list[str]) -> tuple[float, float]: points = [base, *order, base] total = sum( float(distances.loc[(left, right)]) for left, right in zip(points, points[1:]) ) return total, float(distances.loc[(order[-1], base)]) # The final leg is also charged to the next working day's hour budget. def objective(order: list[str]) -> float: total, last = legs(order) return total + last_leg_weight * last if len(remote) <= 9: order = CompetitionPlanner._exact_building_route( remote, distances, base, last_leg_weight, ) else: remaining = set(remote) current = base order = [] while remaining: next_building = min( remaining, key=lambda building: ( float(distances.loc[(current, building)]), building, ), ) order.append(next_building) remaining.remove(next_building) current = next_building improved = True best_cost = objective(order) while improved and len(order) >= 3: improved = False for left in range(len(order) - 1): for right in range(left + 1, len(order)): candidate = ( order[:left] + list(reversed(order[left : right + 1])) + order[right + 1 :] ) candidate_cost = objective(candidate) if candidate_cost + 1e-12 < best_cost: order = candidate best_cost = candidate_cost improved = True cycle, last = legs(order) return order, cycle, last @staticmethod def _route_day( day_frame: pd.DataFrame, travel_costs: pd.DataFrame, base: str, last_leg_weight: float = 0.0, ) -> list[str]: distances = travel_costs.set_index(["from", "to"])["hours"] buildings = set(day_frame["building"].astype(str)) has_base = base in buildings building_order, _, _ = CompetitionPlanner._order_buildings( sorted(buildings - {base}), distances, base, last_leg_weight, ) if has_base: building_order = [base, *building_order] batteries: list[str] = [] for building in building_order: local = day_frame[day_frame["building"].astype(str) == building].sort_values( ["room", "target_day", "battery"] ) batteries.extend(local["battery"].astype(str).tolist()) return batteries @staticmethod def _exact_building_route( buildings: list[str], distances: pd.Series, base: str, last_leg_weight: float = 0.0, ) -> list[str]: if len(buildings) <= 1: return buildings.copy() count = len(buildings) states: dict[tuple[int, int], tuple[float, tuple[str, ...]]] = {} for index, building in enumerate(buildings): states[(1 << index, index)] = ( float(distances.loc[(base, building)]), (building,), ) for mask in range(1, 1 << count): for last in range(count): state = states.get((mask, last)) if state is None: continue cost, path = state for nxt in range(count): bit = 1 << nxt if mask & bit: continue new_mask = mask | bit candidate = ( cost + float(distances.loc[(buildings[last], buildings[nxt])]), (*path, buildings[nxt]), ) key = (new_mask, nxt) current = states.get(key) if current is None or candidate < current: states[key] = candidate full_mask = (1 << count) - 1 best: tuple[float, tuple[str, ...]] | None = None for last in range(count): cost, path = states[(full_mask, last)] candidate = ( cost + (1.0 + last_leg_weight) * float(distances.loc[(buildings[last], base)]), path, ) if best is None or candidate < best: best = candidate assert best is not None return list(best[1])