from enum import Enum
import math
import random

# Tokens on the board are either Noughts, Crosses, or Empty
class Token(Enum):
    Nought = 'O'
    Cross  = 'X'
    Empty  = ' '

# make an empty n * n board
# represented internally as an n*n list
def mk_n_board(n):
    return [Token.Empty for x in range(n*n)]

# inserts Token t at position x in board
# returns None if position is not Empty
# returns None if index is invalid (negative/not in list)
def insert_at(t, x, board):
	elem = board[x] if -1 < x < len(board) else None
	return board[0:x] + [t] + board[x+1:] if elem is Token.Empty else None

# assume all boards are square, and hence we can grab
# the board length and take the sqrt to get n 
def get_n(board):
    return int(math.sqrt(len(board)))

# Gets the rows of n*n board
# assumes all boards are square
def get_rows(board):
	n = get_n(board)
	return [[t for t in board[x:x+n]]
	    for x in range(0, len(board), n)]

# Gets the columns of n*n board
# assumes all boards are square
def get_cols(board):
	return list(map(list, zip(*(get_rows(board)))))

# Gets the diagonals on a board
def get_diags(board):
    n = get_n(board)
    row1 = [get_rows(board)[x][x] for x in range(0,n)]
    row2 = [get_rows(board)[::-1][x][x] for x in range(0,n)]
    return [row1,row2]

# prettify n*n board
def pretty_n_board(board):
	n = get_n(board)
	rows = ["|".join([" %s " % t.value for t in board[x:x+n]])
	    for x in range(0, len(board), n)]
	rows = ["".join(("\n",row,"\n")) for row in rows]
	col = "___ "*n
	return col.join(rows)

# returns True if player wins on n*n board
def winner(board,player):
	rows = get_rows(board)
	cols = get_cols(board)
	diags = get_diags(board)
	def win_check(x, ys):
		return any(all(x is elem for elem in y) for y in ys)
	def check_all(x):
		return win_check(x, rows) or win_check(x, cols) or win_check(x, diags)

	return check_all(player)

# if player is X return O, if player is O return X
# otherwise return Empty
def flip(player):
	if player is Token.Cross:
		return Token.Nought
	elif player is Token.Nought:
		return Token.Cross
	else:
		return Token.Empty

# returns indexes of empty positions in board
def empties(board):
	return [x for cell,x in zip(board,range(0, len(board))) if cell == Token.Empty]

# returns True if the board is draw/full
def is_draw(board):
	return len(empties(board)) == 0

# chaotic player makes a random move on board in one of the empty cells
# returns the cell with a cell filled
# undefined for a full board
def chaos_ai(board, player):
	cell = random.choice(empties(board))
	return insert_at(player, cell, board)

# returns a list of all the boards the player can reach in one valid move
def next_boards(board, player):
	indices = empties(board)
	return [insert_at(player,index,board) for index in indices]

# defines the utility of the board, relative to the maximizing player
# for terminal boards only - those that are winning or drawing boards
def utility(board, max_player):
	num_empties = len(empties(board))
	if winner(board,max_player):
		return 1+num_empties
	elif winner(board, flip(max_player)):
		return -1-num_empties
	else:
		return 0 # draw

# Minimax algorithm, maximizing for X
def minimax(board, cur_player):
	# assume X is maximizing player
	max_player = Token.Cross
	if winner(board, Token.Cross) or winner(board, Token.Nought):
		return (board, [], utility(board, max_player))
	elif is_draw(board): # is it a draw
		return (board, [], 0)
	else:
		moves = [minimax(board, flip(cur_player)) for board in next_boards(board,cur_player)]
		choose = max if cur_player == Token.Cross else min
		best_value = choose([x for _,_,x in moves])
		best_moves = [(move,subtree,x) for (move,subtree,x) in moves if x is best_value]
		return (board,best_moves,best_value)

# Unbeatable AI takes the board, and which player it is, and returns
# the optimal move it can make from a random set of optimal moves
# uses minimax under the hood
def unbeatable_ai(board, player):
	_,t,_ = minimax(board, player)
	return random.choice(t)[0]

# Choice between the random AI and minimax AI
def choose_ai(x):
	if x == 0:
		return chaos_ai
	elif x == 1:
		return unbeatable_ai
	else:
		return None

# Prints possible game modes
def print_game_mode():
	print("Choose between the following:")
	print("0) chaotic AI")
	print("1) unbeatable AI")

# attempt to safely cast to a type, or return None on failure
def safe_cast_type(val, ty):
    try:
        return ty(val)
    except (ValueError, TypeError):
        return None

# parses the input player, if X/x then return Cross
# if O/o then return Nought
# otherwise fails and tries again
def parse_player():
	while True:
		s = safe_cast_type(input("Choose a player (X/O): "),str)
		if s == "X" or s == "x":
			return Token.Cross
		elif s == "O" or s == "o":
			return Token.Nought
		else:
			print("X or O expected, try again")

# parses the board size, should be an integer > 1
# loops and tries again on bad parse
# returns n
def parse_n_board():
	while True:
		n = safe_cast_type(input("Choose the size of the board: "),int)
		if n is None or n < 2: 
			print("integer n > 1 expected, try again")
		else:
			return n

# parses the game mode, should be an integer > 0
# loops and tries again on bad parse
# returns mode
def parse_game_mode():
	while True:
		mode = safe_cast_type(input("Input a number: "), int)
		if mode is None or choose_ai(mode) is None:
			print("0 or 1 expected, try again")
		else:
			return mode

# parses the cell to fill, should be an integer between 0 and n*n-1
# loops and tries again on bad parse
# returns cell
def parse_cell(board):
	valid_cells = empties(board)
	while True:
		cell = safe_cast_type(input("Choose a cell to fill: "), int)
		if cell is None or cell not in valid_cells:
			print("Invalid cell, choose one of", valid_cells, "expected, try again")
		else:
			return cell

# Main Game loop - what runs
def game_loop():
	print("Welcome to tic-tac-toe!")
	player = parse_player()
	adversary = flip(player)
	n = parse_n_board()
	board = mk_n_board(n)
	print_game_mode()
	ai = choose_ai(parse_game_mode())
	turn = Token.Nought
	print()
	print("Nought goes first")
	while True:
		if winner(board, player):
			print("You win!")
			break
		elif winner(board, adversary):
			print("You lose!")
			break
		elif is_draw(board):
			print("It's a Draw!")
			break
		if turn is player:
			cell = parse_cell(board)
			board = insert_at(player, cell, board)
		else:
			print("AI: Hmm, thinking...")
			board = ai(board, adversary)
		print(pretty_n_board(board))
		turn = flip(turn)

# Some tests - incomplete
def test():
	print(pretty_n_board(mk_n_board(3)))
	print(pretty_n_board(insert_at(Token.Nought,1,mk_n_board(3))))
	print(winner([Token.Nought,Token.Nought,Token.Nought,
		          Token.Empty, Token.Empty, Token.Empty,
		          Token.Empty, Token.Empty, Token.Empty], Token.Nought))
	print(winner([Token.Nought,Token.Empty,Token.Empty,
		          Token.Nought, Token.Empty, Token.Empty,
		          Token.Nought, Token.Empty, Token.Empty], Token.Nought))
	print(winner([Token.Nought,Token.Empty,Token.Empty,
		          Token.Empty, Token.Nought, Token.Empty,
		          Token.Nought, Token.Empty, Token.Nought], Token.Nought))
	print(winner([Token.Empty,Token.Empty,Token.Nought,
		          Token.Empty, Token.Nought, Token.Empty,
		          Token.Nought, Token.Empty, Token.Nought], Token.Nought))

# Entry point
game_loop()