from .grasp import GraspStage, PickStage, HookStage
from .place import PlaceStage
from .insert import InsertStage, HitStage
from .slide import SlideStage
from .pour import PourStage
from .pull import PullStage
from .push import PushStage
from .pull_revolute import PullRevoluteStage
from .push_revolute import PushRevoluteStage
from .press import PressPrismaticStage
from .twist import TwistStage

ACTION_STAGE = {
    "grasp": GraspStage,
    "pick": PickStage,
    "hook": HookStage,
    "place": PlaceStage,
    "insert": InsertStage,
    "slide": SlideStage,
    "shave": NotImplemented,
    "brush": NotImplemented,
    "wipe": NotImplemented,
    "hit": NotImplemented,
    "pour": PourStage,
    "push": PushStage,
    'click': PressPrismaticStage,
    'touch': PressPrismaticStage,
    "pull": PullStage,
    # ---- Articulate ----
    "pull_revolute": PullRevoluteStage,
    "push_revolute": PushRevoluteStage,
    "pull_prismatic":PullStage,
    "push_prismatic":PushStage,
    "twist": TwistStage,
    'press_prismatic': PressPrismaticStage,
}

def build_stage(action):
    if action not in ACTION_STAGE:
        raise NotImplementedError
    return ACTION_STAGE[action]