"""Loading helpers shared by run_demo.py and run_experiments.py."""
from __future__ import annotations

import numpy as np

from . import config as C
from .data import load_processed
from .travel_time import TimeModel

MODEL_NAMES = {"nn": "NN (MLP)", "hist": "hist-avg", "dist": "distance-linear"}


def load_all():
    trips, zones, dist = load_processed()
    path = C.MODELS / "time_matrices.npz"
    if not path.exists():
        raise SystemExit("models/time_matrices.npz missing -- run `python train_nn.py` first")
    mats = dict(np.load(path))
    return trips, zones, dist, mats


def time_model(mats, kind: str, weekday: int, factor: float = C.VAN_FACTOR) -> TimeModel:
    return TimeModel(mats[kind][weekday], factor=factor, name=MODEL_NAMES.get(kind, kind))


def zone_weights(zones):
    return zones.sort_values("idx").n_dropoffs.values