"""Public API for the so101-nexus-core package."""

from pathlib import Path

from so101_nexus_core.config import (
    DIRECTION_VECTORS as DIRECTION_VECTORS,
)
from so101_nexus_core.config import (
    EXTENDED_POSE as EXTENDED_POSE,
)
from so101_nexus_core.config import (
    POSES as POSES,
)
from so101_nexus_core.config import (
    REST_POSE as REST_POSE,
)
from so101_nexus_core.config import (
    ROBOT_CAMERA_PRESETS as ROBOT_CAMERA_PRESETS,
)
from so101_nexus_core.config import (
    SO101_JOINT_NAMES as SO101_JOINT_NAMES,
)
from so101_nexus_core.config import (
    ControlMode as ControlMode,
)
from so101_nexus_core.config import (
    EnvironmentConfig as EnvironmentConfig,
)
from so101_nexus_core.config import (
    LookAtConfig as LookAtConfig,
)
from so101_nexus_core.config import (
    MoveConfig as MoveConfig,
)
from so101_nexus_core.config import (
    MoveDirection as MoveDirection,
)
from so101_nexus_core.config import (
    ObsMode as ObsMode,
)
from so101_nexus_core.config import (
    PickAndPlaceConfig as PickAndPlaceConfig,
)
from so101_nexus_core.config import (
    PickConfig as PickConfig,
)
from so101_nexus_core.config import (
    Pose as Pose,
)
from so101_nexus_core.config import (
    ReachConfig as ReachConfig,
)
from so101_nexus_core.config import (
    RenderConfig as RenderConfig,
)
from so101_nexus_core.config import (
    RewardConfig as RewardConfig,
)
from so101_nexus_core.config import (
    RobotCameraPreset as RobotCameraPreset,
)
from so101_nexus_core.config import (
    RobotConfig as RobotConfig,
)
from so101_nexus_core.config import (
    YcbModelId as YcbModelId,
)
from so101_nexus_core.constants import (
    COLOR_MAP as COLOR_MAP,
)
from so101_nexus_core.constants import (
    YCB_OBJECTS as YCB_OBJECTS,
)
from so101_nexus_core.constants import (
    ColorConfig as ColorConfig,
)
from so101_nexus_core.constants import (
    ColorName as ColorName,
)
from so101_nexus_core.constants import (
    sample_color as sample_color,
)
from so101_nexus_core.objects import (  # noqa: F401
    CubeObject,
    MeshObject,
    SceneObject,
    YCBObject,
)
from so101_nexus_core.observations import (  # noqa: F401
    EndEffectorPose,
    GazeDirection,
    GraspState,
    JointPositions,
    ObjectOffset,
    ObjectPose,
    Observation,
    OverheadCamera,
    TargetOffset,
    TargetPosition,
    WristCamera,
)

ASSETS_DIR = Path(__file__).resolve().parent / "assets"
SO_ARM100_DIR = ASSETS_DIR / "SO-ARM100"
SO101_DIR = ASSETS_DIR / "SO101"


def get_so101_simulation_dir() -> Path:
    """Return the path to the SO101 simulation assets directory."""
    return SO101_DIR


def get_so100_simulation_dir() -> Path:
    """Return the path to the SO100 simulation assets directory."""
    return SO_ARM100_DIR / "Simulation" / "SO100"


from so101_nexus_core.rewards import (  # noqa: F401, E402
    orientation_progress,
    reach_progress,
    simple_reward,
)
from so101_nexus_core.ycb_assets import (  # noqa: E402
    ensure_ycb_assets as ensure_ycb_assets,
)
from so101_nexus_core.ycb_assets import (  # noqa: E402
    get_ycb_collision_mesh as get_ycb_collision_mesh,
)
from so101_nexus_core.ycb_assets import (  # noqa: E402
    get_ycb_mesh_dir as get_ycb_mesh_dir,
)
from so101_nexus_core.ycb_assets import (  # noqa: E402
    get_ycb_visual_mesh as get_ycb_visual_mesh,
)
from so101_nexus_core.ycb_geometry import (  # noqa: E402
    get_maniskill_ycb_spawn_z as get_maniskill_ycb_spawn_z,
)
from so101_nexus_core.ycb_geometry import (  # noqa: E402
    get_mujoco_ycb_rest_pose as get_mujoco_ycb_rest_pose,
)
