"""The project's referee. Overcooked is cooperative: a route's bot is a cook, and the season's field are the partners
it cooks with. check: with each partner, on each layout, in both seats, one episode of 400 steps; the score is the
mean reward (20 for each soup served), each partner a subtask. practice: one episode with each partner on the first
layout, a replay of each in replays/.
  python3 -I referee/play.py check|practice <route folder>
"""
import json, os, sys

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(ROOT, 'kit'))
import league  # noqa: E402
from host import Child  # 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  # noqa: E402

LAYOUTS = ['cramped_room', 'asymmetric_advantages', 'coordination_ring']
HORIZON = 400


def cook(folder):
    bot = os.path.join(league.bot_dir(folder), 'bot.py')
    if not os.path.isfile(bot):
        sys.exit(f'No bot in {folder}: add bot.py with make_agent(mdp, mlam) or a class Agent')
    return Child([sys.executable, '-s', os.path.join(ROOT, 'referee', 'child.py'), bot], cwd=os.path.dirname(bot), limit=1.0, start=60)


def episode(cooks, layout, seed):
    mdp = OvercookedGridworld.from_layout_name(layout)
    env = OvercookedEnv.from_mdp(mdp, horizon=HORIZON, info_level=0)
    for seat, c in enumerate(cooks):
        c.ask({'episode': {'layout': layout, 'seat': seat, 'seed': seed + seat, 'horizon': HORIZON}}, limit=120)
    frames, score, done = [], 0, False
    while not done:
        s = env.state.to_dict()
        frames.append({'state': s, 'score': score})
        picks = [c.ask({'state': s}) for c in cooks]
        joint = tuple(Action.INDEX_TO_ACTION[p] if isinstance(p, int) and 0 <= p < len(Action.INDEX_TO_ACTION) else Action.STAY for p in picks)
        _, r, done, _ = env.step(joint)
        score += r
    frames.append({'state': env.state.to_dict(), 'score': score})
    return score, {'game': 'overcooked', 'layout': layout, 'terrain': mdp.terrain_mtx, 'horizon': HORIZON, 'score': score, 'frames': frames}


def main(mode, route):
    me = cook(route)
    name = league.route_name()
    reps = league.Replays() if mode == 'practice' else None
    subtasks = {}
    for partner in league.field(os.path.join(ROOT, 'field')):
        them = cook(os.path.join(ROOT, 'field', partner))
        got = []
        for k, layout in enumerate(LAYOUTS if mode == 'check' else LAYOUTS[:1]):
            for seat in (0, 1) if mode == 'check' else (0,):
                cooks = [me, them] if seat == 0 else [them, me]
                score, replay = episode(cooks, layout, 1000 * k + seat)
                got.append(score)
                if reps:
                    who = [name, partner] if seat == 0 else [partner, name]
                    replay['players'] = who
                    reps.add(f'{partner}-{layout}.json', replay, f'{who[0]} (blue) and {who[1]} (green) in {layout.replace("_", " ")}', who, f'{int(score // 20)} soups served, {int(score)} points')
        them.close()
        subtasks[partner] = sum(got) / len(got)
    me.close()
    league.write_metrics(sum(subtasks.values()) / len(subtasks), subtasks)


if __name__ == '__main__':
    main(sys.argv[1], sys.argv[2])