tests/test_competition_features.py

import numpy as np
import pandas as pd

from batteryswapai.competition_features import (
    build_daily_features,
    numeric_feature_columns,
    scenario_history,
    scenario_history_snapshot,
    scenario_snapshot,
)


def test_scenario_snapshot_never_uses_future_sensor_rows():
    raw = pd.DataFrame(
        {
            "device_id": ["d_a"] * 6,
            "end_time": pd.to_datetime(
                [
                    "2026-01-01 10:00",
                    "2026-01-01 12:00",
                    "2026-01-02 10:00",
                    "2026-01-02 12:00",
                    "2026-01-03 10:00",
                    "2026-01-03 12:00",
                ]
            ),
            "voltage": [3.1, 3.0, 2.9, 2.8, 1.0, 1.0],
            "temperature": [20.0] * 6,
        }
    ).set_index(["device_id", "end_time"])
    locations = pd.DataFrame(
        {
            "battery": ["d_a"],
            "building": ["b_a"],
            "room": ["r_a"],
            "start_time": pd.to_datetime(["2025-01-01"]),
            "end_time": pd.to_datetime(["2026-02-01"]),
        }
    )

    # The 2026-01-02 readings are timestamped after the 00:00 scenario start, so the
    # evaluator's cut excludes them and the freshest usable bucket is the previous day.
    snapshot = scenario_history_snapshot(raw, locations, "s_0", "2026-01-02")
    assert snapshot.loc[0, "day"] == pd.Timestamp("2026-01-01")
    assert np.isclose(snapshot.loc[0, "voltage_median"], 3.05)


def test_temperature_filter_falls_back_when_no_stable_reading_exists():
    raw = pd.DataFrame(
        {
            "device_id": ["d_a", "d_a"],
            "end_time": pd.to_datetime(["2026-01-01 10:00", "2026-01-01 12:00"]),
            "voltage": [3.0, 2.8],
            "temperature": [-5.0, 40.0],
        }
    ).set_index(["device_id", "end_time"])
    daily = build_daily_features(raw)
    assert daily.loc[0, "stable_voltage_median"] == daily.loc[0, "voltage_median"]


def test_model_features_exclude_dataset_boundary_and_calendar_fields():
    snapshot = pd.DataFrame(
        {
            "voltage_median": [3.0],
            "location_age_days": [120.0],
            "censor_proxy_rul_days": [300.0],
            "scenario_month_sin": [0.5],
            "scenario_month_cos": [-0.5],
        }
    )

    columns = numeric_feature_columns(snapshot)

    assert "voltage_median" in columns
    assert "location_age_days" in columns
    assert "censor_proxy_rul_days" not in columns
    assert "scenario_month_sin" not in columns
    assert "scenario_month_cos" not in columns


def _voltage_history(days: int = 420, start: str = "2025-01-01") -> pd.DataFrame:
    stamps = pd.date_range(start, periods=days, freq="D")
    rows = []
    for device, offset in (("d_a", 0.0), ("d_b", 0.12)):
        rows.append(
            pd.DataFrame(
                {
                    "device_id": device,
                    "day": stamps,
                    "stable_voltage_median": np.linspace(2.90, 2.42, days) + offset,
                }
            )
        )
    return pd.concat(rows, ignore_index=True)


def test_trajectory_bins_never_read_past_the_cutoff():
    """The bins must be identical whether or not future rows exist in the frame."""

    from batteryswapai.competition_features import build_trajectory_matrix, trajectory_bins

    history = _voltage_history()
    cutoff = pd.Timestamp("2025-10-01")

    full = build_trajectory_matrix(history.copy())
    truncated = build_trajectory_matrix(history[history["day"] <= cutoff].copy())

    from_full = trajectory_bins(full, ["d_a", "d_b"], cutoff)
    from_truncated = trajectory_bins(truncated, ["d_a", "d_b"], cutoff)
    assert np.allclose(from_full, from_truncated, equal_nan=True)
    assert np.isfinite(from_full).any()


def test_trajectory_bins_are_missing_without_history():
    from batteryswapai.competition_features import (
        MIN_TRAJECTORY_COVERAGE,
        build_trajectory_matrix,
        trajectory_bins,
    )

    matrix = build_trajectory_matrix(_voltage_history(days=40, start="2025-01-01").copy())
    early = trajectory_bins(matrix, ["d_a"], pd.Timestamp("2025-01-20"))
    assert np.isfinite(early).mean() < MIN_TRAJECTORY_COVERAGE
    unknown = trajectory_bins(matrix, ["d_missing"], pd.Timestamp("2025-01-20"))
    assert not np.isfinite(unknown).any()


def test_scenario_history_matches_the_evaluator_cut():
    """scenario_history must reproduce iterate_scenarios' truncation exactly."""

    stamps = pd.date_range("2025-05-01", "2025-06-30", freq="6h")
    frame = pd.DataFrame(
        {
            "device_id": "d_a",
            "end_time": stamps,
            "voltage": np.linspace(2.9, 2.5, len(stamps)),
            "temperature": 20.0,
        }
    ).set_index(["device_id", "end_time"])
    cutoff = pd.Timestamp("2025-06-01")

    visible = scenario_history(frame, cutoff).reset_index()
    assert visible["end_time"].max() == cutoff          # endpoint is inclusive
    assert not (visible["end_time"] > cutoff).any()
    expected = frame.reset_index()
    expected = expected[expected["end_time"] <= cutoff]
    assert len(visible) == len(expected)


def test_features_ignore_readings_taken_after_the_scenario_start():
    cutoff = pd.Timestamp("2025-06-01")
    stamps = pd.date_range("2025-01-01", "2025-06-30", freq="6h")
    frame = pd.DataFrame(
        {
            "device_id": "d_a",
            "end_time": stamps,
            "voltage": np.linspace(2.90, 2.40, len(stamps)),
            "temperature": 20.0,
        }
    ).set_index(["device_id", "end_time"])
    locations = pd.DataFrame(
        {
            "battery": ["d_a"],
            "building": ["b_1"],
            "room": ["r_1"],
            "start_time": [pd.Timestamp("2024-01-01")],
            "end_time": [pd.Timestamp("2025-06-30")],
        }
    )
    flat = frame.reset_index()
    already_cut = flat[flat["end_time"] <= cutoff].set_index(["device_id", "end_time"])

    full = scenario_history_snapshot(frame, locations, "s", cutoff)
    trimmed = scenario_history_snapshot(already_cut, locations, "s", cutoff)
    columns = [c for c in full.columns if pd.api.types.is_numeric_dtype(full[c])]
    assert np.allclose(
        full[columns].astype(float).to_numpy(),
        trimmed[columns].astype(float).to_numpy(),
        equal_nan=True,
    )


def test_trajectory_availability_does_not_depend_on_later_readings():
    """A sensor gap must look the same whether or not the device reports again later."""

    from batteryswapai.competition_features import build_trajectory_matrix, trajectory_bins

    cutoff = pd.Timestamp("2025-06-01")
    early = pd.date_range("2024-01-01", "2025-03-01", freq="D")
    later = pd.date_range("2025-08-01", "2025-09-01", freq="D")   # resumes AFTER the cutoff

    def frame(days):
        return pd.DataFrame(
            {
                "device_id": "d_a",
                "day": days,
                "stable_voltage_median": np.linspace(2.90, 2.60, len(days)),
            }
        )

    without_future = build_trajectory_matrix(frame(early).copy())
    with_future = build_trajectory_matrix(
        pd.concat([frame(early), frame(later)], ignore_index=True).copy()
    )
    a = trajectory_bins(without_future, ["d_a"], cutoff)
    b = trajectory_bins(with_future, ["d_a"], cutoff)
    assert np.allclose(a, b, equal_nan=True)
    assert np.array_equal(np.isfinite(a), np.isfinite(b))