from __future__ import annotations

ACTIONS = ["left", "right", "claim"]


class KeyDoorCorridor:
    def __init__(self, start: int = 4, max_steps: int = 20) -> None:
        self.position = start
        self.has_key = start == 1
        self.max_steps = max_steps
        self.steps = 0
        self.done = False
        self.terminal = "none"

    @property
    def state(self) -> tuple[int, int]:
        return self.position, int(self.has_key)

    def step(self, action: int) -> float:
        if self.done:
            return 0.0
        reward = -0.01
        if action == 0:
            self.position = max(0, self.position - 1)
        elif action == 1:
            if not (self.position == 6 and not self.has_key):
                self.position = min(8, self.position + 1)
        elif action == 2 and self.position == 6:
            reward += 0.4
            self.done = True
            self.terminal = "near_reward"
        elif action == 2 and self.position == 8 and self.has_key:
            reward += 1.0
            self.done = True
            self.terminal = "treasure"
        if self.position == 1:
            self.has_key = True
        self.steps += 1
        if self.steps >= self.max_steps and not self.done:
            self.done = True
            self.terminal = "timeout"
        return reward


def expert_action(environment: KeyDoorCorridor, goal: str) -> int:
    if goal == "treasure":
        if not environment.has_key:
            return 0
        if environment.position < 8:
            return 1
        return 2
    if environment.position < 6:
        return 1
    if environment.position > 6:
        return 0
    return 2