Simple tic-tac-toe AI implementation using statistics-based weights on grids for 3x3 board

tictactoe-ai.py

from random import randrange

# Set of grids
G = {1, 2, 3, 4, 5, 6, 7, 8, 9}
# List of win conditions
W = [{1, 2, 3}, {4, 5, 6}, {7, 8, 9}, {1, 4, 7}, {2, 5, 8}, {3, 6, 9}, {1, 5, 9}, {3, 5, 7}]

# Declare empty sets for players' initial occupied grids
A = set([])
B = set([])

# Initial strategic weights of grids based on how many ways it can be used to win the game
weights = {
  1: 3/8,
  2: 2/8,
  3: 3/8,
  4: 2/8,
  5: 4/8,
  6: 2/8,
  7: 3/8,
  8: 2/8,
  9: 3/8
}

def human_move():
  grid = 0
  while grid not in (G - A - B):
    grid = int(input('Choose grid:'))
  A.add(grid)
  update_weights()

def ai_move():
  print('AI move')
  grid = 0
  curr_weight = -1
  for k in iter(weights):
    if weights[k] > curr_weight:
      grid = k
      curr_weight = weights[k]
  B.add(grid)
  update_weights()

def update_weights():
  # Ignore occupied grids in future moves
  for k in iter(A):
    weights[k] = -1
  for k in iter(B):
    weights[k] = -1

  # Adjust weights according to current available win condition to AI
  available_W = list(W)
  for k in iter(A):
    for w in available_W:
      if k in w:
        available_W.remove(w)
  for k in iter(G - A - B):
    scenario_count = 0
    for w in available_W:
      if k in w:
        scenario_count += 1
    if len(available_W) > 0:
      weights[k] = scenario_count/len(available_W)
    else:
      weights[k] = 0.5

  # Aim for winning move if just one step away
  for w in W:
    intersect_set = w - B
    intersect = list(intersect_set)
    if len(intersect) == 1:
      if intersect[0] not in A:
        weights[intersect[0]] = 1
        return
  # Prevent opponent from winning
  for w in W:
    intersect_set = w - A
    intersect = list(intersect_set)
    if len(intersect) == 1:
      if intersect[0] not in B:
        weights[intersect[0]] = 1
        return

def display_grid(grid_no):
  if int(grid_no) in A:
    return 'X'
  if int(grid_no) in B:
    return 'O'
  return str(grid_no)

def draw_grids():
  print('{} {} {}'.format(display_grid(1), display_grid(2), display_grid(3)))
  print('{} {} {}'.format(display_grid(4), display_grid(5), display_grid(6)))
  print('{} {} {}'.format(display_grid(7), display_grid(8), display_grid(9)))

def game_end():
  for w in W:
    if A.union(w) == A:
      print('Player win')
      return True
    if B.union(w) == B:
      print('AI win')
      return True

  if len(G - A - B) == 0:
    print('Draw')
    return True

  return False

def main():
  whose_move = randrange(2)

  draw_grids()
  while (not game_end()):
    if whose_move == 0:
      human_move()
    else:
      ai_move()
    draw_grids()

    whose_move += 1
    whose_move = whose_move % 2

main()