# Copyright 2025 The RLinf Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# openpi model configs

import dataclasses
import difflib
from typing import Optional

import openpi.models.pi0_config as pi0_config
import openpi.training.optimizer as _optimizer
import openpi.training.weight_loaders as weight_loaders
from openpi.training.config import (
    AssetsConfig,
    DataConfig,
    TrainConfig,
)

from rlinf.models.embodiment.openpi.dataconfig.behavior_dataconfig import (
    LeRobotBehaviorDataConfig,
)
from rlinf.models.embodiment.openpi.dataconfig.calvin_dataconfig import (
    LeRobotCalvinDataConfig,
)
from rlinf.models.embodiment.openpi.dataconfig.dual_franka_tcp_rot6d_dataconfig import (
    DualFrankaTcpRot6dDataConfig,
)
from rlinf.models.embodiment.openpi.dataconfig.franka_co_training_dataconfig import (
    LeRobotFrankaEEDataConfig,
)
from rlinf.models.embodiment.openpi.dataconfig.franka_dataconfig import (
    CustomDataConfig,
)
from rlinf.models.embodiment.openpi.dataconfig.gsenv_dataconfig import (
    LeRobotGSEnvDataConfig,
)
from rlinf.models.embodiment.openpi.dataconfig.isaaclab_dataconfig import (
    LeRobotIsaacLabStackCubeDataConfig,
)
from rlinf.models.embodiment.openpi.dataconfig.libero_dataconfig import (
    LeRobotLiberoDataConfig,
)
from rlinf.models.embodiment.openpi.dataconfig.maniskill_dataconfig import (
    LeRobotManiSkillDataConfig,
)
from rlinf.models.embodiment.openpi.dataconfig.metaworld_dataconfig import (
    LeRobotMetaworldDataConfig,
)
from rlinf.models.embodiment.openpi.dataconfig.polaris_dataconfig import (
    LeRobotPolarisDroidDataConfig,
)
from rlinf.models.embodiment.openpi.dataconfig.realworld_dataconfig import (
    LeRobotRealworldDataConfig,
)
from rlinf.models.embodiment.openpi.dataconfig.robocasa_dataconfig import (
    LeRobotRobocasaDataConfig,
)
from rlinf.models.embodiment.openpi.dataconfig.robotwin_aloha_dataconfig import (
    LeRobotAlohaDataConfig,
)

_CONFIGS = [
    TrainConfig(
        name="pi0_libero",
        model=pi0_config.Pi0Config(),
        data=LeRobotLiberoDataConfig(
            repo_id="physical-intelligence/libero",
            base_config=DataConfig(prompt_from_task=True),
            assets=AssetsConfig(assets_dir="checkpoints/torch/pi0_libero/assets"),
            extra_delta_transform=True,
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader(
            "checkpoints/jax/pi0_base/params"
        ),
        pytorch_weight_path="checkpoints/torch/pi0_base",
    ),
    TrainConfig(
        name="pi05_libero",
        model=pi0_config.Pi0Config(
            pi05=True, action_horizon=10, discrete_state_input=False
        ),
        data=LeRobotLiberoDataConfig(
            repo_id="physical-intelligence/libero",
            base_config=DataConfig(prompt_from_task=True),
            assets=AssetsConfig(assets_dir="checkpoints/torch/pi0_libero/assets"),
            extra_delta_transform=False,
        ),
        batch_size=256,
        lr_schedule=_optimizer.CosineDecaySchedule(
            warmup_steps=10_000,
            peak_lr=5e-5,
            decay_steps=1_000_000,
            decay_lr=5e-5,
        ),
        optimizer=_optimizer.AdamW(clip_gradient_norm=1.0),
        ema_decay=0.999,
        weight_loader=weight_loaders.CheckpointWeightLoader(
            "checkpoints/jax/pi05_base"
        ),
        pytorch_weight_path="checkpoints/torch/pi05_base",
    ),
    TrainConfig(
        name="pi0_maniskill",
        model=pi0_config.Pi0Config(),
        data=LeRobotManiSkillDataConfig(
            repo_id="physical-intelligence/maniskill",
            base_config=DataConfig(prompt_from_task=True),
            assets=AssetsConfig(assets_dir="checkpoints/torch/pi0_base"),
            extra_delta_transform=False,
        ),
        pytorch_weight_path="checkpoints/torch/pi0_base",
        seed=0,
        batch_size=32,
        num_workers=8,
        num_train_steps=200,  # 1_000, #30_000
        log_interval=5,  # 25,
        save_interval=50,  # 200,
    ),
    TrainConfig(
        name="pi05_maniskill",
        model=pi0_config.Pi0Config(
            pi05=True, action_horizon=10, discrete_state_input=False
        ),  # discrete_state_input=False: stateless policy, True: with state policy
        data=LeRobotManiSkillDataConfig(
            repo_id="physical-intelligence/maniskill",
            base_config=DataConfig(prompt_from_task=True),
            assets=AssetsConfig(assets_dir="checkpoints/torch/pi05_maniskill/assets"),
            extra_delta_transform=False,
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader(
            "checkpoints/jax/pi05_base"
        ),
        pytorch_weight_path="checkpoints/torch/pi05_base",
        seed=0,
        batch_size=256,
        optimizer=_optimizer.AdamW(clip_gradient_norm=1.0),
        ema_decay=0.999,
        num_workers=8,
        num_train_steps=5_000,
        log_interval=5,
        save_interval=250,
    ),
    TrainConfig(
        name="pi05_franka",
        model=pi0_config.Pi0Config(
            pi05=True, action_horizon=8, discrete_state_input=False
        ),  # discrete_state_input=False: stateless policy, True: with state policy
        data=LeRobotFrankaEEDataConfig(
            repo_id="physical-intelligence/real_rl",  # Not important
            base_config=DataConfig(prompt_from_task=True),
            assets=AssetsConfig(
                assets_dir="checkpoints/torch/pi05_franka_pretrained/assets"
            ),
            output_action_dim=6,
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader(
            "checkpoints/jax/pi05_base"
        ),
        pytorch_weight_path="checkpoints/torch/pi05_base",
        seed=0,
        batch_size=16,
        optimizer=_optimizer.AdamW(clip_gradient_norm=1.0),
        ema_decay=0.999,
        num_workers=8,
        num_train_steps=5_000,
        log_interval=5,
        save_interval=250,
    ),
    TrainConfig(
        name="pi05_maniskill_sim_real_co_training",
        model=pi0_config.Pi0Config(
            pi05=True, action_horizon=8, discrete_state_input=False
        ),  # discrete_state_input=False: stateless policy, True: with state policy
        data=LeRobotFrankaEEDataConfig(
            repo_id="physical-intelligence/pick_and_place_real",
            default_prompt="default prompt",
            base_config=DataConfig(prompt_from_task=True),
            assets=AssetsConfig(
                assets_dir="checkpoints/torch/pi05_maniskill_sim_real_co_training/assets"
            ),
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader(
            "checkpoints/jax/pi05_base"
        ),
        pytorch_weight_path="checkpoints/torch/pi05_base",
        seed=0,
        batch_size=16,
        optimizer=_optimizer.AdamW(clip_gradient_norm=1.0),
        ema_decay=0.999,
        num_workers=8,
        num_train_steps=5_000,
        log_interval=5,
        save_interval=250,
    ),
    TrainConfig(
        name="pi0_metaworld",
        model=pi0_config.Pi0Config(action_horizon=5),
        data=LeRobotMetaworldDataConfig(
            repo_id="lerobot/metaworld_mt50",
            base_config=DataConfig(prompt_from_task=True),
            assets=AssetsConfig(assets_dir="checkpoints/torch/pi0_metaworld/assets"),
            extra_delta_transform=False,
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader(
            "checkpoints/jax/pi0_base/params"
        ),
        pytorch_weight_path="checkpoints/torch/pi0_base",
    ),
    TrainConfig(
        name="pi05_metaworld",
        model=pi0_config.Pi0Config(
            pi05=True, action_horizon=5, discrete_state_input=False
        ),
        data=LeRobotMetaworldDataConfig(
            repo_id="lerobot/metaworld_mt50",
            base_config=DataConfig(prompt_from_task=True),
            assets=AssetsConfig(assets_dir="checkpoints/torch/pi0_metaworld/assets"),
            extra_delta_transform=False,
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader(
            "checkpoints/jax/pi05_base/params"
        ),
        pytorch_weight_path="checkpoints/torch/pi05_base",
    ),
    TrainConfig(
        name="pi0_calvin",
        model=pi0_config.Pi0Config(action_horizon=5),
        data=LeRobotCalvinDataConfig(
            repo_id="InternRobotics/InternData-Calvin_ABC",
            base_config=DataConfig(
                prompt_from_task=True,
            ),
            assets=AssetsConfig(assets_dir="checkpoints/torch/pi0_calvin/assets"),
            extra_delta_transform=False,
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader(
            "checkpoints/jax/pi0_base/params"
        ),
        pytorch_weight_path="checkpoints/torch/pi0_base",
        num_train_steps=30_000,
    ),
    TrainConfig(
        name="pi05_calvin",
        model=pi0_config.Pi0Config(
            pi05=True, action_horizon=5, discrete_state_input=False
        ),
        data=LeRobotCalvinDataConfig(
            repo_id="InternRobotics/InternData-Calvin_ABC",
            base_config=DataConfig(
                prompt_from_task=True,
            ),
            assets=AssetsConfig(assets_dir="checkpoints/torch/pi0_calvin/assets"),
            extra_delta_transform=False,
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader(
            "checkpoints/jax/pi05_base/params"
        ),
        pytorch_weight_path="checkpoints/torch/pi05_base",
        num_train_steps=30_000,
    ),
    TrainConfig(
        name="pi0_robocasa_human",
        model=pi0_config.Pi0Config(action_horizon=5),
        data=LeRobotRobocasaDataConfig(
            repo_id="daixianjie/robocasa_human_lerobot",
            base_config=DataConfig(prompt_from_task=True),
            assets=AssetsConfig(asset_id="physical-intelligence/robocasa_all_human"),
            extra_delta_transform=False,
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader(
            "checkpoints/jax/pi0_base/params"
        ),
        pytorch_weight_path="checkpoints/torch/pi0_base",
        num_train_steps=100_000,
    ),
    TrainConfig(
        name="pi0_aloha_robotwin",
        model=pi0_config.Pi0Config(discrete_state_input=False),
        data=LeRobotAlohaDataConfig(
            repo_id="physical-intelligence/robotwin",
            adapt_to_pi=False,
            base_config=DataConfig(prompt_from_task=True),
            assets=AssetsConfig(
                assets_dir="checkpoints/torch/pi0_aloha_robotwin/assets"
            ),
            extra_delta_transform=True,  # True for delta action, False for abs_action
        ),
        freeze_filter=pi0_config.Pi0Config().get_freeze_filter(),
        pytorch_weight_path="checkpoints/torch/pi0_base",
        num_train_steps=30_000,
    ),
    TrainConfig(
        name="pi05_aloha_robotwin",
        model=pi0_config.Pi0Config(pi05=True, discrete_state_input=True),
        data=LeRobotAlohaDataConfig(
            repo_id="physical-intelligence/robotwin",
            base_config=DataConfig(prompt_from_task=True),
            assets=AssetsConfig(
                assets_dir="checkpoints/torch/pi05_aloha_robotwin/assets"
            ),
            extra_delta_transform=True,  # True for delta action, False for abs_action
        ),
        pytorch_weight_path="checkpoints/torch/pi05_base",
        num_train_steps=20_000,
    ),
    TrainConfig(
        name="pi0_behavior",
        model=pi0_config.Pi0Config(),
        data=LeRobotBehaviorDataConfig(
            repo_id="physical-intelligence/behavior",
            base_config=DataConfig(prompt_from_task=True),
            assets=AssetsConfig(assets_dir="checkpoints/torch/pi0_behavior/assets"),
            extra_delta_transform=True,
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader(
            "checkpoints/jax/pi0_base/params"
        ),
        pytorch_weight_path="checkpoints/torch/pi0_base",
        num_train_steps=30_000,
    ),
    TrainConfig(
        name="pi05_behavior",
        model=pi0_config.Pi0Config(pi05=True, action_horizon=32),
        data=LeRobotBehaviorDataConfig(
            repo_id="physical-intelligence/behavior",
            base_config=DataConfig(prompt_from_task=True),
            assets=AssetsConfig(assets_dir="checkpoints/torch/pi05_behavior/assets"),
            extra_delta_transform=False,
            extract_state_from_proprio=True,
            use_all_wrist_images=True,
            use_quantile_norm=True,
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader(
            "checkpoints/jax/pi0_base/params"
        ),
        pytorch_weight_path="checkpoints/torch/pi0_base",
        num_train_steps=30_000,
    ),
    TrainConfig(
        name="pi05_gsenv",
        model=pi0_config.Pi0Config(
            pi05=True, action_horizon=5, discrete_state_input=False
        ),
        data=LeRobotGSEnvDataConfig(
            repo_id="RLinf/GSEnv-PutCubeOnPlate-v0",
            base_config=DataConfig(prompt_from_task=True),
            assets=AssetsConfig(assets_dir="checkpoints/torch/pi0_r2s2r/assets"),
            extra_delta_transform=False,
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader(
            "checkpoints/jax/pi05_base/params"
        ),
        pytorch_weight_path="checkpoints/torch/pi05_base",
        num_train_steps=30_000,
    ),
    TrainConfig(
        name="pi0_realworld",
        model=pi0_config.Pi0Config(action_horizon=10),
        data=LeRobotRealworldDataConfig(
            repo_id="realworld_franka_bin_relocation",
            base_config=DataConfig(prompt_from_task=True),
            assets=AssetsConfig(assets_dir="checkpoints/torch/pi0_base/assets"),
            extra_delta_transform=False,
        ),
        pytorch_weight_path="checkpoints/torch/pi0_base",
    ),
    TrainConfig(
        name="pi0_custom",
        model=pi0_config.Pi0Config(),
        data=CustomDataConfig(
            repo_id="physical-intelligence/custom_dataset",
            base_config=DataConfig(
                prompt_from_task=True
            ),  # we need language instruction
            assets=AssetsConfig(assets_dir="checkpoints/torch/pi0_base/assets"),
            extra_delta_transform=False,  # True for delta action, False for abs_action
            action_train_with_rotation_6d=False,  # User can add extra config in custom dataset
        ),
        pytorch_weight_path="checkpoints/torch/pi0_base",
    ),
    TrainConfig(
        name="pi05_isaaclab_stack_cube",
        model=pi0_config.Pi0Config(
            pi05=True, action_horizon=10, discrete_state_input=False
        ),
        data=LeRobotIsaacLabStackCubeDataConfig(
            repo_id="RLinf/IsaacLab-Stack-Cube-Data",
            base_config=DataConfig(prompt_from_task=False),
            assets=AssetsConfig(assets_dir="checkpoints/torch/pi0_isaaclab/assets"),
        ),
        weight_loader=weight_loaders.CheckpointWeightLoader(
            "checkpoints/jax/pi05_base/params"
        ),
        pytorch_weight_path="checkpoints/torch/pi05_base",
    ),
    TrainConfig(
        name="pi05_dualfranka_tcp_rot6d",
        model=pi0_config.Pi0Config(
            pi05=True, action_horizon=20, discrete_state_input=False
        ),
        data=DualFrankaTcpRot6dDataConfig(
            repo_id="",
            base_config=DataConfig(prompt_from_task=True),
            assets=AssetsConfig(assets_dir="checkpoints/torch/pi05_base/assets"),
            extra_delta_transform=True,
        ),
        pytorch_weight_path="checkpoints/torch/pi05_base",
    ),
    TrainConfig(
        name="pi0_droid_polaris",
        model=pi0_config.Pi0Config(
            action_horizon=10,
            max_token_len=48,
        ),
        data=LeRobotPolarisDroidDataConfig(
            repo_id="physical-intelligence/droid",
            base_config=DataConfig(prompt_from_task=True),
            assets=AssetsConfig(asset_id="assets/droid"),
        ),
        pytorch_weight_path="checkpoints/torch/pi0_droid_polaris",
    ),
    TrainConfig(
        name="pi05_droid_polaris",
        model=pi0_config.Pi0Config(
            action_horizon=15,
            pi05=True,
            max_token_len=200,
        ),
        data=LeRobotPolarisDroidDataConfig(
            repo_id="physical-intelligence/droid",
            base_config=DataConfig(prompt_from_task=True),
            assets=AssetsConfig(asset_id="assets/droid"),
        ),
        pytorch_weight_path="checkpoints/torch/pi05_droid_polaris",
    ),
]

if len({config.name for config in _CONFIGS}) != len(_CONFIGS):
    raise ValueError("Config names must be unique.")
_CONFIGS_DICT = {config.name: config for config in _CONFIGS}


def _override_with_model_path(config: TrainConfig, model_path: str) -> TrainConfig:
    """Return a copy of the config with assets/weight paths set from model_path."""
    data_config = config.data
    if (
        dataclasses.is_dataclass(data_config)
        and hasattr(data_config, "assets")
        and dataclasses.is_dataclass(data_config.assets)
    ):
        data_config = dataclasses.replace(
            data_config,
            assets=dataclasses.replace(data_config.assets, assets_dir=model_path),
        )

    replace_kwargs = {
        "data": data_config,
        "pytorch_weight_path": model_path,
    }
    if dataclasses.is_dataclass(config) and any(
        field.name == "assets_dirs" for field in dataclasses.fields(config)
    ):
        replace_kwargs["assets_dirs"] = model_path

    return dataclasses.replace(config, **replace_kwargs)


def _override_with_data_kwargs(config: TrainConfig, data_kwargs: dict) -> TrainConfig:
    """Return a copy of the config with data_config set from openpi_data."""
    data_config = dataclasses.replace(config.data, **data_kwargs)
    replace_kwargs = {"data": data_config}
    return dataclasses.replace(config, **replace_kwargs)


def get_openpi_config(
    config_name: str,
    model_path: Optional[str] = None,
    batch_size: Optional[int] = None,
    repo_id: Optional[str] = None,
    data_kwargs: Optional[dict] = None,
) -> TrainConfig:
    """Get a config by name.

    Args:
        config_name: Name of the config to load.
        model_path: Optional path to override model weights and assets.
        batch_size: Optional batch size override.
        repo_id: Optional LeRobot repo_id or local data path to override.
            When using a local path, the original asset_id is preserved so
            that norm_stats can still be loaded from the model checkpoint.
    """
    if config_name not in _CONFIGS_DICT:
        closest = difflib.get_close_matches(
            config_name, _CONFIGS_DICT.keys(), n=1, cutoff=0.0
        )
        closest_str = f" Did you mean '{closest[0]}'? " if closest else ""
        raise ValueError(f"Config '{config_name}' not found.{closest_str}")

    config = _CONFIGS_DICT[config_name]
    if model_path is not None:
        config = _override_with_model_path(config, model_path)
    if data_kwargs is not None:
        config = _override_with_data_kwargs(config, data_kwargs)
    if batch_size is not None:
        config = dataclasses.replace(config, batch_size=batch_size)

    if repo_id is not None:
        original_repo_id = config.data.repo_id
        new_assets = dataclasses.replace(config.data.assets, asset_id=original_repo_id)
        new_data = dataclasses.replace(config.data, repo_id=repo_id, assets=new_assets)
        config = dataclasses.replace(config, data=new_data)

    return config
