from tugo.rl import experience
from tugo.game_logic import scoring
from tugo.game_logic import goboard_fast as goboard
from tugo.game_logic.gotypes import Player

from collections import namedtuple


class GameRecord(namedtuple('GameRecord', 'moves winner margin')):
    pass


def simulate_game(black_player, white_player):
    moves = []
    game = goboard.GameState.new_game(19)
    agents = {
        Player.black: black_player,
        Player.white: white_player,
    }
    while not game.is_over():
        next_move = agents[game.next_player].select_move(game)
        moves.append(next_move)
        game = game.apply_move(next_move)

    game_result = scoring.compute_game_result(game)
    # print(game_result)

    return GameRecord(
        moves=moves,
        winner=game_result.winner,
        margin=game_result.winning_margin,
    )


def experience_simulation(num_games, agent1, agent2, progress_callback=None):
    collectors = [experience.ExperienceCollector() for _ in range(2)]
    agents = [agent1, agent2]

    color1 = Player.black
    for i in range(num_games):
        for collector, agent in zip(collectors, agents):
            collector.begin_episode()
            agent.set_collector(collector)

        if color1 == Player.black:
            black_player, white_player = agent1, agent2
        else:
            white_player, black_player = agent2, agent1
        game_record = simulate_game(black_player, white_player)

        winner_reward = 1 if game_record.winner == color1 else -1
        for collector in collectors:
            collector.complete_episode(reward=winner_reward)
            winner_reward *= -1

        color1 = color1.other

        if progress_callback:
            progress_callback()

    return experience.combine_experience(collectors)