Swap OrderedPlanner heuristic for CP-SAT MILP planner

batteryswap_example/planners/best.pickle

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:098ee847de3b9cbe57b2ccdadc65213625cd06585bd2fcb492bb8f8eed56693c
-size 201765
+oid sha256:d298671d2c8f3c4374b5f2d75edae3e7eceff08985f857bd406ebc9e3f450df4
+size 6525860

batteryswap_example/train.py

@@ -1,14 +1,21 @@
+
+import pandas
+import numpy
 import pickle
-import os
-from pathlib import Path
 from typing import Optional, Sequence
+from pathlib import Path
+import pathlib
+import os
 
-import pandas
+
+from sklearn.ensemble import ExtraTreesRegressor
 import pandas as pd
+import numpy
 import numpy as np
+import pandas
+from pydantic import Field
 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
@@ -16,1067 +23,323 @@ 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)}")
+class OrderedPlanner(Planner):
+    def __init__(self, rul_estimator):
+        self.rul_estimator = rul_estimator
 
-    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 plan(self, battery_data, locations, travel_costs, settings):
+
+        # Remaining Useful Life estimation
+        # FIXME: consider balance of over/under-estimation
+        percentile = 'p50'
+        rul = self.rul_estimator.predict(battery_data)
+        rul_days = rul[percentile]
+        # convert to a date
+        start_time = battery_data.reset_index()['end_time'].max().normalize()
+        predict_eol = start_time + pandas.to_timedelta(rul_days, unit='D')
+
+        loc = locations.copy().set_index('battery')
+        loc['eol_time'] = predict_eol
+        order = loc.sort_values('eol_time', ascending=True)
+
+        # Planner
+        # Stupid heuristic: Do one swap per day
+        # FIXME: take travel distances into account
+        # FIXME: take co-location into account
+        # FIXME: take daily and weekly limits into account
+        days = start_time + pandas.to_timedelta(numpy.arange(len(order)), unit='D')
+        plan = pandas.DataFrame({
+            'day': days,
+            'battery': order.index,
+        })
+
+        check_plan_valid(plan, locations, start_time=start_time)
 
+        return plan
 
-def _setting(settings, name, default):
-    if isinstance(settings, dict):
-        return settings.get(name, default)
-    return getattr(settings, name, default)
 
+class DummyRULModel(RULModel):
+    # RUL model that predicts (no information rate)
+    # FIXME: make a model that actually uses the data to improve predictions
 
-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,
+        time_col: str = 'end_time',
+        group_col: str = 'device_id',
+        value_cols: Sequence[str] = ('voltage', 'temperature'),
+        quantiles: Sequence[float] = (0.5, ),
     ):
-        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)
-                )
+        self.time_col = time_col
+        self.group_col = group_col
+        self.value_cols = list(value_cols)
+        self.quantiles = sorted(quantiles)
+        self.quantile_cols = [f"p{round(q * 100):02d}" for q in self.quantiles]
 
-        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
-            )
+        self.model = None
+        self.use_total_elapsed_days = True
 
-        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
+    def _compute_features(self, unit_df: pd.DataFrame) -> np.ndarray:
+        unit_df = unit_df.sort_index(level=self.time_col).copy()
+        if len(unit_df) == 0:
+            raise ValueError("Received empty battery time series")
 
-        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])
+        all_stats = {}
+
+        # Convert timestamps to elapsed days
+        times = pd.to_datetime(
+            unit_df.index.get_level_values(self.time_col)
             )
-            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
+        elapsed_days = (
+            times - times[0]
+            ).total_seconds().to_numpy() / 86400.0
+        
+
+        # General history features
+        all_stats["n_obs"] = float(len(unit_df))
+        all_stats["history_days"] = (
+            float(elapsed_days[-1]) if len(elapsed_days) > 1 else 0.0
             )
-            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,
+        # Features for voltage and temperature
+        for col in self.value_cols:
+            values = pd.to_numeric(
+                unit_df[col],
+                errors="coerce"
                 )
-                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)
 
+            valid = values.notna()
+
+            if valid.sum() == 0:
+                all_stats[f"{col}_latest"] = 0.0
+                all_stats[f"{col}_mean"] = 0.0
+                all_stats[f"{col}_std"] = 0.0
+                all_stats[f"{col}_min"] = 0.0
+                all_stats[f"{col}_max"] = 0.0
+                all_stats[f"{col}_change"] = 0.0
+                all_stats[f"{col}_slope"] = 0.0
+                continue
+
+            x = values[valid].to_numpy(dtype=float)
+            t = elapsed_days[valid.to_numpy()]
+
+            all_stats[f"{col}_latest"] = float(x[-1])
+            all_stats[f"{col}_mean"] = float(np.mean(x))
+            all_stats[f"{col}_std"] = float(np.std(x))
+            all_stats[f"{col}_min"] = float(np.min(x))
+            all_stats[f"{col}_max"] = float(np.max(x))
+            all_stats[f"{col}_change"] = float(x[-1] - x[0])
+            
+            # Recent-window features
+            for window_days in (7, 14, 30):
+                cutoff = t[-1] - window_days
+                recent_mask = t >= cutoff
+
+                recent_x = x[recent_mask]
+                recent_t = t[recent_mask]
+
+                prefix = f"{col}_{window_days}d"
+
+                if len(recent_x) > 0:
+                    all_stats[f"{prefix}_mean"] = float(np.mean(recent_x))
+                    all_stats[f"{prefix}_std"] = float(np.std(recent_x))
+                    all_stats[f"{prefix}_min"] = float(np.min(recent_x))
+                    all_stats[f"{prefix}_max"] = float(np.max(recent_x))
+                    all_stats[f"{prefix}_change"] = float(
+                        recent_x[-1] - recent_x[0]
+                        )
+
+                    if len(recent_x) >= 2 and np.ptp(recent_t) > 0:
+                        recent_slope = np.polyfit(
+                            recent_t,
+                            recent_x,
+                            1
+                            )[0]
+                    else:
+                        recent_slope = 0.0
+
+                    all_stats[f"{prefix}_slope"] = float(recent_slope)
+                else:
+                    all_stats[f"{prefix}_mean"] = 0.0
+                    all_stats[f"{prefix}_std"] = 0.0
+                    all_stats[f"{prefix}_min"] = 0.0
+                    all_stats[f"{prefix}_max"] = 0.0
+                    all_stats[f"{prefix}_change"] = 0.0
+                    all_stats[f"{prefix}_slope"] = 0.0
+            
+            # Recent level compared with overall level
+            if len(x) > 0:
+                recent_7_mask = t >= (t[-1] - 7)
+                recent_7 = x[recent_7_mask]
+
+                if len(recent_7) > 0:
+                    all_stats[f"{col}_recent7_vs_mean"] = float(
+                        np.mean(recent_7) - np.mean(x)
+                        )
+                else:
+                    all_stats[f"{col}_recent7_vs_mean"] = 0.0
+                    
+            # Trend per day
+            if len(x) >= 2 and np.ptp(t) > 0:
+                slope = np.polyfit(t, x, 1)[0]
+            else:
+                slope = 0.0
 
-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)
+            all_stats[f"{col}_slope"] = float(slope)
 
+        feature_names = sorted(all_stats.keys())
+        self._feature_names_ = feature_names
 
-def _rul_finite(value):
-    return float(value) if np.isfinite(value) else 0.0
+        return np.array(
+            [all_stats[k] for k in feature_names],
+            dtype=float,
+            )
 
+    def _build_feature_matrix(self, timeseries: pd.DataFrame) -> tuple[np.ndarray, list]:
+        rows, ids = [], []
+        for unit_id, unit_df in timeseries.groupby(
+            self.group_col,
+            observed=True
+            ):
+            if len(unit_df) < 2:
+                continue
+            rows.append(self._compute_features(unit_df))
+            ids.append(unit_id)
+        return np.vstack(rows), ids
+
+    def fit(self, timeseries: pd.DataFrame, rul: pd.Series):
+        X, ids = self._build_feature_matrix(timeseries)
+        y = np.array([rul[unit_id] for unit_id in ids])
+
+        # FIXME: actually use an estimator that learns
+        self.model = ExtraTreesRegressor(
+            n_estimators=300,
+            min_samples_leaf=2,
+            random_state=42,
+            n_jobs=-1,
+            )
+        self.model.fit(X, y)
+        return self
 
-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))
+    def predict(self, timeseries: pd.DataFrame) -> pd.DataFrame:
+        rows, ids = [], []
+        for unit_id, unit_df in timeseries.groupby(
+            self.group_col,
+            observed=True
+            ):
+            rows.append(self._compute_features(unit_df))
+            ids.append(unit_id)
 
-    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])
-            )
+        X = np.vstack(rows)
 
-        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])
-                )
+        # every quantile column just gets the single point prediction.
+        point_pred = self.model.predict(X)
+        preds = {col: point_pred for col in self.quantile_cols}
 
-        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
-    )
+        out = pd.DataFrame(preds, index=pd.Index(ids, name=self.group_col))
+        return out[self.quantile_cols]
 
 
-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 train_rul_model(locations, timeseries, eol_times, scenarios, limit_scenarios=None):
 
-    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)
+    # Collect training data
+    # FIXME: train/validate/test split to estimate generalized predictive performance
+    gen = iterate_scenarios(locations, timeseries, eol_times, scenarios)
+    cut_dfs = []
+    cut_ruls = []
+    scenarios_loaded = 0
+    for scenario, locs, cut, eol in gen:
+        print('load-scenario')
+        cut = cut.reset_index()
+        cut['device_id'] = cut['device_id'].astype(str) + scenario['name']
+        cut = cut.set_index(['device_id', 'end_time'])
+        cut_dfs.append(cut)
+        plan_start = pandas.Timestamp(scenario['start_time'])
+        unobserved_rul = 120
 
+        # Convert EOL datetimes to RUL in days relative to planning time
+        rul_days = (eol - plan_start) / pandas.Timedelta(days=1)
+        rul_days.index = pandas.Series(rul_days.index) + scenario['name']
+        rul_days = rul_days.fillna(unobserved_rul)
+        cut_ruls.append(rul_days)
 
-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
+        assert set(rul_days.index) == set(cut.index.get_level_values('device_id'))
 
-    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]
+        if limit_scenarios is not None:
+            if scenarios_loaded > limit_scenarios:
+                break
+        scenarios_loaded += 1
 
+        
+    # Train a RUL prediction model
+    rul_model = DummyRULModel()
+    X = pandas.concat(cut_dfs)
+    Y = pandas.concat(cut_ruls)
 
-def build_training_snapshots(
-    locations, timeseries, eol_times, scenarios, limit_scenarios=None
-):
-    feature_parts = []
-    duration_parts = []
-    event_parts = []
+    print(X.head())
+    print(Y.head())
 
-    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),
-    )
+    rul_model.fit(X, Y)
 
+    # FIXME: do model selection
 
-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)
+    return rul_model
 
 
 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
-
+    split : str = 'train'
 
 def main():
     cfg = Config()
 
     if cfg.dataset_path is None:
-        dataset_path = os.environ.get("BATTERYSWAP_DATASET_PATH", 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)
+    locations, timeseries, eol_times, scenarios  = load_dataset(split_path)
 
-    log.info("evaluate-load-data", path=dataset_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)
+    # Prediction model training
+    #rul_model = DummyRULModel()
+    rul_model = train_rul_model(locations, timeseries, eol_times, scenarios, limit_scenarios=1)
+    log.info('train-done')
+
+    log.info('evaluate')
+    # Evaluate on the planning scenarios
+    gen = iterate_scenarios(locations, timeseries, eol_times, 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,
-        )
+        scenario_name = scenario['name']
+        travel_costs = scenario['travel_costs']
+        settings = scenario['settings']
+
+        planner = OrderedPlanner(rul_model)
         plan = planner.plan(cut, locs, travel_costs, settings)
 
-        start_time = pandas.Timestamp(scenario["start_time"])
+        start_time = pandas.Timestamp(scenario['start_time'])
 
-        transitions, daily, overall = evaluate_plan(
-            plan, locs, travel_costs, settings, eol_times=eol, start_time=start_time
-        )
+        transitions, daily, overall = evaluate_plan(plan, locs, travel_costs, settings, eol_times=eol, start_time=start_time)
 
-        print("scores", scenario_name, overall)
+        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 = OrderedPlanner(rul_model)
 
     planner_path = 'batteryswap_example/planners/best.pickle' 
     with open(planner_path, "wb") as f: