'''
Module for the connect four game

Written by aozdemir for Foundations of Computing, Loomis Chaffee Fall 2017
'''

import multiprocessing
import time
import random

ROWS = 6
COLUMNS = 7
STREAK = 4
TIMEOUT = 1 # seconds
PLY = 3

# A board is a list of 6 rows
# A row is a list of 7 characters
# A character is 'X', 'O', or ' '

def print_board(board):
    ''' prints the board with column numbers '''
    for row in board:
        print('|%s|' % ' '.join(row))
    print('+-------------+')
    print('|0 1 2 3 4 5 6|')
    print('+-------------+')


def _column_from_user(board, piece):
    ''' Gets a move from the user that is valid for `board` '''
    print('%s to move' % piece)
    while True:
        column = raw_input('Enter a column number: ')
        if column not in ['0', '1', '2', '3', '4', '5', '6']:
            print('Invalud column number!')
            continue
        column = int(column)
        if _is_valid_move(board, column):
            return column
        else:
            print('That column is full!')

def get_move(queue, board, piece, ply, ai):
    queue.put(ai(_copy_board_as_tuple(board), piece, ply))

def _column_from_ai(board, piece, ai, ply):
    ''' Gets a move from `ai` that is valid for `board`
    `ai` should be a function that takes `board`, `piece`, and `ply`,
    and returns a column number in range(7).
    If `ai` fails to do so, or takes too long, a random (valid)
    choice will be made for it.
    '''

    queue = multiprocessing.Queue(1)
    p = multiprocessing.Process(target=get_move, args=(queue,board,piece,ply,ai))
    p.start()
    start = time.time()
    p.join(TIMEOUT)
    if p.is_alive():
        p.terminate()
        print('The AI took longer than %s seconds and was killed' % TIMEOUT)
        p.join()
        return _random_move(board, piece)
    if p.exitcode != 0:
        print('The AI threw an exception')
        return _random_move(board, piece)
    move = queue.get_nowait()
    if move not in range(COLUMNS) or not _is_valid_move(board, move):
        print_board(board)
        print('The AI made the move %s, which is invalid on the above board!'% move)
        return _random_move(board, piece)
    secs = time.time() - start
    print('The AI for %s took %.3f seconds' % (piece, secs))
    return move

def _random_move(board, piece):
    ''' returns a random valid move on `board` as a column number '''
    m = random.choice(_valid_moves(board))
    print('Making a random valid move for %s: %d' % (piece, m))
    return m

def _copy_board_as_tuple(board):
    ''' copies board '''
    return tuple(tuple(row) for row in board)

def _copy_board(board):
    ''' copies board '''
    return [list(row) for row in board]

def _valid_moves(board):
    ''' returns a list of valid moves, in the form of column numbers '''
    return [i for i in xrange(COLUMNS) if board[0][i] == ' ']

def _is_valid_move(board, column):
    ''' returns whether a piece may be dropped in `column` '''
    return board[0][column] == ' '

def _make_move_mut(board, column, piece):
    ''' drops `piece` in `column`, modifying `board`
    Does nothing for a full column.
    '''
    for i in reversed(xrange(ROWS)):
        if board[i][column] == ' ':
            board[i][column] = piece
            return

def make_move(board, column, piece):
    ''' drops `piece` in `column`, forming a new `board` '''
    b2 = _copy_board(board)
    _make_move_mut(b2, column, piece)
    return b2

def has_won(board, piece):
    ''' returns whether `piece` has won the `board` '''
    # Check column wins
    for c in xrange(COLUMNS):
        for r in xrange(ROWS - STREAK + 1):
            if all(board[r+i][c] == piece for i in xrange(STREAK)):
                return True
    # Check row wins
    for r in xrange(ROWS):
        for c in xrange(COLUMNS - STREAK + 1):
            if all(board[r][c+i] == piece for i in xrange(STREAK)):
                return True
    # Check for back-diagonal wins
    for r in xrange(ROWS - STREAK + 1):
        for c in xrange(COLUMNS - STREAK + 1):
            if all(board[r+i][c+i] == piece for i in xrange(STREAK)):
                return True
    # Check for forward-diagonal wins
    for r in xrange(STREAK - 1, ROWS):
        for c in xrange(COLUMNS - STREAK + 1):
            if all(board[r-i][c+i] == piece for i in xrange(STREAK)):
                return True
    return False

def _empty_board():
    ''' returns an empty board '''
    return [ [ ' ' for i in xrange(COLUMNS) ] for i in xrange(ROWS) ]

def next_piece(piece):
    ''' returns the piece that comes after `piece` '''
    if piece == 'X': return 'O'
    if piece == 'O': return 'X'

def human_v_human():
    ''' Plays a human v. human game of connect-4 '''
    board = _empty_board()
    piece = 'X'
    for move_number in xrange(COLUMNS * ROWS):
        print_board(board)
        print('')
        column = _column_from_user(board, piece)
        _make_move_mut(board, column, piece)
        if has_won(board, piece):
            print_board(board)
            print('\n%s won!\n' % piece)
            return
        piece = next_piece(piece)
    print_board(board)
    print('\nIt was a draw')

def human_v_ai(ai, ply=3):
    ''' Plays a human v. computer game of connect 4 where
    `ai` is the AI and it can think
    `ply` moves ahead
    '''
    board = _empty_board()
    piece = 'X'
    print('The %s AI will be playing \'O\'' % ai.__doc__.split('\n')[0].strip())
    for move_number in xrange(COLUMNS * ROWS):
        print_board(board)
        print('')
        if piece == 'X':
            column = _column_from_user(board, piece)
        else:
            column = _column_from_ai(board, piece, ai, ply)
        _make_move_mut(board, column, piece)
        if has_won(board, piece):
            print_board(board)
            print('\n%s won!\n' % piece)
            return
        piece = next_piece(piece)
    print_board(board)
    print('\nIt was a draw')

def ai_v_ai(ai1, ai2, ply=3):
    ''' Plays a computer v. computer game of connect 4 where they think
    `ply` moves ahead.

    `ai1` is first to move, as 'X'
    '''
    board = _empty_board()
    piece = 'X'
    print('The %s AI will be playing \'X\'' % ai1.__doc__.split('\n')[0].strip())
    print('The %s AI will be playing \'O\'' % ai2.__doc__.split('\n')[0].strip())
    for move_number in xrange(COLUMNS * ROWS):
        print_board(board)
        print('')
        if piece == 'X':
            column = _column_from_ai(board, piece, ai1, ply)
        else:
            column = _column_from_ai(board, piece, ai2, ply)
        time.sleep(1)
        _make_move_mut(board, column, piece)
        if has_won(board, piece):
            print_board(board)
            print('\n%s won!\n' % piece)
            return
        piece = next_piece(piece)
    print_board(board)
    print('\nIt was a draw')



BOARD_X_THREAT = _copy_board(['       ',
                            '       ',
                            '       ',
                            '   X   ',
                            '   XO  ',
                            '   XOO '])
BOARD_O_THREAT = _copy_board(['       ',
                            '       ',
                            '       ',
                            '       ',
                            '    XX ',
                            '  X OOO'])
BOARD_BOTH_THREAT = _copy_board(['       ',
                            '       ',
                            '       ',
                            '  X    ',
                            '  X    ',
                            '  X OOO'])
BOARD_X_WILL_WIN   = _copy_board(['       ',
                            '       ',
                            '       ',
                            '       ',
                            '  OOO  ',
                            '  XXX  '])
BOARD_EMPTY  = _copy_board(['       ',
                            '       ',
                            '       ',
                            '       ',
                            '       ',
                            '       '])
BOARD_X_BLOCKS_THEN_WINS = _copy_board([reversed(' OOOXX '),
                                        reversed('XOXOOO '),
                                        reversed('OXOXXX '),
                                        reversed('OXOXXX '),
                                        reversed('OOXXOO '),
                                        reversed('XOOOXOO')])

PIECES = ['X', 'O']