connect-four-minimax.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Connect 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_COLUMNS = 7\n",
    "COLUMN_HEIGHT = 6\n",
    "FOUR = 4\n",
    "\n",
    "# Board can be initialized with `board = np.zeros((NUM_COLUMNS, COLUMN_HEIGHT), dtype=np.byte)`\n",
    "# Notez Bien: Connect 4 \"columns\" are actually NumPy \"rows\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def valid_moves(board):\n",
    "    \"\"\"Returns columns where a disc may be played\"\"\"\n",
    "    return [n for n in range(NUM_COLUMNS) if board[n, COLUMN_HEIGHT - 1] == 0]\n",
    "\n",
    "\n",
    "def play(board, column, player):\n",
    "    \"\"\"Updates `board` as `player` drops a disc in `column`\"\"\"\n",
    "    (index,) = next((i for i, v in np.ndenumerate(board[column]) if v == 0))\n",
    "    board[column, index] = player\n",
    "\n",
    "\n",
    "def take_back(board, column):\n",
    "    \"\"\"Updates `board` removing top disc from `column`\"\"\"\n",
    "    (index,) = [i for i, v in np.ndenumerate(board[column]) if v != 0][-1]\n",
    "    board[column, index] = 0\n",
    "\n",
    "\n",
    "def four_in_a_row(board, player):\n",
    "    \"\"\"Checks if `player` has a 4-piece line\"\"\"\n",
    "    return (\n",
    "        any(\n",
    "            all(board[c, r] == player)\n",
    "            for c in range(NUM_COLUMNS)\n",
    "            for r in (list(range(n, n + FOUR)) for n in range(COLUMN_HEIGHT - FOUR + 1))\n",
    "        )\n",
    "        or any(\n",
    "            all(board[c, r] == player)\n",
    "            for r in range(COLUMN_HEIGHT)\n",
    "            for c in (list(range(n, n + FOUR)) for n in range(NUM_COLUMNS - FOUR + 1))\n",
    "        )\n",
    "        or any(\n",
    "            np.all(board[diag] == player)\n",
    "            for diag in (\n",
    "                (range(ro, ro + FOUR), range(co, co + FOUR))\n",
    "                for ro in range(0, NUM_COLUMNS - FOUR + 1)\n",
    "                for co in range(0, COLUMN_HEIGHT - FOUR + 1)\n",
    "            )\n",
    "        )\n",
    "        or any(\n",
    "            np.all(board[diag] == player)\n",
    "            for diag in (\n",
    "                (range(ro, ro + FOUR), range(co + FOUR - 1, co - 1, -1))\n",
    "                for ro in range(0, NUM_COLUMNS - FOUR + 1)\n",
    "                for co in range(0, COLUMN_HEIGHT - FOUR + 1)\n",
    "            )\n",
    "        )\n",
    "    )\n",
    "\n",
    "def _mc(board, player):\n",
    "    p = -player\n",
    "    while valid_moves(board):\n",
    "        p = -p\n",
    "        c = np.random.choice(valid_moves(board))\n",
    "        play(board, c, p)\n",
    "        if four_in_a_row(board, p):\n",
    "            return p\n",
    "    return 0\n",
    "\n",
    "\n",
    "def montecarlo(board, player):\n",
    "    montecarlo_samples = 100\n",
    "    cnt = Counter(_mc(np.copy(board), player) for _ in range(montecarlo_samples))\n",
    "    return (cnt[1] - cnt[-1]) / montecarlo_samples\n",
    "\n",
    "\n",
    "def eval_board(board, player):\n",
    "    if four_in_a_row(board, 1):\n",
    "        # Alice won\n",
    "        return 1\n",
    "    elif four_in_a_row(board, -1):\n",
    "        # Bob won\n",
    "        return -1\n",
    "    else:\n",
    "        # Not terminal, let's simulate...\n",
    "        return montecarlo(board, player)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Author: Fra98\n",
    "### Four-Connect solver with MinMax (+AlphaBeta Pruning) + Montecarlo tree exploration\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "PLAYERS = {1: \"A\", -1: \"B\"}\n",
    "MAX_ROUNDS = NUM_COLUMNS * COLUMN_HEIGHT\n",
    "MAX_DEPTH = 3\n",
    "MC_ITERATIONS = 20\n",
    "SEARCH_ORDER = [3, 2, 4, 1, 5, 0, 6]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Auxiliary Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def initialize_board():\n",
    "    return np.zeros((NUM_COLUMNS, COLUMN_HEIGHT), dtype=np.byte)\n",
    "    \n",
    "def display(board):\n",
    "    for j in reversed(range(COLUMN_HEIGHT)):\n",
    "        for i in range(NUM_COLUMNS):\n",
    "            cell = board[i][j]\n",
    "            if cell == 1:\n",
    "                print(PLAYERS[1], end=\" \")\n",
    "            elif cell == -1:\n",
    "                print(PLAYERS[-1], end=\" \")\n",
    "            else:\n",
    "                print(\"-\", end=\" \")\n",
    "        print()\n",
    "\n",
    "\n",
    "def round_number(board):\n",
    "    return np.count_nonzero(board)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### MiniMax (with Pruning) + Montecarlo Simulation "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def can_win_next_move(board, player, moves=None, round_num=None):\n",
    "    if moves is None:\n",
    "        moves = valid_moves(board)\n",
    "    \n",
    "    if round_num is None:\n",
    "        round_num = round_number(board)\n",
    "\n",
    "    for m in moves:\n",
    "        play(board, m, player)\n",
    "        score = None\n",
    "        if four_in_a_row(board, player):\n",
    "            score = (MAX_ROUNDS - (round_num - 1)) // 2\n",
    "        take_back(board, m)\n",
    "        if score:\n",
    "            return score, m\n",
    "\n",
    "    return None, None\n",
    "\n",
    "\n",
    "def mc_simulation(board: np.ndarray, player: int):\n",
    "    best_score = -MAX_ROUNDS\n",
    "    best_move = None\n",
    "    \n",
    "    for _ in range(MC_ITERATIONS):\n",
    "        move = np.random.choice(valid_moves(board))\n",
    "        \n",
    "        tmp_board = np.copy(board)\n",
    "        tmp_player = player\n",
    "        tmp_move = move\n",
    "        while True:\n",
    "            play(tmp_board, tmp_move, tmp_player)\n",
    "\n",
    "            # Terminal conditions\n",
    "            if round_number(tmp_board) == MAX_ROUNDS:\n",
    "                score = 0\n",
    "                break\n",
    "            if four_in_a_row(tmp_board, tmp_player):\n",
    "                if tmp_player == player:\n",
    "                    score = (MAX_ROUNDS - (round_number(tmp_board) - 1)) // 2\n",
    "                else:\n",
    "                    score = -((MAX_ROUNDS - (round_number(tmp_board) - 1)) // 2)\n",
    "                break\n",
    "\n",
    "            tmp_move = np.random.choice(valid_moves(tmp_board))\n",
    "            tmp_player = -tmp_player\n",
    "\n",
    "        if score > best_score:\n",
    "            best_score = score\n",
    "            best_move = move\n",
    "        \n",
    "    return best_score, best_move\n",
    "\n",
    "\n",
    "def minimax(board: np.ndarray, player: int, depth: int, alpha: int, beta: int, max_depth: int = MAX_DEPTH):\n",
    "    # CHECK FOR TERMINAL CONDITIONS (DRAW OR WIN WITHIN NEXT MOVE)\n",
    "    round_num = round_number(board)\n",
    "\n",
    "    # Case draw\n",
    "    if round_num == MAX_ROUNDS:\n",
    "        return 0, None\n",
    "\n",
    "    moves = valid_moves(board)\n",
    "\n",
    "    # Check if current player can win with the next move\n",
    "    score, m = can_win_next_move(board, player, moves, round_num)\n",
    "    if score:\n",
    "        return score, m\n",
    "\n",
    "    # NON TERMINAL\n",
    "    \n",
    "    # Cap beta to maximum possible score (already considering next round -> + 1)\n",
    "    max_score = (MAX_ROUNDS - (round_num + 1)) // 2\n",
    "    if beta > max_score:\n",
    "        beta = max_score\n",
    "    \n",
    "    # Prune\n",
    "    if alpha >= beta:\n",
    "        return beta, None\n",
    "\n",
    "    # Montecarlo simulation if DEPTH too high\n",
    "    if depth > max_depth:\n",
    "        return mc_simulation(board, player)\n",
    "\n",
    "    # MinMax exploration if DEPTH is low enough\n",
    "    # start searching from middle columns because there is more probability of higher score, therefore more pruning\n",
    "    best_move = None\n",
    "    for m in SEARCH_ORDER:\n",
    "        if m in moves:\n",
    "            play(board, m, player)\n",
    "            score, _ = minimax(board, -player, depth+1, -beta, -alpha)\n",
    "            score = -score\n",
    "            take_back(board, m)\n",
    "\n",
    "            if score >= beta:\n",
    "                return score, m\n",
    "\n",
    "            if score > alpha:\n",
    "                alpha = score\n",
    "                best_move = m\n",
    "\n",
    "    return alpha, best_move"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Choosing Move"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def choose_move(board: np.ndarray, player: int, max_depth: int = MAX_DEPTH):\n",
    "    # FIRST/SECOND MOVE -> always play central\n",
    "    if not board.any() or round_number(board) == 1:\n",
    "        play(board, 3, player)\n",
    "        return 3\n",
    "\n",
    "    # THIRD MOVE -> always play in one of the 3 cells in the center\n",
    "    if round_number(board) == 2:\n",
    "        move = np.random.choice([2, 3, 4])\n",
    "        play(board, move, player)\n",
    "        return move\n",
    "\n",
    "    # MINMAX \n",
    "    _, move = minimax(board, player, depth=1, alpha=-1000, beta=1000, max_depth=max_depth)\n",
    "\n",
    "    if move is not None:\n",
    "        play(board, move, player)\n",
    "\n",
    "    return move"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulation examples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### AI vs AI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A TURN -> 4\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - A - - - \n",
      "\n",
      "B TURN -> 4\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - B - - - \n",
      "- - - A - - - \n",
      "\n",
      "A TURN -> 5\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - B - - - \n",
      "- - - A A - - \n",
      "\n",
      "B TURN -> 3\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - B - - - \n",
      "- - B A A - - \n",
      "\n",
      "A TURN -> 6\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - B - - - \n",
      "- - B A A A - \n",
      "\n",
      "B TURN -> 7\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - B - - - \n",
      "- - B A A A B \n",
      "\n",
      "A TURN -> 3\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - A B - - - \n",
      "- - B A A A B \n",
      "\n",
      "B TURN -> 4\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - B - - - \n",
      "- - A B - - - \n",
      "- - B A A A B \n",
      "\n",
      "A TURN -> 4\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - A - - - \n",
      "- - - B - - - \n",
      "- - A B - - - \n",
      "- - B A A A B \n",
      "\n",
      "B TURN -> 3\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - A - - - \n",
      "- - B B - - - \n",
      "- - A B - - - \n",
      "- - B A A A B \n",
      "\n",
      "A TURN -> 7\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - A - - - \n",
      "- - B B - - - \n",
      "- - A B - - A \n",
      "- - B A A A B \n",
      "\n",
      "B TURN -> 2\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - A - - - \n",
      "- - B B - - - \n",
      "- - A B - - A \n",
      "- B B A A A B \n",
      "\n",
      "A TURN -> 7\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - A - - - \n",
      "- - B B - - A \n",
      "- - A B - - A \n",
      "- B B A A A B \n",
      "\n",
      "B TURN -> 7\n",
      "- - - - - - - \n",
      "- - - - - - - \n",
      "- - - A - - B \n",
      "- - B B - - A \n",
      "- - A B - - A \n",
      "- B B A A A B \n",
      "\n",
      "A TURN -> 4\n",
      "- - - - - - - \n",
      "- - - A - - - \n",
      "- - - A - - B \n",
      "- - B B - - A \n",
      "- - A B - - A \n",
      "- B B A A A B \n",
      "\n",
      "B TURN -> 3\n",
      "- - - - - - - \n",
      "- - - A - - - \n",
      "- - B A - - B \n",
      "- - B B - - A \n",
      "- - A B - - A \n",
      "- B B A A A B \n",
      "\n",
      "A TURN -> 3\n",
      "- - - - - - - \n",
      "- - A A - - - \n",
      "- - B A - - B \n",
      "- - B B - - A \n",
      "- - A B - - A \n",
      "- B B A A A B \n",
      "\n",
      "B TURN -> 2\n",
      "- - - - - - - \n",
      "- - A A - - - \n",
      "- - B A - - B \n",
      "- - B B - - A \n",
      "- B A B - - A \n",
      "- B B A A A B \n",
      "\n",
      "A TURN -> 2\n",
      "- - - - - - - \n",
      "- - A A - - - \n",
      "- - B A - - B \n",
      "- A B B - - A \n",
      "- B A B - - A \n",
      "- B B A A A B \n",
      "\n",
      "B TURN -> 6\n",
      "- - - - - - - \n",
      "- - A A - - - \n",
      "- - B A - - B \n",
      "- A B B - - A \n",
      "- B A B - B A \n",
      "- B B A A A B \n",
      "\n",
      "A TURN -> 4\n",
      "- - - A - - - \n",
      "- - A A - - - \n",
      "- - B A - - B \n",
      "- A B B - - A \n",
      "- B A B - B A \n",
      "- B B A A A B \n",
      "\n",
      "B TURN -> 3\n",
      "- - B A - - - \n",
      "- - A A - - - \n",
      "- - B A - - B \n",
      "- A B B - - A \n",
      "- B A B - B A \n",
      "- B B A A A B \n",
      "\n",
      "A TURN -> 6\n",
      "- - B A - - - \n",
      "- - A A - - - \n",
      "- - B A - - B \n",
      "- A B B - A A \n",
      "- B A B - B A \n",
      "- B B A A A B \n",
      "\n",
      "B TURN -> 6\n",
      "- - B A - - - \n",
      "- - A A - - - \n",
      "- - B A - B B \n",
      "- A B B - A A \n",
      "- B A B - B A \n",
      "- B B A A A B \n",
      "\n",
      "A TURN -> 6\n",
      "- - B A - - - \n",
      "- - A A - A - \n",
      "- - B A - B B \n",
      "- A B B - A A \n",
      "- B A B - B A \n",
      "- B B A A A B \n",
      "\n",
      "B TURN -> 2\n",
      "- - B A - - - \n",
      "- - A A - A - \n",
      "- B B A - B B \n",
      "- A B B - A A \n",
      "- B A B - B A \n",
      "- B B A A A B \n",
      "\n",
      "A TURN -> 2\n",
      "- - B A - - - \n",
      "- A A A - A - \n",
      "- B B A - B B \n",
      "- A B B - A A \n",
      "- B A B - B A \n",
      "- B B A A A B \n",
      "\n",
      "B TURN -> 2\n",
      "- B B A - - - \n",
      "- A A A - A - \n",
      "- B B A - B B \n",
      "- A B B - A A \n",
      "- B A B - B A \n",
      "- B B A A A B \n",
      "\n",
      "A TURN -> 6\n",
      "- B B A - A - \n",
      "- A A A - A - \n",
      "- B B A - B B \n",
      "- A B B - A A \n",
      "- B A B - B A \n",
      "- B B A A A B \n",
      "\n",
      "B TURN -> 1\n",
      "- B B A - A - \n",
      "- A A A - A - \n",
      "- B B A - B B \n",
      "- A B B - A A \n",
      "- B A B - B A \n",
      "B B B A A A B \n",
      "\n",
      "A TURN -> 1\n",
      "- B B A - A - \n",
      "- A A A - A - \n",
      "- B B A - B B \n",
      "- A B B - A A \n",
      "A B A B - B A \n",
      "B B B A A A B \n",
      "\n",
      "B TURN -> 5\n",
      "- B B A - A - \n",
      "- A A A - A - \n",
      "- B B A - B B \n",
      "- A B B - A A \n",
      "A B A B B B A \n",
      "B B B A A A B \n",
      "\n",
      "A TURN -> 5\n",
      "- B B A - A - \n",
      "- A A A - A - \n",
      "- B B A - B B \n",
      "- A B B A A A \n",
      "A B A B B B A \n",
      "B B B A A A B \n",
      "\n",
      "B TURN -> 5\n",
      "- B B A - A - \n",
      "- A A A - A - \n",
      "- B B A B B B \n",
      "- A B B A A A \n",
      "A B A B B B A \n",
      "B B B A A A B \n",
      "\n",
      "A TURN -> 5\n",
      "- B B A - A - \n",
      "- A A A A A - \n",
      "- B B A B B B \n",
      "- A B B A A A \n",
      "A B A B B B A \n",
      "B B B A A A B \n",
      "\n",
      "Player A WON\n"
     ]
    }
   ],
   "source": [
    "def main_AI_vs_AI():\n",
    "    board = initialize_board()\n",
    "    player = 1\n",
    "\n",
    "    while True:\n",
    "        move = choose_move(board, player, max_depth=MAX_DEPTH)\n",
    "\n",
    "        if move is None:\n",
    "            print(\"DRAW\")\n",
    "            return\n",
    "\n",
    "        print(f\"{PLAYERS[player]} TURN -> {move + 1}\")\n",
    "        display(board)\n",
    "\n",
    "        if four_in_a_row(board, player):\n",
    "            print(f\"\\nPlayer {PLAYERS[player]} WON\")\n",
    "            return\n",
    "\n",
    "        print()\n",
    "        player = -player\n",
    "\n",
    "\n",
    "main_AI_vs_AI()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Human vs AI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def main_human_vs_AI():\n",
    "    board = initialize_board()\n",
    "    player = np.random.choice([-1, 1])\n",
    "\n",
    "    while True:\n",
    "        if player == 1:\n",
    "            move = choose_move(board, player, max_depth=MAX_DEPTH)\n",
    "        else:\n",
    "            move = int(input(\"Enter column number (1-7): \"))\n",
    "            play(board, move-1, player)\n",
    "\n",
    "        if move is None:\n",
    "            print(\"DRAW\")\n",
    "            return\n",
    "\n",
    "        print(f\"{PLAYERS[player]} TURN -> {move + 1}\")\n",
    "        display(board)\n",
    "\n",
    "        if four_in_a_row(board, player):\n",
    "            print(f\"\\nPlayer {PLAYERS[player]} WON\")\n",
    "            return\n",
    "\n",
    "        print()\n",
    "        player = -player\n",
    "\n",
    "\n",
    "main_human_vs_AI()"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "371627151b0642f6c185c280fa16312776bda89bbb063ce0f2d8135c657194f4"
  },
  "kernelspec": {
   "display_name": "Python 3.8.2 64-bit ('ci2021': conda)",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}