"""Load the released RoboTwin Pi05 checkpoint with portable asset paths."""

from pathlib import Path

from openpi.models.pi0_config import Pi0Config
from openpi.policies.policy_config import create_trained_policy
from openpi.training.config import AssetsConfig, DataConfig, LeRobotAlohaDataConfig, TrainConfig

ASSET_ID = "RoboTwin-pi05_900-aloha_agilex-joint"
CONFIG_NAME = "pi05_base_aloha_full_sim_robotwin_900_defaults"


def make_config(checkpoint_dir):
    checkpoint_dir = Path(checkpoint_dir).resolve()
    return TrainConfig(
        name=CONFIG_NAME,
        model=Pi0Config(pi05=True),
        data=LeRobotAlohaDataConfig(
            repo_id=ASSET_ID,
            assets=AssetsConfig(
                assets_dir=str(checkpoint_dir / "assets"),
                asset_id=ASSET_ID,
            ),
            base_config=DataConfig(prompt_from_task=True),
            use_delta_joint_actions=True,
            adapt_to_pi=False,
        ),
    )


def load_policy(checkpoint_dir="14999"):
    checkpoint_dir = Path(checkpoint_dir).resolve()
    if not (checkpoint_dir / "params").is_dir():
        raise FileNotFoundError(f"Missing parameters: {checkpoint_dir / 'params'}")
    stats = checkpoint_dir / "assets" / ASSET_ID / "norm_stats.json"
    if not stats.is_file():
        raise FileNotFoundError(f"Missing normalization statistics: {stats}")
    return create_trained_policy(make_config(checkpoint_dir), str(checkpoint_dir))