##############################
# Populate winning positions #
##############################

get_winning_combinations <- function(k = 3) {

  # Sanity check
  if (k%%1 != 0 || k < 3) stop ('k is not an integer or less than 3')

  # Create a k by k matrix
  winning_mat<-matrix(seq(1:k^2), nrow = k, byrow = TRUE)
  list_winning_combinations<-vector("list", length = 2*k +2)

  # Assign winning combinations
  for (i in 1:k){
    list_winning_combinations[[i]]<-winning_mat[i,] # Rows
    list_winning_combinations[[i+k]]<-winning_mat[,i] # Columns
  }
  list_winning_combinations[[2*k+1]]<-diag(winning_mat) # Diagonal
  list_winning_combinations[[2*k+2]]<-diag(winning_mat[nrow(winning_mat):1, ]) # Reverse-diagonal
  return(list_winning_combinations)
}

#################################################
# Display current board in a k by k matrix form #
#################################################

display_board <- function(current_board, k = 3) {

  # Sanity check
  if (k%%1 != 0 || k < 3) stop ('k is not an integer or less than 3')

  # Display the current board as a k by k matrix
  current_board[which(is.na(current_board))]<-which(is.na(current_board))
  mat<-matrix(current_board, nrow = k, byrow = TRUE)
  rownames(mat)<-colnames(mat)<-rep("", k)
  print(noquote(mat))
}

############################
# Function to update board #
############################

update_board <- function(current_board, player, position) {
  if (is.na(current_board[position])) {
    current_board[position] <- player
  } else {
    print("This position is already taken. Please choose another one!")
  }
  return(current_board)
}

#######################################
# Function to check and decide winner #
#######################################

check_winner <- function(current_board, k=3) {

  # Extract winning combinations
  winning_combinations<-get_winning_combinations(k)

  # Extract Players' moves
  x.move <- as.numeric(which(current_board == "x"))
  o.move <- as.numeric(which(current_board == "o"))

  # Match
  for (i in 1:length(winning_combinations)) {
    if (sum(winning_combinations[[i]] %in% x.move) == k) {
      print("Player 1 (x) wins!")
      return (FALSE)
    }
    if (sum(winning_combinations[[i]] %in% o.move) == k) {
      print("Player 2 (o) wins!")
      return (FALSE)
    }
    if (!(NA %in% current_board)){
      print("It's a draw!")
      return (FALSE)
    }
  }
  return(TRUE)
}

#####################################
# Check if a proposed move is valid #
#####################################

check.move <- function(position) {
if (!grepl("^[1-9]$", position)) {
  return(FALSE)
}
return(TRUE)
}

#######################################
# Check if the player number is valid #
#######################################

check.player.number <- function() {
  input <- readline(prompt = "How many players (Answer 1 or 2):")
  if (!grepl("^[1-2]", input)) {
    stop("Invalid input. Try again!")
  } else {
    return (as.integer(input))
  }
}

#############################
# Main function to play TTT #
#############################

play_TTT <- function(k = 3){

  # Initialize game conditons: 1 or 2 players.
  # If AI plays, is it player 1 or 2?
  n <- check.player.number()

  # Register names of the players
  if (n == 1) {
    if ((readline(prompt = "Do you want to play first? (Answer Y or N) ")) == "Y") {
      player1 <- readline(prompt = "Please enter your name: ")
      player2 <- "AI"
    } else if (readline(prompt = "Do you want to play first? (Answer Y or N) ") == "N") {
      player1 <- "AI"
      player2 <- readline(prompt = "Please enter your name: ")
    } else {
      stop("Invalid input. Please try again!")
    }
  } else if (n == 2){
    player1 <- readline(prompt = "Who wants to go first: ")
    player2 <- readline(prompt = "Who wants to go next: ")
  } else {
    stop("Invalid input. Please try again!")
  }

  # Initialize game board
  current_board <- rep(NA, k^2)
  repeat{

    ###################
    # Player 1's turn #
    ###################

    display_board(current_board, k = k) # Display board

    if (player1 == "AI") {  # Player 1 chooses where to play. prompt user or AI_turn()
      move <- AI_turn(current_board)
      current_board <- update_board(current_board, "x", move)
    } else {
      move <- readline(prompt = cat(player1, "Which position do you want to play: ", sep = ' '))
      if (check.move(move)) {
        move <- as.integer(move)
        current_board <- update_board(current_board, "x", move) # update board
      } else {break}
    }
    if (!check_winner(current_board)) {
      break    # if Player 1 wins - quit loop
    }

    ###################
    # Player 2's turn #
    ###################

    display_board(current_board, k = k) # Display board

    # o chooses where to play. prompt user or AI_turn()
    if (player2 == "AI") {
      move <- AI_turn(current_board)
      current_board <- update_board(current_board, "o", move)
    } else {
      move <- readline(prompt = cat(player2, "Which position do you want to play: ", sep = " "))
      if (check.move(move)) {
        move <- as.integer(move)
        current_board <- update_board(current_board, "o", move) # update board
      } else {break}
    }
    if (!check_winner(current_board)) {
      break   # if Player 2 wins - quit loop
    }
  }

  # Final Step - Display final board state and declare winner
  display_board(current_board)
}