playground.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "37474969-e894-45b5-beb8-62fc522afb2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import envs\n",
    "import agents\n",
    "import random\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9d6e0aa9-b8d2-455f-b8a6-b5a142ad6752",
   "metadata": {},
   "outputs": [],
   "source": [
    "env = envs.WordleEnv4.create_default(render=True, max_sub_attempts=10, prefill=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "97faa240-0ccd-4123-9f56-300e9fe9afd3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "attempt 0\n",
      "slan \n",
      "     \n",
      "     \n",
      "     \n",
      "     \n",
      "     \n"
     ]
    }
   ],
   "source": [
    "env.reset()\n",
    "env.print_render()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a178a1f1-7351-4a69-8329-d2a51ae4e01a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'slant'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "env.word"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "6649cd97-40b4-45de-bdec-1d63e195a9dc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([18, 11,  0, 13,  6,  3,  3,  3,  3,  1, 18, 11,  0, 13, 10,  3,  3,\n",
       "         3,  3,  1, 18, 11,  0, 13,  4,  3,  3,  3,  3,  1, 18, 11,  0, 13,\n",
       "        19,  3,  3,  3,  3,  3, 19, 11,  0, 13, 19,  1,  1,  1,  1,  0, 26,\n",
       "        26, 26, 26, 26,  1,  1,  1,  1,  1]),\n",
       " 0,\n",
       " False,\n",
       " {'attempt': 4, 'sub_attempt': 15, 'index': 4})"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "env.step(env.index_map.get('t'), skip_render=True)\n",
    "# env.print_render()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "14fed364-a61b-4bb7-acf5-b40901f0dc8e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "attempt 4\n",
      "slang\n",
      "slank\n",
      "slane\n",
      "slant\n",
      "tlant\n",
      "     \n"
     ]
    }
   ],
   "source": [
    "env.print_render()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "1cfa837c-e448-47c0-899a-be8b76628f48",
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = agents.WordleAgent1.create_default()\n",
    "env = envs.WordleEnv2.create_default(flatten=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "00e81fa6-bfa5-4e20-9dd5-e9be9f1423ab",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "attempt 1\n",
      "pacts\n",
      "     \n",
      "     \n",
      "     \n",
      "     \n",
      "     \n",
      "0\n",
      "\n",
      "attempt 2\n",
      "pacts\n",
      "bemad\n",
      "     \n",
      "     \n",
      "     \n",
      "     \n",
      "0\n",
      "\n",
      "attempt 3\n",
      "pacts\n",
      "bemad\n",
      "gnarl\n",
      "     \n",
      "     \n",
      "     \n",
      "0\n",
      "\n",
      "attempt 4\n",
      "pacts\n",
      "bemad\n",
      "gnarl\n",
      "angel\n",
      "     \n",
      "     \n",
      "10\n"
     ]
    }
   ],
   "source": [
    "observation = env.reset()\n",
    "done = False\n",
    "while not done:\n",
    "    action = agent.act(observation)\n",
    "    observation, reward, done, info = env.step(env._translate_word(action))\n",
    "    env.print_render()\n",
    "    print(reward)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ebf92e6f-936b-49e2-8376-fb23ac29ac02",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'env' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<timed exec>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'env' is not defined"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "rewards = []\n",
    "attempts = []\n",
    "for _ in range(1000):\n",
    "    observation = env.reset()\n",
    "    done = False\n",
    "    while not done:\n",
    "        action = agent.act(observation)\n",
    "        observation, reward, done, info =env.step(env._translate_word(action))\n",
    "    rewards.append(1 if reward == 1000 else 0)\n",
    "    attempts.append(info[\"attempt\"])\n",
    "print(f'average wins: {np.mean(rewards)}')\n",
    "print(f'average attempts: {np.mean(attempts)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9de0f826-8e5b-4cec-8fc2-f1045f0d5e97",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Stable Baselines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3039af12-099b-48fb-b5f6-dfc74a033053",
   "metadata": {},
   "outputs": [],
   "source": [
    "import envs\n",
    "from stable_baselines3 import PPO, A2C\n",
    "env = envs.WordleEnv7.create_default()\n",
    "model = PPO.load('bonus_model', env=env)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "089bf0ef-0bd6-48c0-af27-e9d4e5319f62",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "acrid\n",
      "{'attempt': 6, 'sub_attempt': 6} 19\n",
      "\n",
      "attempt 6\n",
      "tales\n",
      "carer\n",
      "raper\n",
      "raker\n",
      "raser\n",
      "rarer\n"
     ]
    }
   ],
   "source": [
    "observation = env.reset()\n",
    "print(env.word)\n",
    "done = False\n",
    "rewards = 0\n",
    "actions = []\n",
    "while not done:# and env.attempt < 3:\n",
    "    action, _state = model.predict(observation, deterministic=True)\n",
    "    actions.append(action)\n",
    "    observation, reward, done, info = env.step(action)\n",
    "    rewards += reward\n",
    "    # print(env.attempt, reward, done, env.letter_map[action])\n",
    "    # print(reward)\n",
    "print(env._get_info(), rewards)\n",
    "env.print_render()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cef5a066-f09b-4614-8cbf-ccb1bea32a5f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "average wins: 0.78\n",
      "average attempts: 6.0\n",
      "average sub_attempts: 6.0\n",
      "CPU times: user 720 ms, sys: 4.31 ms, total: 725 ms\n",
      "Wall time: 726 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "rewards = []\n",
    "attempts = []\n",
    "sub_attempts = []\n",
    "for _ in range(100):\n",
    "    observation = env.reset()\n",
    "    done = False\n",
    "    # observation, reward, done, info = env.step(env._translate_word('alter'))\n",
    "    while not done:\n",
    "        action, _state = model.predict(observation, deterministic=True)\n",
    "        # print(env._translate_action(action))\n",
    "        # env.print_render()\n",
    "        observation, reward, done, info = env.step(action)\n",
    "    rewards.append(1 if reward == 10 else 0)\n",
    "    attempts.append(info[\"attempt\"])\n",
    "    sub_attempts.append(info[\"sub_attempt\"])\n",
    "    # print(f'scored {reward} in {info[\"attempt\"]} attempt(s)')\n",
    "print(f'average wins: {np.mean(rewards)}')\n",
    "print(f'average attempts: {np.mean(attempts)}')\n",
    "print(f'average sub_attempts: {np.mean(sub_attempts)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b37ed0e-d30e-4bb6-afd3-a1c2cf7393c3",
   "metadata": {},
   "source": [
    "# Manual Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "0048896b-254e-464a-b253-a08cf7d4d9ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "import string\n",
    "import pandas as pd\n",
    "from tabulate import tabulate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b441b2a5-0f94-451b-aad5-77e2ee056b81",
   "metadata": {},
   "outputs": [],
   "source": [
    "wbm = np.array([list(word) for word in word_bank])\n",
    "vwm = np.array([list(word) for word in valid_words])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0b39d5d-6828-4162-9724-02ea031669a0",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Letter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "8e5f9e44-f46c-4c12-be16-bb3d6260f62a",
   "metadata": {},
   "outputs": [],
   "source": [
    "stats = []\n",
    "for l in string.ascii_lowercase:\n",
    "    lm = wbm == l\n",
    "    stats.append({\n",
    "        'letter': l,\n",
    "        'words_contain': lm.max(axis=1).sum(),\n",
    "        'words_contain_percent': lm.max(axis=1).mean(),\n",
    "        'freq': lm.sum(axis=1).mean(),\n",
    "        'freq_if_present': lm.sum(axis=1)[lm.sum(axis=1) > 0].mean(),\n",
    "        'max_freq': lm.sum(axis=1).max(),\n",
    "    })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "id": "a13ed6ca-4e89-4e2f-b106-8c4617267f81",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "| letter   |   words_contain |   words_contain_percent |      freq |   freq_if_present |   max_freq |\n",
      "|:---------|----------------:|------------------------:|----------:|------------------:|-----------:|\n",
      "| a        |             909 |               0.392657  | 0.422894  |           1.07701 |          2 |\n",
      "| b        |             267 |               0.115335  | 0.121382  |           1.05243 |          3 |\n",
      "| c        |             448 |               0.193521  | 0.206048  |           1.06473 |          2 |\n",
      "| d        |             370 |               0.159827  | 0.169762  |           1.06216 |          3 |\n",
      "| e        |            1056 |               0.456156  | 0.532613  |           1.16761 |          3 |\n",
      "| f        |             207 |               0.0894168 | 0.0993521 |           1.11111 |          3 |\n",
      "| g        |             300 |               0.12959   | 0.134341  |           1.03667 |          2 |\n",
      "| h        |             379 |               0.163715  | 0.168035  |           1.02639 |          2 |\n",
      "| i        |             647 |               0.279482  | 0.289849  |           1.03709 |          2 |\n",
      "| j        |              27 |               0.0116631 | 0.0116631 |           1       |          1 |\n",
      "| k        |             202 |               0.087257  | 0.0907127 |           1.0396  |          2 |\n",
      "| l        |             648 |               0.279914  | 0.310583  |           1.10957 |          2 |\n",
      "| m        |             298 |               0.128726  | 0.136501  |           1.0604  |          3 |\n",
      "| n        |             550 |               0.237581  | 0.24838   |           1.04545 |          3 |\n",
      "| o        |             673 |               0.290713  | 0.325702  |           1.12036 |          2 |\n",
      "| p        |             346 |               0.14946   | 0.158531  |           1.06069 |          3 |\n",
      "| q        |              29 |               0.012527  | 0.012527  |           1       |          1 |\n",
      "| r        |             837 |               0.361555  | 0.388337  |           1.07407 |          3 |\n",
      "| s        |             618 |               0.266955  | 0.288985  |           1.08252 |          3 |\n",
      "| t        |             667 |               0.288121  | 0.314903  |           1.09295 |          3 |\n",
      "| u        |             457 |               0.197408  | 0.201728  |           1.02188 |          2 |\n",
      "| v        |             149 |               0.0643629 | 0.0660907 |           1.02685 |          2 |\n",
      "| w        |             194 |               0.0838013 | 0.0842333 |           1.00515 |          2 |\n",
      "| x        |              37 |               0.0159827 | 0.0159827 |           1       |          1 |\n",
      "| y        |             417 |               0.18013   | 0.183585  |           1.01918 |          2 |\n",
      "| z        |              35 |               0.0151188 | 0.0172786 |           1.14286 |          2 |\n"
     ]
    }
   ],
   "source": [
    "print(tabulate(pd.DataFrame(stats).set_index('letter'), tablefmt=\"pipe\", headers=\"keys\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "ca23ce23-422c-4dc1-a16d-5189bfb5caed",
   "metadata": {},
   "outputs": [],
   "source": [
    "contain_rate= {}\n",
    "for l in string.ascii_lowercase:\n",
    "    lm = wbm == l\n",
    "    contain_rate[l] = lm.max(axis=1).mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "0d6e7195-f651-43b6-a1fc-29d3f9269fb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "scores = []\n",
    "for word in valid_words:\n",
    "    score = 0\n",
    "    for l in set(list(word)):\n",
    "        score += contain_rate[l]\n",
    "    scores.append(score)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "d398ad09-b125-4eef-a376-de237f4cf262",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame({'word': valid_words, 'score': scores})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "107f2dbc-94e9-488b-83c5-b7f9dbbd87b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "df['in_bank'] = df['word'].isin(word_bank)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "id": "a1bfd584-2247-4e59-96ec-fc5c6394f04f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "| word   |   score | in_bank   |\n",
      "|:-------|--------:|:----------|\n",
      "| alter  | 1.7784  | True      |\n",
      "| later  | 1.7784  | True      |\n",
      "| alert  | 1.7784  | True      |\n",
      "| irate  | 1.77797 | True      |\n",
      "| arose  | 1.76803 | True      |\n",
      "| stare  | 1.76544 | True      |\n",
      "| raise  | 1.7568  | True      |\n",
      "| arise  | 1.7568  | True      |\n",
      "| learn  | 1.72786 | True      |\n",
      "| renal  | 1.72786 | True      |\n",
      "| saner  | 1.7149  | True      |\n",
      "| snare  | 1.7149  | True      |\n",
      "| cater  | 1.69201 | True      |\n",
      "| trace  | 1.69201 | True      |\n",
      "| react  | 1.69201 | True      |\n",
      "| crate  | 1.69201 | True      |\n",
      "| stale  | 1.6838  | True      |\n",
      "| steal  | 1.6838  | True      |\n",
      "| least  | 1.6838  | True      |\n",
      "| slate  | 1.6838  | True      |\n"
     ]
    }
   ],
   "source": [
    "print(tabulate(df[df['in_bank']].sort_values(by='score', ascending=False).head(20).set_index('word'), tablefmt=\"pipe\", headers=\"keys\"))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "d6336cc7-0f31-47ae-b129-90e0ef7f263c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "| word   |   score | in_bank   |\n",
      "|:-------|--------:|:----------|\n",
      "| oater  | 1.7892  | False     |\n",
      "| roate  | 1.7892  | False     |\n",
      "| orate  | 1.7892  | False     |\n",
      "| realo  | 1.78099 | False     |\n",
      "| taler  | 1.7784  | False     |\n",
      "| artel  | 1.7784  | False     |\n",
      "| ratel  | 1.7784  | False     |\n",
      "| alert  | 1.7784  | True      |\n",
      "| alter  | 1.7784  | True      |\n",
      "| later  | 1.7784  | True      |\n",
      "| terai  | 1.77797 | False     |\n",
      "| irate  | 1.77797 | True      |\n",
      "| retia  | 1.77797 | False     |\n",
      "| raile  | 1.76976 | False     |\n",
      "| ariel  | 1.76976 | False     |\n",
      "| arose  | 1.76803 | True      |\n",
      "| aeros  | 1.76803 | False     |\n",
      "| soare  | 1.76803 | False     |\n",
      "| taser  | 1.76544 | False     |\n",
      "| strae  | 1.76544 | False     |\n"
     ]
    }
   ],
   "source": [
    "print(tabulate(df.sort_values(by='score', ascending=False).head(20).set_index('word'), tablefmt=\"pipe\", headers=\"keys\"))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "id": "0c6f1cfa-a572-4101-a7a4-921fc80b4d53",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "| word   |   score | in_bank   |\n",
      "|:-------|--------:|:----------|\n",
      "| sonic  | 1.26825 | True      |\n",
      "| scion  | 1.26825 | True      |\n",
      "| noisy  | 1.25486 | True      |\n",
      "| disco  | 1.1905  | True      |\n",
      "| bison  | 1.19006 | True      |\n",
      "| sound  | 1.15248 | True      |\n",
      "| synod  | 1.13521 | True      |\n",
      "| shiny  | 1.12786 | True      |\n",
      "| spiny  | 1.11361 | True      |\n",
      "| suing  | 1.11102 | True      |\n",
      "| using  | 1.11102 | True      |\n",
      "| minus  | 1.11015 | True      |\n",
      "| bonus  | 1.10799 | True      |\n",
      "| doing  | 1.09719 | True      |\n",
      "| dingo  | 1.09719 | True      |\n",
      "| spicy  | 1.06955 | True      |\n",
      "| music  | 1.06609 | True      |\n",
      "| snowy  | 1.05918 | True      |\n",
      "| bingo  | 1.0527  | True      |\n",
      "| hound  | 1.04924 | True      |\n"
     ]
    }
   ],
   "source": [
    "print(tabulate(df[(df['in_bank']) & (~df.word.str.contains('[alter]', regex=True))].sort_values(by='score', ascending=False).head(20).set_index('word'), tablefmt=\"pipe\", headers=\"keys\"))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "id": "ec7619a2-ff76-4c24-bb26-4a0668947c61",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "| word   |   score | in_bank   |\n",
      "|:-------|--------:|:----------|\n",
      "| lysin  | 1.24406 | False     |\n",
      "| linds  | 1.22376 | False     |\n",
      "| sulci  | 1.21728 | False     |\n",
      "| sling  | 1.19352 | True      |\n",
      "| lings  | 1.19352 | False     |\n",
      "| limns  | 1.19266 | False     |\n",
      "| hilus  | 1.18747 | False     |\n",
      "| blins  | 1.17927 | False     |\n",
      "| incus  | 1.17495 | False     |\n",
      "| pilus  | 1.17322 | False     |\n",
      "| pulis  | 1.17322 | False     |\n",
      "| shily  | 1.17019 | False     |\n",
      "| clips  | 1.16933 | False     |\n",
      "| idyls  | 1.16631 | False     |\n",
      "| unlid  | 1.15421 | False     |\n",
      "| linch  | 1.15421 | False     |\n",
      "| gusli  | 1.15335 | False     |\n",
      "| iglus  | 1.15335 | False     |\n",
      "| muils  | 1.15248 | False     |\n",
      "| simul  | 1.15248 | False     |\n"
     ]
    }
   ],
   "source": [
    "print(tabulate(df[(~df.word.str.contains('[oater]', regex=True))].sort_values(by='score', ascending=False).head(20).set_index('word'), tablefmt=\"pipe\", headers=\"keys\"))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b2ac4cb-c4d7-4bad-bb07-3b3711aa458f",
   "metadata": {},
   "source": [
    "## Word"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "673f9483-6593-4464-a9ea-82d57022d72d",
   "metadata": {},
   "outputs": [],
   "source": [
    "stats =[]\n",
    "for word in vwm:\n",
    "    present_letter_msk = np.isin(wbm, word)\n",
    "    known_letter_msk = wbm[:, :] == word\n",
    "\n",
    "    score = present_letter_msk * 1\n",
    "    score[known_letter_msk] = 1\n",
    "    stats.append({\n",
    "        'avg_score': score.sum(axis=1).mean(),\n",
    "    })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "0e768871-7478-4b97-ab13-dbb45de6143d",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.DataFrame(stats)\n",
    "df['word'] = valid_words\n",
    "df['in_bank'] = df['word'].isin(word_bank)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "cfd5d804-24f5-43c7-b07a-7a7e0742e31d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "| word   |   avg_score | in_bank   |\n",
      "|:-------|------------:|:----------|\n",
      "| roate  |     1.98445 | False     |\n",
      "| oater  |     1.98445 | False     |\n",
      "| orate  |     1.98445 | False     |\n",
      "| realo  |     1.98013 | False     |\n",
      "| artel  |     1.96933 | False     |\n",
      "| later  |     1.96933 | True      |\n",
      "| taler  |     1.96933 | False     |\n",
      "| ratel  |     1.96933 | False     |\n",
      "| alter  |     1.96933 | True      |\n",
      "| alert  |     1.96933 | True      |\n",
      "| soare  |     1.95853 | False     |\n",
      "| aeros  |     1.95853 | False     |\n",
      "| arose  |     1.95853 | True      |\n",
      "| irate  |     1.9486  | True      |\n",
      "| retia  |     1.9486  | False     |\n",
      "| terai  |     1.9486  | False     |\n",
      "| reast  |     1.94773 | False     |\n",
      "| stare  |     1.94773 | True      |\n",
      "| strae  |     1.94773 | False     |\n",
      "| teras  |     1.94773 | False     |\n"
     ]
    }
   ],
   "source": [
    "print(tabulate(df.sort_values(by='avg_score', ascending=False).head(20).set_index('word'), tablefmt=\"pipe\", headers=\"keys\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "0d978774-f946-4c38-bc8f-abd207152328",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "| word   |   avg_score | in_bank   |\n",
      "|:-------|------------:|:----------|\n",
      "| alter  |     1.96933 | True      |\n",
      "| alert  |     1.96933 | True      |\n",
      "| later  |     1.96933 | True      |\n",
      "| arose  |     1.95853 | True      |\n",
      "| irate  |     1.9486  | True      |\n",
      "| stare  |     1.94773 | True      |\n",
      "| raise  |     1.92268 | True      |\n",
      "| arise  |     1.92268 | True      |\n",
      "| learn  |     1.90281 | True      |\n",
      "| renal  |     1.90281 | True      |\n",
      "| snare  |     1.88121 | True      |\n",
      "| saner  |     1.88121 | True      |\n",
      "| least  |     1.86998 | True      |\n",
      "| stale  |     1.86998 | True      |\n",
      "| slate  |     1.86998 | True      |\n",
      "| steal  |     1.86998 | True      |\n",
      "| react  |     1.86479 | True      |\n",
      "| crate  |     1.86479 | True      |\n",
      "| trace  |     1.86479 | True      |\n",
      "| cater  |     1.86479 | True      |\n"
     ]
    }
   ],
   "source": [
    "print(tabulate(df[df[\"in_bank\"]].sort_values(by='avg_score', ascending=False).head(20).set_index('word'), tablefmt=\"pipe\", headers=\"keys\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "97b19174-b8f5-4e21-b1ef-b6fc83f40d84",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "|    | index   |            0 | index   |           1 | index   |           2 | index   |             3 | index   |             4 |\n",
      "|---:|:--------|-------------:|:--------|------------:|:--------|------------:|:--------|--------------:|:--------|--------------:|\n",
      "|  0 | s       |   0.158099   | a       | 0.131317    | a       | 0.132613    | e       |   0.137365    | e       |   0.183153    |\n",
      "|  1 | c       |   0.0855292  | o       | 0.120518    | i       | 0.114903    | n       |   0.0786177   | y       |   0.157235    |\n",
      "|  2 | b       |   0.07473    | r       | 0.115335    | o       | 0.1054      | s       |   0.0738661   | t       |   0.109287    |\n",
      "|  3 | t       |   0.0643629  | e       | 0.104536    | e       | 0.0764579   | a       |   0.0704104   | r       |   0.0915767   |\n",
      "|  4 | p       |   0.0613391  | i       | 0.087257    | u       | 0.0712743   | l       |   0.0699784   | l       |   0.0673866   |\n",
      "|  5 | a       |   0.0609071  | l       | 0.0868251   | r       | 0.0704104   | i       |   0.0682505   | h       |   0.0600432   |\n",
      "|  6 | f       |   0.0587473  | u       | 0.0803456   | n       | 0.0600432   | c       |   0.0656587   | n       |   0.0561555   |\n",
      "|  7 | g       |   0.049676   | h       | 0.062203    | l       | 0.0483801   | r       |   0.0656587   | d       |   0.0509719   |\n",
      "|  8 | d       |   0.0479482  | n       | 0.037581    | t       | 0.0479482   | t       |   0.0600432   | k       |   0.0488121   |\n",
      "|  9 | m       |   0.0462203  | t       | 0.0332613   | s       | 0.0345572   | o       |   0.0570194   | a       |   0.0276458   |\n",
      "| 10 | r       |   0.0453564  | p       | 0.0263499   | d       | 0.0323974   | u       |   0.0354212   | o       |   0.025054    |\n",
      "| 11 | l       |   0.038013   | w       | 0.0190065   | g       | 0.0289417   | g       |   0.0328294   | p       |   0.0241901   |\n",
      "| 12 | w       |   0.0358531  | c       | 0.0172786   | m       | 0.0263499   | d       |   0.0298056   | m       |   0.0181425   |\n",
      "| 13 | e       |   0.0311015  | m       | 0.0164147   | p       | 0.025054    | m       |   0.0293737   | g       |   0.0177106   |\n",
      "| 14 | h       |   0.0298056  | y       | 0.00993521  | b       | 0.024622    | k       |   0.0237581   | s       |   0.0155508   |\n",
      "| 15 | v       |   0.0185745  | d       | 0.00863931  | c       | 0.0241901   | p       |   0.0215983   | c       |   0.0133909   |\n",
      "| 16 | o       |   0.0177106  | b       | 0.00691145  | v       | 0.0211663   | v       |   0.0198704   | f       |   0.0112311   |\n",
      "| 17 | n       |   0.0159827  | s       | 0.00691145  | y       | 0.012527    | f       |   0.0151188   | w       |   0.00734341  |\n",
      "| 18 | i       |   0.0146868  | v       | 0.00647948  | w       | 0.0112311   | h       |   0.012095    | i       |   0.00475162  |\n",
      "| 19 | u       |   0.0142549  | x       | 0.00604752  | f       | 0.0107991   | w       |   0.0107991   | b       |   0.00475162  |\n",
      "| 20 | q       |   0.00993521 | g       | 0.00518359  | k       | 0.00518359  | b       |   0.0103672   | x       |   0.00345572  |\n",
      "| 21 | k       |   0.00863931 | k       | 0.00431965  | x       | 0.00518359  | z       |   0.00863931  | z       |   0.00172786  |\n",
      "| 22 | j       |   0.00863931 | f       | 0.00345572  | z       | 0.00475162  | x       |   0.0012959   | u       |   0.000431965 |\n",
      "| 23 | y       |   0.00259179 | q       | 0.00215983  | h       | 0.00388769  | y       |   0.0012959   | nan     | nan           |\n",
      "| 24 | z       |   0.0012959  | z       | 0.000863931 | j       | 0.0012959   | j       |   0.000863931 | nan     | nan           |\n",
      "| 25 | nan     | nan          | j       | 0.000863931 | q       | 0.000431965 | nan     | nan           | nan     | nan           |\n"
     ]
    }
   ],
   "source": [
    "print(tabulate(pd.concat([pd.DataFrame(wbm)[i].value_counts(normalize=True).to_frame().reset_index() for i in range(5)], axis=1), tablefmt=\"pipe\", headers=\"keys\"))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e5877110-ea4c-41cc-90a7-af93806b330f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import string"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7d21df65-6af3-4b2f-be2f-989eda5f3a27",
   "metadata": {},
   "outputs": [],
   "source": [
    "env = envs.WordleEnv2.create_default(flatten=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0fa8281d-61b0-4dc3-889c-6ea59529e1d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "env.reset()\n",
    "obs, _, _, _ =env.step(env._translate_word('alter'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "02fb94c5-4aa9-4f7e-ae06-f7f1f071cd6b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2, 5)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "obs.shape[1:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "d9ed9b76-8070-4167-a49a-5d84f06453fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "letter_map = {i: l for i, l in enumerate(string.ascii_lowercase)}\n",
    "index_map = {l: i for i, l in enumerate(string.ascii_lowercase)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "01116d67-0895-45bf-a0c9-eee3095cd2a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "letters = pd.DataFrame(obs[:, 0, :]).applymap(letter_map.get)\n",
    "scores = pd.DataFrame(obs[:, 1, :])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "2f5ca544-fb7f-48bb-922f-7a73cb13da8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('valid-words.csv') as f:\n",
    "    valid_words = f.readlines()\n",
    "valid_words = [w.strip() for w in valid_words]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "1e621bf1-ace5-4a38-99a0-90d4501f4719",
   "metadata": {},
   "outputs": [],
   "source": [
    "words = pd.DataFrame([list(w) for w in valid_words])\n",
    "words['whole'] = valid_words"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "8085b1c7-ff78-4452-bbf2-bb5520f12d84",
   "metadata": {},
   "outputs": [],
   "source": [
    "freq = {i: words[0].value_counts(normalize=True).to_dict() for i in range(5)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "31d12c55-9c92-4ed5-a4f2-e96e5162af41",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_score(row):\n",
    "    return sum([freq[i].get(row[i], 0) for i in range(5)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "fd110d82-870d-49ee-a3d7-1b08ca128feb",
   "metadata": {},
   "outputs": [],
   "source": [
    "words['score'] = words.apply(get_score, axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 133,
   "id": "28c77f84-2d94-45e1-8677-f231db2e3d49",
   "metadata": {},
   "outputs": [],
   "source": [
    "# known absent letters\n",
    "known_absent_letters = letters[scores == 0].melt()['value'].dropna().unique().tolist()\n",
    "# known present letters\n",
    "known_present_letters = letters[scores > 0].melt()['value'].dropna().unique().tolist()\n",
    "# known exact letters\n",
    "known_exact_letters = letters[scores == 2].melt().dropna().set_index('variable')['value'].to_dict()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 134,
   "id": "498e2c80-d2aa-477b-99eb-b366e4c63c7f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['a', 'l', 'o', 'e']"
      ]
     },
     "execution_count": 134,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "known_absent_letters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "id": "e14e0cc0-ea72-4227-872c-8b63f5a5ccb2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{}"
      ]
     },
     "execution_count": 122,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "known_exact_letters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 135,
   "id": "c00e1603-5a78-4d43-bd4a-2f4a23b30d69",
   "metadata": {},
   "outputs": [],
   "source": [
    "msk = words[0].astype(bool)\n",
    "\n",
    "if known_absent_letters:\n",
    "    msk = msk & (~words.whole.str.contains(f'[{\"\".join(known_absent_letters)}]'))\n",
    "    \n",
    "if known_present_letters:\n",
    "    msk = msk & (words.whole.str.contains(f'[{\"\".join(known_present_letters)}]'))\n",
    "    \n",
    "if known_present_letters:\n",
    "    for k, v in known_exact_letters.items():\n",
    "        msk = msk & (words[k] == v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 136,
   "id": "bdedcc40-8a0a-4d0b-b979-96ab222feb64",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'swiss'"
      ]
     },
     "execution_count": 136,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "words[msk].sort_values(by='score', ascending=False).iloc[0].whole"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 137,
   "id": "8dd1e201-78b1-45bd-b3e9-d3be027ae98d",
   "metadata": {},
   "outputs": [],
   "source": [
    "obs, _, _, _ =env.step(env._translate_word('swiss'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 138,
   "id": "5f45c2a0-9ade-4013-acfc-d7a928c2efc0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[[ 0, 11, 19,  4, 17],\n",
       "        [ 0,  0,  1,  0,  1]],\n",
       "\n",
       "       [[18, 19, 14, 18, 18],\n",
       "        [ 1,  1,  0,  2,  1]],\n",
       "\n",
       "       [[18, 22,  8, 18, 18],\n",
       "        [ 1,  0,  0,  2,  1]],\n",
       "\n",
       "       [[26, 26, 26, 26, 26],\n",
       "        [ 0,  0,  0,  0,  0]],\n",
       "\n",
       "       [[26, 26, 26, 26, 26],\n",
       "        [ 0,  0,  0,  0,  0]],\n",
       "\n",
       "       [[26, 26, 26, 26, 26],\n",
       "        [ 0,  0,  0,  0,  0]]])"
      ]
     },
     "execution_count": 138,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "obs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "id": "26ea1c88-9514-4971-9989-f3c5a84e9c65",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([18, 22,  8, 18, 18])"
      ]
     },
     "execution_count": 144,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "env._translate_word('swiss')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e83e165a-3d8b-4288-aa77-4bcf88ea4cd8",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "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.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}