import numpy as np
import random


# Custom functions
def menu():
    global computer
    end = False
    while not end:
        print "Tic Tac Toe\n1. New game PvP\n2. New game PvC\n0. Exit"
        option = input("Select an option: ")
        while not(option in range(0, 3)):
            option = input("Select a correct option: ")
        if option == 1:
            computer = False
            end = setup_game()
        if option == 2:
            computer = True
            end = setup_game()
        if option == 0:
            end = True
    return 0


def setup_game():
    print "Setting up new game..."
    #grid = np.zeros((3, 3))
    grid = np.chararray((3, 3))
    grid[:] = '-'
    turn = random.randint(1, 2)
    while turn != 0:
        print "\nPlayer " + str(turn) + "' turn!"
        turn = turn_game(grid, turn)
    global winner
    if winner == 0:
        print "Draw :("
    else:
        draw_game(grid)
        print "\n/////  Winner Player " + str(winner) + "  /////\n"
        answer = raw_input("Play again? (y/n):  ").lower()
        while answer != "y" and answer != 'n':
            print "Select a correct answer."
            answer = raw_input("Play again? (y/n):  ").lower()
        if answer == 'n':
            return True
        return False


def turn_game(grid, player):
    finish = False
    while not finish:
        draw_game(grid)
        if (player == 2 and not computer) or player == 1:
            column = input("Select a column: ") - 1
            row = input("Select a row: ") - 1
            if column in xrange(0, 3) and row in xrange(0, 3):
                # Check if position is available
                if grid[row, column] == '-':
                    grid[row, column] = player_char(player)
                    finish = True
                else:
                    print "Move not possible, try again."
            else:
                print "Impossible move, try again."
        else:
            # Computer makes a move
            finish = computer_turn(grid)
    player = check_game(grid, player)
    return player


def computer_turn(grid):
    move = False
    column = 0
    row = 0
    while not move:
        column = random.randint(0, 2)
        row = random.randint(0, 2)
        if grid[row, column] == "-":
            move = True
    print "Computer movement."
    grid[row, column] = "O"
    return True


def check_game(grid, player):
    win = False
    # Get width and height
    w = grid.shape[0]
    h = grid.shape[1]
    # Horizontal check
    for i in xrange(h):
        count = 0
        for j in xrange(w):
            if grid[i, j] == player_char(player):
                count += 1
        if count >= 3:
            win = True
    # Vertical check
    for j in xrange(w):
        count = 0
        for i in xrange(h):
            if grid[i, j] == player_char(player):
                count += 1
        if count >= 3:
            win = True
    # Diagonal check
    count = 0
    for i in xrange(h):
        if grid[i, i] == player_char(player):
            count += 1
        if count >= 3:
            win = True
    count = 0
    for i in xrange(h):
        if grid[abs(i - 2), i] == player_char(player):
            count += 1
        if count >= 3:
            win = True

    # Change player or set winner
    if not win:
        if player == 1:
            player = 2
        else:
            player = 1
    else:
        global winner
        winner = player
        player = 0
    return player


def player_char(player):
    if player == 1:
        return "X"
    else:
        return "O"


def draw_game(grid):
    # Get width and height
    w = grid.shape[0]
    h = grid.shape[1]
    # Draw grid
    for i in xrange(h):
        row = "\t"
        for j in xrange(w):
            row += grid[i, j] + "\t"
        print row

# Main
winner = 0
computer = False
menu()
print "Program ended."