"""dwarf warm-cache — precompute `giant train`'s setup-stage sidecar ahead of time.

Thin wrapper around `giant.pipeline.run_setup_stage` so a dataset's vocab
maps, event-id split index, and normalizer stats can be warmed once — e.g.
right after `dwarf convert`, or before kicking off a `dwarf hparam-scan`
sweep — without needing to also start training. See giant/data/setup_cache.py
for the sidecar itself.
"""

from pathlib import Path

from giant.pipeline import run_setup_stage


def run_warm_setup_cache(
    data: str,
    val_fraction: float = 0.1,
    seed: int = 0,
    conditioning: str = "physical",
    router_enabled: bool = False,
    router_type: str = "energy",
    n_experts: int = 4,
    rebuild: bool = False,
    echo=print,
) -> None:
    """Populate (or refresh) the setup cache sidecar for `data`.

    `val_fraction`/`seed`/`conditioning` select the normalizer cache entry
    (`giant.data.setup_cache.normalizer_key`) — pass the same values a later
    `giant train` invocation will use so it hits this warmed entry.
    `router_enabled`/`router_type`/`n_experts` only matter for
    `router_type == "process"` (warms that `n_experts`'s process map); the
    energy-router quantile summary is always collected regardless, so a
    later `--router-type energy` run never needs to rescan just to seed
    centers.
    """
    router_cfg = {
        "enabled": router_enabled,
        "type": router_type,
        "n_experts": n_experts,
    }
    run_setup_stage(
        Path(data),
        val_fraction=val_fraction,
        seed=seed,
        conditioning=conditioning,
        router_cfg=router_cfg,
        cache_setup=True,
        rebuild_setup_cache=rebuild,
        echo=echo,
    )
    echo("setup cache warmed.")