"""Load this repository offline after snapshot_download, using official LeRobot."""
from pathlib import Path


def load_policy(checkpoint_dir, device="cpu"):
    from lerobot.configs.policies import PreTrainedConfig
    from lerobot.policies.factory import get_policy_class, make_pre_post_processors

    root = Path(checkpoint_dir).resolve()
    config = PreTrainedConfig.from_pretrained(root)
    config.device = device
    config.pretrained_path = root
    overrides = {"device_processor": {"device": device}}
    if config.type == "smolvla":
        assets = root / "vlm_assets"
        config.vlm_model_name = str(assets)
        config.load_vlm_weights = False
        overrides["tokenizer_processor"] = {"tokenizer_name": str(assets)}
    elif config.type == "act":
        # The full trained backbone is in model.safetensors; avoid fetching its
        # initialization weights before restoring that complete state dictionary.
        config.pretrained_backbone_weights = None
    policy = get_policy_class(config.type).from_pretrained(root, config=config).eval().to(device)
    pre, post = make_pre_post_processors(
        config,
        pretrained_path=root,
        preprocessor_overrides=overrides,
        postprocessor_overrides={"device_processor": {"device": "cpu"}},
    )
    return policy, pre, post