Tic-Tac-Toe AI

ttt_ai.py

#!/usr/bin/env python
"""
Tic-Tac-Toe

2 players 'a' and 'b' are playing the game. Player 'a' always moves first.
board_state is a 9-character string (row-major), each character representing who marked that cell.
'0' means that the cell is unmarked.

outcome is a one-character string containing 'w', 'l' or '0'.
"""

from __future__ import print_function
import sys
import random
from collections import defaultdict

try:
    from typing import Dict, Generator, Iterable, List, Optional, Tuple, Union
except ImportError:
    pass

if sys.version_info[0] == 2:
    input = raw_input # type: ignore

def validate_board(board):
    # type: (str) -> Optional[str]
    """
    Validates the board string and returns the current player
    If the board string is invalid, None will be returned.
    """
    if len(board) != 9:
        return None
    counta = board.count('a')
    countb = board.count('b')
    count0 = board.count('0')
    if counta + countb + count0 != 9 or counta - countb not in (0, 1):
        return None
    current_player = 'a' if counta == countb else 'b'
    return current_player

def get_winner(raw_board):
    # type: (Union[str, Iterable[str]]) -> Optional[str]
    if isinstance(raw_board, str):
        board = raw_board
    else:
        board = "".join(raw_board)
    rows = [board[0:3], board[3:6], board[6:9]]
    cols = []
    for j in range(3):
        cols.append("".join([board[j], board[3+j], board[6+j]]))
    diagonal1 = "".join([board[0], board[4], board[8]])
    diagonal2 = "".join([board[2], board[4], board[6]])
    winner = None
    for player in ('a', 'b'):
        if player * 3 in rows + cols + [diagonal1, diagonal2]:
            winner = player
    if winner is not None:
        return winner

    if board.count('0') == 0:
        return 't'
    else:
        return None

outcome_map = {} # type: Dict[str, str]
moves_map = defaultdict(list) # type: Dict[str, List[int]]

def next_moves(board, current_player):
    # type: (str, str) -> Generator[Tuple[int, str], None, None]
    zeros = [i for i, ch in enumerate(board) if ch == '0']
    for z in zeros:
         new_board = board[:z] + current_player + board[z+1:]
         yield (z, new_board)

def get_keys(d, x):
    # type: (Dict[int, str], str) -> Generator[int, None, None]
    for k, v in d.items():
        if v == x:
            yield k

def analyze_board(board):
    # type: (str) -> str
    """
    Analyzes a board and returns an outcome
    Also memoizes the outcome in outcome_map and best possible moves in moves_map
    """

    # return if already computed
    if board in outcome_map:
        return outcome_map[board]

    # check if input is valid
    current_player = validate_board(board)
    assert current_player is not None

    # Check if someone has won or there is a tie
    winner = get_winner(board)
    if winner is not None:
        assert winner != current_player
        if winner == 't':
            outcome_map[board] = 't'
        else:
            outcome_map[board] = 'l'
        return outcome_map[board]

    # recurse
    results = {space: analyze_board(new_board) for space, new_board in next_moves(board, current_player)}
    if 'l' in results.values():
        outcome_map[board] = 'w'
        moves_map[board] = list(get_keys(results, 'l'))
    elif 't' in results.values():
        outcome_map[board] = 't'
        moves_map[board] = list(get_keys(results, 't'))
    else:
        outcome_map[board] = 'l'
        moves_map[board] = list(results.keys())
    return outcome_map[board]

def print_board(b):
    # type: (Iterable[str]) -> None
    b = "".join(b)
    hstr = "+---+---+---+"
    b = b.replace('a', 'X').replace('b', 'O').replace('0', ' ')
    print()
    print(hstr)
    for i in range(3):
        print('|', b[i*3], '|', b[i*3+1], '|', b[i*3+2], '|')
        print(hstr)
    print()

def get_valid_move(board, message="Enter move position: "):
    # type: (Iterable[str], str) -> int
    zeros = [i for i, ch in enumerate(board) if ch == '0']
    if len(zeros) == 1:
        return zeros[0]
    while True:
        n = int(input(message))
        if n in zeros:
            break
        else:
            print("Invalid choice - choose again")
    return n

intro = """
The first player uses 'X' and the second player uses 'O'.

Cells are numbered like this:
0 1 2
3 4 5
6 7 8

In each move, you have to enter the number of the cell you wish to mark X.
"""

def main():
    # type: () -> None
    print(intro)

    human_first = input("Do you want to move first? (y/n) ").lower() == 'y'
    if human_first:
        bot_name = 'b'
        human_name = 'a'
    else:
        bot_name = 'a'
        human_name = 'b'

    winner = None
    board = ['0'] * 9
    current_player = 'a'
    while winner is None:
        if current_player == bot_name:
            board_str = "".join(board)
            analyze_board(board_str)
            moves = moves_map[board_str]
            if outcome_map[board_str] == 'w':
                print("Now bot will surely win")
            if outcome_map[board_str] == 'l':
                print("The bot is feeling nervous")
            n = random.choice(moves)
        else:
            print_board(board)
            n = get_valid_move(board)
        board[n] = current_player
        winner = get_winner(board)
        current_player = 'a' if current_player == 'b' else 'b'

    print_board(board)
    if winner == 't':
        print("Tie")
    elif winner == human_name:
        print("You won!")
    elif winner == bot_name:
        print("Bot won")

if __name__ == "__main__":
    main()