"""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.constants import K_MAX
from giant.pipeline import run_setup_stage


def run_warm_setup_cache(
    data: str,
    val_fraction: float = 0.1,
    seed: int = 0,
    particle_conditioning: str = "physical",
    material_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`/`particle_conditioning`/`material_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. The two
    conditioning axes are independent and may differ.
    `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,
    }
    # A minimal v0.3 cfg — only the keys run_setup_stage actually reads
    # (conditioning.{particle,material}.type, stage{1,2}_model.router). This
    # CLI only ever configures one router (matching today's single
    # --router-type flag), so it's placed on stage1_model; stage2_model's
    # stays disabled.
    cfg = {
        "conditioning": {
            "particle": {"type": particle_conditioning},
            "material": {"type": material_conditioning},
        },
        "stage1_model": {"router": router_cfg},
        "stage2_model": {"router": {"enabled": False}, "k_max": K_MAX},
    }
    run_setup_stage(
        Path(data),
        val_fraction=val_fraction,
        seed=seed,
        cfg=cfg,
        cache_setup=True,
        rebuild_setup_cache=rebuild,
        echo=echo,
    )
    echo("setup cache warmed.")