line.py

from copy import deepcopy

RED = 1  # max
BLACK = 2  # min


class Board(object):

    def __init__(self, width=5, height=4, win_len=4, depth=2):
        self.width = width
        self.height = height
        self.win_len = win_len
        self.depth = depth

        self.grid = [[0] * width for _ in range(height)]
        self.current_player = RED
        self.last_move = (-1, -1)
        self.ended = 0

        self.red_segments = []
        self.black_segments = []

    def print_board(self):
        for row in self.grid:
            for cell in row:
                print cell,
            print
        print

    def place_in_column(self, row, col):
        if row == self.height:
            raise IndexError
        if self.grid[row][col] != 0:
            raise IndexError
        self.grid[row][col] = self.current_player
        self.last_move = (row, col)

        if self.check_win(row, col) != 0:
            self.ended = self.current_player

        self.current_player ^= 3

    def check_win(self, row, col):
        color = self.grid[row][col]
        return any([self.__num_in_direction(color, row, col, i, j) >=
                    self.win_len for (i, j) in [(1, 0), (0, 1), (1, 1), (1, -1)]])

    def __num_in_direction(self, color, row, col, deltax, deltay):
        def inner(dx, dy):
            count = 0
            x, y = row, col
            for i in range(self.width + self.height):
                x, y = x + dx, y + dy
                try:
                    assert(x >= 0 and y >= 0)
                    if self.grid[x][y] != color:
                        return count
                except IndexError:
                    return count
                except AssertionError:
                    return count
                count += 1
            return count
        return 1 + inner(deltax, deltay) + inner(-deltax, -deltay)

    def children(self):
        chiidren_boards = []
        for row in range(self.height):
            for column in range(self.width):
                chiidren_board = deepcopy(self)
                try:
                    chiidren_board.place_in_column(row, column)
                except IndexError:
                    continue
                    print "hit"
                else:
                    chiidren_boards.append(chiidren_board)
        return chiidren_boards

    def get_segments(self, row, col):
        segments = []
        rmin, cmin = max(0, row + 1 - self.win_len), max(
            0, col + 1 - self.win_len)
        rmax, cmax = min(self.height - self.win_len, row), min(
            self.width - self.win_len, col)

        # rows
        c = cmin
        while c <= cmax:
            segments.append([(row, i) for i in range(c, c + self.win_len)])
            c += 1

        # columns
        r = rmin
        while r <= rmax:
            segments.append([(i, col) for i in range(r, r + self.win_len)])
            r += 1

        # up-right diagonals
        r, c = row, col
        while r >= rmin and c >= cmin:
            seg = []
            i, j = r, c
            while i < r + self.win_len and j < c + self.win_len:
                if i in range(self.height) and j in range(self.width):
                    seg.append((i, j))
                else:
                    break
                i, j = i + 1, j + 1
            if len(seg) == self.win_len:
                segments.append(seg)
            r, c = r - 1, c - 1

        # up-left diagonals
        r, c = row, col
        while r <= rmax + self.win_len and c >= cmin:
            seg = []
            i, j = r, c
            while i > r - self.win_len and j < c + self.win_len:
                if i in range(self.height) and j in range(self.width):
                    seg.append((i, j))
                else:
                    break
                i, j = i - 1, j + 1
            if len(seg) == self.win_len:
                segments.append(seg)
            r, c = r + 1, c - 1
        return segments

    def score(self):
        s = 0
        if self.ended == RED:
            return 1000
        elif self.ended == BLACK:
            return -1000

        for row in range(self.height):
            for col in range(self.width):
                segments = self.get_segments(row, col)
                for seg in segments:
                    num_red = sum(map(lambda x: 1 if self.grid[
                                  x[0]][x[1]] == RED else 0, seg))
                    num_black = sum(map(lambda x: 1 if self.grid[
                                    x[0]][x[1]] == BLACK else 0, seg))
                    if num_red == 0 or num_black == 0:
                        s += num_red ** 2
                        s -= num_black ** 2

        return s

    def score2(self):
        s = 0
        if self.ended == RED:
            return 1000
        elif self.ended == BLACK:
            return -1000

        for row in range(self.height):
            for col in range(self.width):
                if self.grid[row][col] == RED:
                    t = self.__num_in_direction(RED, row, col, 1, 0)
                    s += t**2
                    t = self.__num_in_direction(RED, row, col, 1, 1)
                    s += t**2
                    t = self.__num_in_direction(RED, row, col, 0, 1)
                    s += t**2
                    t = self.__num_in_direction(RED, row, col, 1, -1)
                    s += t**2
                elif self.grid[row][col] == BLACK:
                    t = self.__num_in_direction(BLACK, row, col, 1, 0)
                    t -= t**2
                    t = self.__num_in_direction(BLACK, row, col, 1, 1)
                    t -= t**2
                    t = self.__num_in_direction(BLACK, row, col, 0, 1)
                    t -= t**2
                    t = self.__num_in_direction(BLACK, row, col, 1, -1)
                    t -= t**2

        return s


def alphabeta(board, depth, alpha, beta, player, move):
    if depth == 0 or board.ended != 0:
        return board.score(), move
    if player == RED:
        for child in reversed(board.children()):
            alpha, move = max((alpha, move), (alphabeta(
                child, depth - 1, alpha, beta, player ^ 3, move)[0],
                child.last_move))
            if beta <= alpha:
                break
        return alpha, move
    else:
        for child in board.children():
            beta, move = min((beta, move), (alphabeta(
                child, depth - 1, alpha, beta, player ^ 3, move)[0],
                child.last_move))
            if beta <= alpha:
                break
        return beta, move


def get_move(board):
    return alphabeta(board, board.depth, -99999999, 99999999,
                     board.current_player, (-1, -1))


def get_human_move(board):
    board.print_board()
    row, column = raw_input('Enter (column row)').split(' ')
    row, column = int(row), int(column)
    if column not in range(board.width) or row not in range(board.height):
        raise Exception
    board.place_in_column(row, column)
    board.print_board()


def get_ai_move(board):
    move = get_move(board)
    print move
    board.place_in_column(*move[1])
    print 'My move has a score of %d' % move[0]


if __name__ == "__main__":
    level = input('Type 1-3 to select level')
    board = Board(depth=int(level) + 2)
    human = input('Type 1 to go first, 2 to go second: ')
    if human not in [1, 2]:
        raise Exception
    if human == 1:
        get_human_move(board)

    while not board.ended != 0:
        get_ai_move(board)
        if board.ended > 0:
            break
        get_human_move(board)

    print 'Thank you for playing! Player %d wins.' % board.ended
    board.print_board()