"""A bot's own process, in Overcooked-AI's own interface. The bot file defines make_agent(mdp, mlam), returning an
overcooked_ai_py Agent, or a class Agent; each step it gets the OvercookedState, rebuilt here from the referee's line,
and answers with one of Action.ALL_ACTIONS. mlam (the medium-level action manager, which plans motions) is built when
the bot first asks for it.
  python3 -I referee/child.py <bot.py>
"""
import os, sys

sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'kit'))
import numpy as np  # noqa: E402
from host import load, serve  # noqa: E402
from overcooked_ai_py.mdp.actions import Action  # noqa: E402
from overcooked_ai_py.mdp.overcooked_env import OvercookedEnv  # noqa: E402
from overcooked_ai_py.mdp.overcooked_mdp import OvercookedGridworld, OvercookedState  # noqa: E402

make = load(sys.argv[1], 'make_agent', 'Agent')
agent = None


def handle(msg):
    global agent
    if 'episode' in msg:  # a new episode: its layout, this seat and a seed
        e = msg['episode']
        np.random.seed(e['seed'])
        mdp = OvercookedGridworld.from_layout_name(e['layout'])
        env = OvercookedEnv.from_mdp(mdp, horizon=e['horizon'], info_level=0)
        agent = make(mdp, env.mlam) if make.__name__ == 'make_agent' else make()
        agent.set_agent_index(e['seat'])
        agent.set_mdp(mdp)
        return 'ready'
    a, _ = agent.action(OvercookedState.from_dict(msg['state']))
    return Action.ACTION_TO_INDEX[a]


serve(handle)