scripts/validate_submission.py

from __future__ import annotations

import argparse
import itertools
import json
from pathlib import Path

import numpy as np
import pandas as pd
from batteryswap_public.evaluate import evaluate_plan
from batteryswap_public.utils import iterate_scenarios, load_dataset
from sklearn.metrics import average_precision_score, roc_auc_score
from sklearn.model_selection import GroupKFold

from batteryswapai.competition_features import (
    build_trajectory_matrix,
    scenario_history_snapshot,
    attach_training_targets,
    build_daily_features,
    scenario_snapshot,
)
from batteryswapai.competition_model import (
    blend_risk_components,
    blend_trajectory,
    fit_event_time_model,
)
from batteryswapai.competition_planner import CompetitionPlanner, PlannerPolicy
from train_submission import DATASET_REVISION


def _numbers(value: str) -> list[float]:
    return [float(item.strip()) for item in value.split(",") if item.strip()]


def _integers(value: str) -> list[int]:
    return [int(item.strip()) for item in value.split(",") if item.strip()]


def _schedule_plans(
    fractions: str, quotas: str, bands: str
) -> list[tuple[float | None, int, int | None]]:
    plans: list[tuple[float | None, int, int | None]] = [
        (fraction, 8, 24) for fraction in _numbers(fractions)
    ]
    plans.extend((None, minimum, maximum) for minimum, maximum in _schedule_bands(quotas, bands))
    return plans


def _schedule_bands(quotas: str, bands: str) -> list[tuple[int, int]]:
    values = [(quota, quota) for quota in _integers(quotas)]
    for item in bands.split(","):
        item = item.strip()
        if not item:
            continue
        minimum, maximum = item.split(":", maxsplit=1)
        values.append((int(minimum), int(maximum)))
    return list(dict.fromkeys(values))


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset-path", type=Path, default=Path("data/raw/train"))
    parser.add_argument("--folds", type=int, default=5)
    parser.add_argument("--quantile", type=float, default=0.05)
    parser.add_argument("--inner-calibration-folds", type=int, default=3)
    parser.add_argument(
        "--outer-group",
        choices=("battery", "building"),
        default="building",
        help=(
            "Hold out entire buildings by default to emulate hidden-domain "
            "generalization; use battery only to reproduce legacy reports."
        ),
    )
    parser.add_argument("--risk-scales", default="1.50")
    parser.add_argument("--gain-margins", default="10")
    parser.add_argument("--service-costs", default="2")
    parser.add_argument("--emergency-operational-scales", default="0")
    parser.add_argument("--offsets", default="-5")
    parser.add_argument("--schedule-fractions", default="0.038")
    parser.add_argument("--schedule-quotas", default="")
    parser.add_argument("--schedule-bands", default="")
    parser.add_argument("--capacity-lookback-days", type=int, default=42)
    parser.add_argument("--weekly-guard-fractions", default="0.95")
    parser.add_argument("--hard-limit-penalty-multipliers", default="1.5")
    parser.add_argument(
        "--predictions-csv",
        type=Path,
        help="Reuse previously generated OOF predictions for policy-only sweeps.",
    )
    parser.add_argument(
        "--output", type=Path, default=Path("artifacts/submission_cv_final.json")
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    locations, timeseries, eol_times, scenarios = load_dataset(args.dataset_path)
    daily = build_daily_features(timeseries)

    snapshots = []
    scenario_inputs = {}
    for scenario, locs, visible, not_dead in iterate_scenarios(
        locations, timeseries, eol_times, scenarios
    ):
        snapshot = scenario_history_snapshot(
            visible, locs, scenario["name"], scenario["start_time"]
        )
        snapshot = attach_training_targets(
            snapshot,
            eol_times,
            unobserved_eol_days=float(scenario["settings"].unobserved_eol_days),
        )
        snapshots.append(snapshot)
        scenario_inputs[scenario["name"]] = (scenario, locs, not_dead)
    training = pd.concat(snapshots, ignore_index=True)

    oof_risk = np.full(len(training), np.nan)
    oof_rul = np.full(len(training), np.nan)
    oof_survivor = np.full(len(training), np.nan)
    oof_components: dict[str, np.ndarray] = {}
    oof_residual = np.full(len(training), np.nan)
    oof_residual_weight = np.zeros(len(training))
    blend_weights = (1.0, 0.0, 0.0, 0.0)
    outer_fold_audits: list[dict[str, int]] = []
    if args.predictions_csv is not None:
        cached = pd.read_csv(args.predictions_csv).set_index(["scenario", "battery"])
        keys = pd.MultiIndex.from_frame(training[["scenario", "battery"]])
        aligned = cached.reindex(keys)
        if aligned[["oof_event_risk", "oof_rul_days"]].isna().any().any():
            raise ValueError("Cached OOF predictions do not cover the validation rows")
        oof_risk[:] = aligned["oof_event_risk"].to_numpy(dtype=float)
        oof_rul[:] = aligned["oof_rul_days"].to_numpy(dtype=float)
        if "oof_survivor_rul" in aligned.columns:
            oof_survivor[:] = aligned["oof_survivor_rul"].to_numpy(dtype=float)
    else:
        splitter = GroupKFold(n_splits=args.folds)
        groups = training[args.outer_group].astype(str)
        for fold_number, (train_index, valid_index) in enumerate(
            splitter.split(training, groups=groups), start=1
        ):
            train_rows = training.iloc[train_index]
            valid_rows = training.iloc[valid_index]
            train_batteries = set(train_rows["battery"].astype(str))
            valid_batteries = set(valid_rows["battery"].astype(str))
            train_buildings = set(train_rows["building"].astype(str))
            valid_buildings = set(valid_rows["building"].astype(str))
            if train_batteries & valid_batteries:
                raise AssertionError(f"outer fold {fold_number} leaks a battery")
            if args.outer_group == "building" and train_buildings & valid_buildings:
                raise AssertionError(f"outer fold {fold_number} leaks a building")

            # The dense trajectory forecasters have future-voltage and EOL targets.
            # Their matrix and labels must therefore obey the same outer split as the
            # snapshot model. Passing the full matrix here made the old OOF report
            # optimistic even though hidden inference itself remained train-only.
            fold_daily = daily.loc[
                daily["device_id"].astype(str).isin(train_batteries)
            ].copy()
            fold_trajectory = build_trajectory_matrix(fold_daily)
            trajectory_batteries = set(fold_trajectory["index"])
            if trajectory_batteries & valid_batteries:
                raise AssertionError(
                    f"outer fold {fold_number} leaks a trajectory battery"
                )
            fold_eol_times = eol_times.reindex(sorted(train_batteries)).copy()
            if set(fold_eol_times.index.astype(str)) & valid_batteries:
                raise AssertionError(f"outer fold {fold_number} leaks an EOL label")

            model = fit_event_time_model(
                train_rows,
                quantile=args.quantile,
                dataset_revision=DATASET_REVISION,
                random_state=2026 + fold_number,
                calibration_folds=args.inner_calibration_folds,
                trajectory=fold_trajectory,
                eol_times=fold_eol_times,
            )
            valid = valid_rows
            for name, values in model.predict_risk_components(valid).items():
                oof_components.setdefault(
                    name, np.full(len(training), np.nan)
                )[valid_index] = values
            blend_weights = model.blend_weights
            residual, residual_weight = model._residual_risk(valid)
            if residual is not None:
                oof_residual[valid_index] = residual
                oof_residual_weight[valid_index] = residual_weight
            oof_rul[valid_index] = model.predict_rul(valid)
            survivor = model.predict_survivor_rul(valid)
            if survivor is not None:
                oof_survivor[valid_index] = survivor
            outer_fold_audits.append(
                {
                    "fold": fold_number,
                    "train_batteries": len(train_batteries),
                    "valid_batteries": len(valid_batteries),
                    "train_buildings": len(train_buildings),
                    "valid_buildings": len(valid_buildings),
                    "trajectory_batteries": len(trajectory_batteries),
                    "battery_overlap": 0,
                    "building_overlap": 0 if args.outer_group == "building" else -1,
                }
            )

    scenario_groups = training["scenario"].astype(str).to_numpy()
    if oof_components:
        oof_risk[:] = blend_risk_components(oof_components, blend_weights, scenario_groups)
    # Mirror inference: the trajectory nudge is applied over complete scenarios.
    if np.any(oof_residual_weight > 0.0):
        oof_risk[:] = blend_trajectory(
            oof_risk, np.nan_to_num(oof_residual), oof_residual_weight, scenario_groups
        )

    due = (
        training["event_observed"].astype(bool)
        & training["target_rul_days"].between(0.0, 42.0)
    ).astype(int)
    classification = {
        "positive_rows": int(due.sum()),
        "total_rows": len(due),
        "roc_auc": float(roc_auc_score(due, oof_risk)),
        "average_precision": float(average_precision_score(due, oof_risk)),
    }
    adjusted_oof_risk = oof_risk.copy()
    gaps = pd.to_numeric(training["data_gap_days"], errors="coerce").to_numpy(dtype=float)
    recent_gap = (gaps > 0.0) & (gaps <= 7.0)
    stale = (gaps > 7.0) | ~np.isfinite(gaps)
    adjusted_oof_risk[recent_gap] *= 0.75
    adjusted_oof_risk[stale] *= 0.10
    classification["freshness_adjusted_roc_auc"] = float(
        roc_auc_score(due, adjusted_oof_risk)
    )
    classification["freshness_adjusted_average_precision"] = float(
        average_precision_score(due, adjusted_oof_risk)
    )

    observed = training["event_observed"].astype(bool)
    observed_rul_error = np.abs(
        oof_rul[observed] - training.loc[observed, "target_rul_days"].to_numpy()
    )
    due_mask = due.astype(bool).to_numpy()
    due_rul_error = np.abs(
        oof_rul[due_mask] - training.loc[due_mask, "target_rul_days"].to_numpy()
    )
    classification["all_observed_rul_mae_days"] = float(np.mean(observed_rul_error))
    classification["due_within_horizon_rul_mae_days"] = float(np.mean(due_rul_error))

    policy_results = []
    case_results = []
    placeholder_model = None
    grid = itertools.product(
        _numbers(args.risk_scales),
        _numbers(args.gain_margins),
        _numbers(args.service_costs),
        _numbers(args.emergency_operational_scales),
        _numbers(args.weekly_guard_fractions),
        _numbers(args.hard_limit_penalty_multipliers),
        _numbers(args.offsets),
        _schedule_plans(
            args.schedule_fractions,
            args.schedule_quotas,
            args.schedule_bands,
        ),
    )
    for (
        risk_scale,
        gain_margin,
        service_cost,
        emergency_scale,
        weekly_guard,
        hard_multiplier,
        offset,
        (fraction, minimum, maximum),
    ) in grid:
        policy_id = f"p{len(policy_results):03d}"
        planner = CompetitionPlanner(
            placeholder_model,
            PlannerPolicy(
                event_risk_threshold=0.50,
                prediction_offset_days=offset,
                use_expected_cost=True,
                risk_calibration_scale=risk_scale,
                expected_service_cost_hours=service_cost,
                expected_gain_margin=gain_margin,
                emergency_operational_scale=emergency_scale,
                capacity_lookahead_days=21,
                capacity_lookback_days=args.capacity_lookback_days,
                capacity_weekly_limit_fraction=weekly_guard,
                capacity_limit_penalty_multiplier=hard_multiplier,
                scheduled_fraction=fraction,
                minimum_scheduled_batteries=minimum,
                maximum_scheduled_batteries=maximum,
            ),
        )
        scores = []
        scheduled_counts = []
        for scenario_name, (scenario, locs, not_dead) in scenario_inputs.items():
            mask = training["scenario"].eq(scenario_name).to_numpy()
            snapshot = training.loc[mask]
            risk = oof_risk[mask]
            rul = oof_rul[mask]
            survivor = oof_survivor[mask]
            plan = planner.plan_snapshot(
                snapshot,
                locs,
                scenario["travel_costs"],
                scenario["settings"],
                scenario["start_time"],
                predicted_rul=rul,
                predicted_risk=risk,
                predicted_survivor_rul=(
                    None if np.isnan(survivor).all() else survivor
                ),
            )
            start = pd.Timestamp(scenario["start_time"])
            horizon_end = start + pd.Timedelta(
                days=scenario["settings"].planning_window_days
            )
            scheduled = plan["day"].le(horizon_end)
            scheduled_count = int(scheduled.sum())
            scheduled_counts.append(scheduled_count)
            _, _, score = evaluate_plan(
                plan,
                locs,
                scenario["travel_costs"],
                scenario["settings"],
                eol_times=not_dead,
                start_time=start,
                verbose=0,
            )
            scores.append(score)
            required = pd.to_datetime(not_dead).between(
                start,
                horizon_end,
                inclusive="right",
            )
            required_ids = set(not_dead.index[required].astype(str))
            scheduled_ids = set(plan.loc[scheduled, "battery"].astype(str))
            case_results.append(
                {
                    "policy_id": policy_id,
                    "scenario": scenario_name,
                    "start_time": start.isoformat(),
                    "scheduled_count": scheduled_count,
                    "required_count": len(required_ids),
                    "true_positive_count": len(required_ids & scheduled_ids),
                    "missed_count": len(required_ids - scheduled_ids),
                    **{key: float(value) for key, value in score.items()},
                }
            )
        mean_score = pd.concat(scores, axis=1).mean(axis=1)
        policy_results.append(
            {
                "policy_id": policy_id,
                "risk_calibration_scale": risk_scale,
                "expected_gain_margin": gain_margin,
                "expected_service_cost_hours": service_cost,
                "emergency_operational_scale": emergency_scale,
                "prediction_offset_days": offset,
                "scheduled_fraction": fraction,
                "minimum_scheduled_batteries": minimum,
                "maximum_scheduled_batteries": maximum,
                "capacity_weekly_limit_fraction": weekly_guard,
                "capacity_limit_penalty_multiplier": hard_multiplier,
                "stale_risk_cutoff_days": planner.policy.stale_risk_cutoff_days,
                "recent_gap_risk_factor": planner.policy.recent_gap_risk_factor,
                "stale_risk_factor": planner.policy.stale_risk_factor,
                "building_batch_window_days": planner.policy.building_batch_window_days,
                "capacity_operational_cost_weight": planner.policy.capacity_operational_cost_weight,
                "mean_scheduled_batteries": float(np.mean(scheduled_counts)),
                **{key: float(value) for key, value in mean_score.items()},
            }
        )

    policy_results.sort(key=lambda item: item["total_cost"])
    report = {
        "dataset_revision": DATASET_REVISION,
        "folds": args.folds,
        "inner_calibration_folds": args.inner_calibration_folds,
        "outer_group": args.outer_group,
        "prediction_source": (
            args.predictions_csv.as_posix()
            if args.predictions_csv is not None
            else "generated by this run"
        ),
        "outer_fold_trajectory_scope": (
            "caller-supplied predictions; not refit"
            if args.predictions_csv is not None
            else "outer-train batteries only"
        ),
        "outer_fold_eol_scope": (
            "caller-supplied predictions; not refit"
            if args.predictions_csv is not None
            else "outer-train batteries only"
        ),
        "outer_fold_audits": outer_fold_audits,
        "quantile": args.quantile,
        "classification": classification,
        "policies": policy_results,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(report, indent=2), encoding="utf-8")
    predictions = training[["scenario", "battery", "target_rul_days", "event_observed"]].copy()
    predictions["oof_event_risk"] = oof_risk
    predictions["oof_rul_days"] = oof_rul
    predictions["oof_survivor_rul"] = oof_survivor
    predictions.to_csv(args.output.with_suffix(".csv"), index=False)
    pd.DataFrame(case_results).to_csv(
        args.output.with_suffix(".cases.csv"), index=False
    )
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()