from functools import wraps
from enum import Enum


class Action(Enum):
    Idle = 0
    Left = 1
    TopLeft = 2
    Top = 3
    TopRight = 4
    Right = 5
    BottomRight = 6
    Bottom = 7
    BottomLeft = 8
    LongPass= 9
    HighPass = 10
    ShortPass = 11
    Shot = 12
    Sprint = 13
    ReleaseDirection = 14
    ReleaseSprint = 15
    Slide = 16
    Dribble = 17
    ReleaseDribble = 18


sticky_index_to_action = [
    Action.Left,
    Action.TopLeft,
    Action.Top,
    Action.TopRight,
    Action.Right,
    Action.BottomRight,
    Action.Bottom,
    Action.BottomLeft,
    Action.Sprint,
    Action.Dribble
]


class PlayerRole(Enum):
    GoalKeeper = 0
    CenterBack = 1
    LeftBack = 2
    RightBack = 3
    DefenceMidfield = 4
    CentralMidfield = 5
    LeftMidfield = 6
    RIghtMidfield = 7
    AttackMidfield = 8
    CentralFront = 9


class GameMode(Enum):
    Normal = 0
    KickOff = 1
    GoalKick = 2
    FreeKick = 3
    Corner = 4
    ThrowIn = 5
    Penalty = 6


def human_readable_agent(agent):
    """
    Decorator allowing for more human-friendly implementation of the agent function.

    @human_readable_agent
    def my_agent(obs):
        ...
        return football_action_set.action_right
    """
    @wraps(agent)
    def agent_wrapper(obs):
        # Extract observations for the first (and only) player we control.
        obs = obs['players_raw'][0]
        # Turn 'sticky_actions' into a set of active actions (strongly typed).
        obs['sticky_actions'] = { sticky_index_to_action[nr] for nr, action in enumerate(obs['sticky_actions']) if action }
        # Turn 'game_mode' into an enum.
        obs['game_mode'] = GameMode(obs['game_mode'])
        # In case of single agent mode, 'designated' is always equal to 'active'.
        if 'designated' in obs:
            del obs['designated']
        # Conver players' roles to enum.
        obs['left_team_roles'] = { PlayerRole(role) for role in obs['left_team_roles'] }
        obs['right_team_roles'] = { PlayerRole(role) for role in obs['right_team_roles'] }

        action = agent(obs)
        return [action.value]

    return agent_wrapper

@@ -43,16 +43,37 @@
 rt Enum%0A
+from typing import *%0A
 %0A%0Aclass 
@@ -1042,16 +1042,42 @@
 nt(agent
+: Callable%5B%5BDict%5D, Action%5D
 ):%0A    %22
@@ -1322,24 +1322,37 @@
 wrapper(obs)
+ -%3E List%5Bint%5D
 :%0A        # 

Add types to Football helpers