tests/test_competition_planner.py

from pathlib import Path
from types import SimpleNamespace

import pytest
import numpy as np
import pandas as pd
from batteryswap_public.evaluate import EvaluationSettings, check_plan_valid

from batteryswapai.competition_planner import CompetitionPlanner, PlannerPolicy


class DummyEventTimeModel:
    def predict_event_risk(self, snapshot):
        return np.array([0.9, 0.8, 0.01])

    def predict_rul(self, snapshot):
        return np.array([20.0, 22.0, 5.0])


def _inputs():
    locations = pd.DataFrame(
        {
            "battery": ["d_a", "d_b", "d_c"],
            "building": ["b_remote", "b_remote", "b_base"],
            "room": ["r_1", "r_1", "r_base"],
            "start_time": pd.to_datetime(["2025-01-01"] * 3),
            "end_time": pd.to_datetime(["2026-12-01"] * 3),
        }
    )
    travel = pd.DataFrame(
        [
            {"from": "b_base", "to": "b_base", "hours": 0.0},
            {"from": "b_base", "to": "b_remote", "hours": 1.0},
            {"from": "b_remote", "to": "b_base", "hours": 1.0},
            {"from": "b_remote", "to": "b_remote", "hours": 0.0},
        ]
    )
    settings = EvaluationSettings(base_location="b_base", base_room="r_base")
    return locations, travel, settings


def _scheduler_settings(**overrides):
    values = {
        "base_location": "b_base",
        "base_room": "r_base",
        "worker_limit_daily_hours": 4.0,
        "worker_limit_weekly_hours": 100.0,
        "worker_limit_daily_penalty": 100.0,
        "worker_limit_weekly_penalty": 100.0,
        "overtime_start": 4.0,
        "overtime_penalty_factor": 10.0,
        "time_per_building_change_hours": 0.0,
        "time_per_room_change_hours": 0.0,
        "time_per_battery_hours": 1.0,
        "early_replacement_penalty_daily": 0.5,
        "late_replacement_penalty_daily": 2.0,
    }
    values.update(overrides)
    return SimpleNamespace(**values)


def test_plan_is_complete_and_batches_same_building():
    locations, travel, settings = _inputs()
    planner = CompetitionPlanner(
        DummyEventTimeModel(),
        PlannerPolicy(
            event_risk_threshold=0.1,
            prediction_offset_days=0.0,
            building_batch_window_days=3,
        ),
    )
    plan = planner.plan_snapshot(
        locations,
        locations,
        travel,
        settings,
        "2026-01-01",
    )
    check_plan_valid(plan, locations, start_time=pd.Timestamp("2026-01-01"))
    days = plan.set_index("battery")["day"]
    assert days["d_a"] == days["d_b"]
    assert days["d_c"] > pd.Timestamp("2026-02-12")


def test_prediction_arrays_can_be_supplied_for_oof_validation():
    locations, travel, settings = _inputs()
    planner = CompetitionPlanner(DummyEventTimeModel())
    plan = planner.plan_snapshot(
        locations,
        locations,
        travel,
        settings,
        "2026-01-01",
        predicted_rul=np.array([5.0, 6.0, 7.0]),
        predicted_risk=np.array([0.9, 0.01, 0.01]),
    )
    inside = plan[plan["day"] <= pd.Timestamp("2026-02-12")]
    assert inside["battery"].tolist() == ["d_a"]


def test_stale_sensor_history_reduces_proactive_risk():
    locations, travel, settings = _inputs()
    snapshot = locations.copy()
    snapshot["data_gap_days"] = [0.0, 8.0, 0.0]
    planner = CompetitionPlanner(
        DummyEventTimeModel(),
        PlannerPolicy(event_risk_threshold=0.5, prediction_offset_days=0.0),
    )

    plan = planner.plan_snapshot(
        snapshot,
        locations,
        travel,
        settings,
        "2026-01-01",
        predicted_rul=np.array([5.0, 5.0, 5.0]),
        predicted_risk=np.array([0.6, 0.6, 0.01]),
    )

    inside = plan[plan["day"] <= pd.Timestamp("2026-02-12")]
    assert inside["battery"].tolist() == ["d_a"]


def test_schedule_quota_uses_risk_ranking_when_threshold_drifts():
    locations, travel, settings = _inputs()
    planner = CompetitionPlanner(
        DummyEventTimeModel(),
        PlannerPolicy(
            event_risk_threshold=0.99,
            prediction_offset_days=0.0,
            minimum_scheduled_batteries=2,
            maximum_scheduled_batteries=2,
        ),
    )

    plan = planner.plan_snapshot(
        locations,
        locations,
        travel,
        settings,
        "2026-01-01",
        predicted_rul=np.array([20.0, 22.0, 5.0]),
        predicted_risk=np.array([0.9, 0.8, 0.01]),
    )

    inside = plan[plan["day"] <= pd.Timestamp("2026-02-12")]
    assert set(inside["battery"]) == {"d_a", "d_b"}


def test_capacity_balancing_accounts_for_marginal_route_cost():
    locations = pd.DataFrame(
        {
            "battery": ["d_a", "d_b"],
            "building": ["b_a", "b_b"],
            "room": ["r_a", "r_b"],
            "start_time": pd.to_datetime(["2025-01-01"] * 2),
            "end_time": pd.to_datetime(["2026-12-01"] * 2),
        }
    )
    travel = pd.DataFrame(
        [
            {
                "from": left,
                "to": right,
                "hours": 0.0 if left == right else 0.1 if "b_base" not in (left, right) else 1.0,
            }
            for left in ("b_base", "b_a", "b_b")
            for right in ("b_base", "b_a", "b_b")
        ]
    )
    settings = EvaluationSettings(base_location="b_base", base_room="r_base")
    planner = CompetitionPlanner(
        DummyEventTimeModel(),
        PlannerPolicy(
            prediction_offset_days=0.0,
            building_batch_window_days=0,
        ),
    )

    plan = planner.plan_snapshot(
        locations,
        locations,
        travel,
        settings,
        "2026-01-01",
        predicted_rul=np.array([1.0, 2.0]),
        predicted_risk=np.array([0.9, 0.9]),
    )

    days = plan.set_index("battery")["day"]
    assert days["d_a"] == days["d_b"]


def test_scheduler_splits_visit_before_daily_limit_penalty():
    candidates = pd.DataFrame(
        {
            "battery": [f"d_{index}" for index in range(6)],
            "building": ["b_remote"] * 6,
            "room": ["r_1"] * 6,
            "predicted_rul": [0.0] * 6,
            "predicted_risk": [0.9] * 6,
            "target_day": [0] * 6,
        }
    )
    travel = pd.DataFrame(
        [
            {"from": left, "to": right, "hours": 0.0}
            for left in ("b_base", "b_remote")
            for right in ("b_base", "b_remote")
        ]
    )
    planner = CompetitionPlanner(
        DummyEventTimeModel(),
        PlannerPolicy(capacity_lookahead_days=7),
    )
    scheduled = planner._schedule_candidates(
        candidates,
        travel,
        _scheduler_settings(),
        horizon=7,
    )

    counts = scheduled.groupby("assigned_day").size()
    assert counts.max() <= 4
    assert len(counts) >= 2


def test_scheduler_batches_trip_when_savings_exceed_early_cost():
    candidates = pd.DataFrame(
        {
            "battery": ["d_a", "d_b"],
            "building": ["b_remote", "b_remote"],
            "room": ["r_1", "r_1"],
            "predicted_rul": [0.0, 1.0],
            "predicted_risk": [0.9, 0.9],
            "target_day": [0, 1],
        }
    )
    travel = pd.DataFrame(
        [
            {"from": "b_base", "to": "b_base", "hours": 0.0},
            {"from": "b_base", "to": "b_remote", "hours": 2.0},
            {"from": "b_remote", "to": "b_base", "hours": 2.0},
            {"from": "b_remote", "to": "b_remote", "hours": 0.0},
        ]
    )
    planner = CompetitionPlanner(
        DummyEventTimeModel(),
        PlannerPolicy(capacity_lookahead_days=7),
    )
    scheduled = planner._schedule_candidates(
        candidates,
        travel,
        _scheduler_settings(
            worker_limit_daily_hours=20.0,
            overtime_start=20.0,
            time_per_battery_hours=0.25,
        ),
        horizon=7,
    )

    days = scheduled.set_index("battery")["assigned_day"]
    assert days["d_a"] == days["d_b"] == 0


def test_exact_building_route_beats_greedy_trap():
    buildings = ["b_a", "b_b", "b_c"]
    nodes = ["b_base", *buildings]
    values = {}
    for left in nodes:
        for right in nodes:
            values[(left, right)] = 0.0 if left == right else 10.0
    values[("b_base", "b_a")] = 1.0
    values[("b_a", "b_c")] = 1.0
    values[("b_c", "b_b")] = 1.0
    values[("b_b", "b_base")] = 1.0
    distances = pd.Series(values)

    route = CompetitionPlanner._exact_building_route(
        buildings,
        distances,
        "b_base",
    )

    assert route == ["b_a", "b_c", "b_b"]


def test_official_planner_interface_has_a_working_fallback():
    locations, travel, settings = _inputs()
    battery_data = pd.DataFrame(
        {
            "device_id": ["d_a", "d_b", "d_c"],
            "end_time": pd.to_datetime(["2026-01-01"] * 3),
            "voltage": [2.7, 2.8, 3.0],
            "temperature": [20.0, 20.0, 20.0],
        }
    ).set_index(["device_id", "end_time"])
    planner = CompetitionPlanner(
        DummyEventTimeModel(),
        PlannerPolicy(prediction_offset_days=0.0),
    )
    plan = planner.plan(battery_data, locations, travel, settings)
    check_plan_valid(plan, locations, start_time=pd.Timestamp("2026-01-01"))


def test_return_leg_is_charged_to_the_next_working_day():
    """Characterises the evaluator rule the capacity model exists to track."""

    from batteryswap_public.evaluate import evaluate_plan

    travel = pd.DataFrame(
        [
            ("b_base", "b_base", 0.0),
            ("b_base", "b_far", 5.0), ("b_far", "b_base", 5.0),
            ("b_base", "b_near", 0.1), ("b_near", "b_base", 0.1),
            ("b_far", "b_near", 5.0), ("b_near", "b_far", 5.0),
            ("b_far", "b_far", 0.0), ("b_near", "b_near", 0.0),
        ],
        columns=["from", "to", "hours"],
    )
    near = [f"d_near_{index}" for index in range(30)]
    batteries = ["d_far", *near]
    locations = pd.DataFrame(
        {
            "battery": batteries,
            "building": ["b_far"] + ["b_near"] * len(near),
            "room": ["r_far"] + ["r_near"] * len(near),
            "start_time": pd.to_datetime(["2025-01-01"] * len(batteries)),
            "end_time": pd.to_datetime(["2026-12-01"] * len(batteries)),
        }
    )
    settings = EvaluationSettings(base_location="b_base", base_room="r_base")
    start_time = pd.Timestamp("2025-09-01")
    eol = pd.Series({battery: pd.NaT for battery in batteries})

    def overtime(first_day_batteries, second_day_batteries):
        rows = [(1, battery) for battery in first_day_batteries]
        rows += [(2, battery) for battery in second_day_batteries]
        plan = pd.DataFrame(
            {
                "day": [start_time + pd.Timedelta(days=day) for day, _ in rows],
                "battery": [battery for _, battery in rows],
            }
        )
        _, _, scores = evaluate_plan(
            plan, locations, travel, settings,
            eol_times=eol, start_time=start_time, verbose=0,
        )
        return float(scores.overtime)

    far_first = overtime(["d_far"], near)
    far_last = overtime(near, ["d_far"])
    # Both days run into overtime, so the extra return leg is billed twice.
    assert far_first - far_last == pytest.approx(
        settings.overtime_penalty_factor * (5.0 - 0.1), abs=1e-6
    )


def test_route_finishes_at_the_building_nearest_base():
    distances = pd.DataFrame(
        [
            ("b_base", "b_base", 0.0),
            ("b_base", "b_a", 1.0), ("b_a", "b_base", 0.2),
            ("b_base", "b_b", 1.5), ("b_b", "b_base", 1.0),
            ("b_a", "b_b", 0.5), ("b_b", "b_a", 0.8),
            ("b_a", "b_a", 0.0), ("b_b", "b_b", 0.0),
        ],
        columns=["from", "to", "hours"],
    ).set_index(["from", "to"])["hours"]

    # Both orders cost 2.5 in travel; only the final leg differs.
    plain, plain_cycle, plain_last = CompetitionPlanner._order_buildings(
        ["b_a", "b_b"], distances, "b_base", last_leg_weight=0.0
    )
    weighted, weighted_cycle, weighted_last = CompetitionPlanner._order_buildings(
        ["b_a", "b_b"], distances, "b_base", last_leg_weight=2.0
    )
    assert plain_cycle == pytest.approx(weighted_cycle) == pytest.approx(2.5)
    assert plain[-1] == "b_b" and plain_last == pytest.approx(1.0)
    assert weighted[-1] == "b_a" and weighted_last == pytest.approx(0.2)


def test_scheduler_merges_far_trips_that_would_breach_the_carried_daily_limit():
    from batteryswap_public.evaluate import evaluate_plan

    travel = pd.DataFrame(
        [
            ("b_base", "b_base", 0.0),
            ("b_base", "b_f1", 9.0), ("b_f1", "b_base", 9.0),
            ("b_base", "b_f2", 8.0), ("b_f2", "b_base", 8.0),
            ("b_f1", "b_f2", 0.5), ("b_f2", "b_f1", 0.5),
            ("b_f1", "b_f1", 0.0), ("b_f2", "b_f2", 0.0),
        ],
        columns=["from", "to", "hours"],
    )
    locations = pd.DataFrame(
        {
            "battery": ["d_1", "d_2"],
            "building": ["b_f1", "b_f2"],
            "room": ["r_1", "r_2"],
            "start_time": pd.to_datetime(["2025-01-01"] * 2),
            "end_time": pd.to_datetime(["2026-12-01"] * 2),
        }
    )
    settings = EvaluationSettings(base_location="b_base", base_room="r_base")
    snapshot = locations.assign(data_gap_days=0.0)
    planner = CompetitionPlanner(None, PlannerPolicy(event_risk_threshold=0.1))
    plan = planner.plan_snapshot(
        snapshot,
        locations,
        travel,
        settings,
        "2025-09-01",
        predicted_rul=np.array([12.0, 17.0]),
        predicted_risk=np.array([0.9, 0.9]),
    )
    scheduled = plan[plan["day"] <= pd.Timestamp("2025-09-01") + pd.Timedelta(days=42)]
    assert len(scheduled) == 2
    assert scheduled["day"].nunique() == 1

    _, _, scores = evaluate_plan(
        plan,
        locations,
        travel,
        settings,
        eol_times=pd.Series({"d_1": pd.NaT, "d_2": pd.NaT}),
        start_time=pd.Timestamp("2025-09-01"),
        verbose=0,
    )
    assert float(scores.daily_limit) == 0.0


def test_survivor_horizon_defaults_to_the_window_and_never_falls_below_it():
    planner = CompetitionPlanner(None, PlannerPolicy(survivor_horizon_scale=0.1))
    snapshot = pd.DataFrame({"battery": ["d_a", "d_b"]})
    assert list(planner._survivor_rul(snapshot, 42, None)) == [42.0, 42.0]
    values = planner._survivor_rul(snapshot, 42, np.array([500.0, np.nan]))
    assert values[0] == pytest.approx(50.0)
    assert values[1] == pytest.approx(42.0)


def test_scheduled_fraction_tracks_fleet_size_and_ignores_risk_rescaling():
    size = 40
    locations = pd.DataFrame(
        {
            "battery": [f"d_{index}" for index in range(size)],
            "building": ["b_base"] * size,
            "room": ["r_base"] * size,
            "start_time": pd.to_datetime(["2025-01-01"] * size),
            "end_time": pd.to_datetime(["2026-12-01"] * size),
        }
    )
    travel = pd.DataFrame([("b_base", "b_base", 0.0)], columns=["from", "to", "hours"])
    settings = EvaluationSettings(base_location="b_base", base_room="r_base")
    snapshot = locations.assign(data_gap_days=0.0)
    risk = np.linspace(0.02, 0.6, size)
    rul = np.linspace(5.0, 40.0, size)
    horizon_end = pd.Timestamp("2025-09-01") + pd.Timedelta(days=42)

    def count(scale):
        planner = CompetitionPlanner(
            None,
            PlannerPolicy(
                use_expected_cost=True,
                risk_calibration_scale=scale,
                scheduled_fraction=0.25,
                minimum_scheduled_batteries=2,
                maximum_scheduled_batteries=30,
            ),
        )
        plan = planner.plan_snapshot(
            snapshot, locations, travel, settings, "2025-09-01",
            predicted_rul=rul, predicted_risk=risk,
        )
        return int(plan["day"].le(horizon_end).sum())

    assert count(1.0) == 10
    assert count(1.0) == count(1.9) == count(0.5)


def test_plan_ignores_label_columns_and_is_scenario_independent():
    """Guards rule 10: a plan must depend only on one scenario's own inputs."""

    locations, travel, settings = _inputs()
    snapshot = locations.assign(data_gap_days=0.0)
    planner = CompetitionPlanner(DummyEventTimeModel(), PlannerPolicy())

    first = planner.plan_snapshot(snapshot, locations, travel, settings, "2025-09-01")

    other = snapshot.assign(battery=["d_x", "d_y", "d_z"])
    planner.plan_snapshot(other, other, travel, settings, "2025-11-01")
    repeat = planner.plan_snapshot(snapshot, locations, travel, settings, "2025-09-01")
    pd.testing.assert_frame_equal(first, repeat)

    leaked = snapshot.assign(
        effective_eol=pd.to_datetime(["2025-09-05", "2025-09-06", "2025-09-07"]),
        target_rul_days=[4.0, 5.0, 6.0],
        event_observed=[1, 1, 1],
    )
    pd.testing.assert_frame_equal(
        first,
        planner.plan_snapshot(leaked, locations, travel, settings, "2025-09-01"),
    )


def test_selection_ignores_the_dataset_boundary():
    """V1/V2 priced healthy swaps off locations.end_time, so workload moved with the split."""

    size = 24
    locations = pd.DataFrame(
        {
            "battery": [f"d_{index}" for index in range(size)],
            "building": ["b_base"] * size,
            "room": ["r_base"] * size,
            "start_time": pd.to_datetime(["2025-01-01"] * size),
            "end_time": pd.to_datetime(["2026-08-01"] * size),
        }
    )
    travel = pd.DataFrame([("b_base", "b_base", 0.0)], columns=["from", "to", "hours"])
    settings = EvaluationSettings(base_location="b_base", base_room="r_base")
    risk = np.linspace(0.05, 0.7, size)
    rul = np.linspace(4.0, 40.0, size)
    survivor = np.full(size, 180.0)
    horizon_end = pd.Timestamp("2025-09-01") + pd.Timedelta(days=42)
    planner = CompetitionPlanner(
        None,
        PlannerPolicy(
            use_expected_cost=True,
            scheduled_fraction=0.25,
            minimum_scheduled_batteries=2,
            maximum_scheduled_batteries=20,
        ),
    )

    plans = []
    for shift in (-180, 0, 365):
        snapshot = locations.assign(
            data_gap_days=0.0,
            end_time=locations["end_time"] + pd.Timedelta(days=shift),
            censor_proxy_rul_days=float(shift) + 365.0,
        )
        plan = planner.plan_snapshot(
            snapshot, locations, travel, settings, "2025-09-01",
            predicted_rul=rul, predicted_risk=risk, predicted_survivor_rul=survivor,
        )
        plans.append(plan[plan["day"].le(horizon_end)].reset_index(drop=True))

    for plan in plans[1:]:
        pd.testing.assert_frame_equal(plans[0], plan)


def test_planner_runs_with_an_artifact_that_has_no_survivor_head():
    locations, travel, settings = _inputs()
    snapshot = locations.assign(data_gap_days=0.0)
    planner = CompetitionPlanner(DummyEventTimeModel(), PlannerPolicy(use_expected_cost=True))
    plan = planner.plan_snapshot(snapshot, locations, travel, settings, "2025-09-01")
    check_plan_valid(plan, locations, start_time=pd.Timestamp("2025-09-01"))


def test_weekly_guard_band_reserves_headroom_below_the_hard_limit():
    """The evaluator charges the weekly penalty at >= 24h, so the planner keeps a margin."""

    policy = PlannerPolicy(capacity_weekly_limit_fraction=0.95,
                           capacity_limit_penalty_multiplier=1.5)
    settings = EvaluationSettings(base_location="b_base", base_room="r_base")
    guarded = float(settings.worker_limit_weekly_hours) * policy.capacity_weekly_limit_fraction
    assert guarded < float(settings.worker_limit_weekly_hours)
    assert guarded == pytest.approx(22.8)
    assert policy.capacity_limit_penalty_multiplier > 1.0


def _train_split_available() -> bool:
    from pathlib import Path

    return Path("data/raw/train/scenarios.json").exists() and Path(
        "submission_artifacts/planner.joblib"
    ).exists()


@pytest.mark.skipif(not _train_split_available(), reason="train split or artifact absent")
def test_official_interface_matches_our_direct_scenario_path():
    """make_submissions -> Planner.plan must reproduce our iterate_scenarios path exactly."""

    import joblib
    from batteryswap_public.utils import iterate_scenarios, load_dataset
    from batteryswapai.competition_features import scenario_history_snapshot

    locations, timeseries, eol_times, scenarios = load_dataset(Path("data/raw/train"))
    subset = scenarios[:3]
    planner = joblib.load("submission_artifacts/planner.joblib")

    for scenario, locs, cut, _ in iterate_scenarios(locations, timeseries, eol_times, subset):
        start = pd.Timestamp(scenario["start_time"])
        rows = cut.reset_index()
        # the planner is handed exactly the scenario-visible history
        assert rows["end_time"].max() == start
        assert not (rows["end_time"] > start).any()

        official = planner.plan(cut, locs, scenario["travel_costs"], scenario["settings"])
        snapshot = scenario_history_snapshot(cut, locs, scenario["name"], start)
        direct = planner.plan_snapshot(
            snapshot, locs, scenario["travel_costs"], scenario["settings"], start
        )
        pd.testing.assert_frame_equal(official, direct)

        # ranking is deterministic across repeated calls
        again = planner.plan_snapshot(
            snapshot, locs, scenario["travel_costs"], scenario["settings"], start
        )
        pd.testing.assert_frame_equal(direct, again)