generate_improved_dataset.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/NicolasChagnet/pokemon-team-optimization/blob/main/generate_improved_dataset.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate improved dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import sys; sys.path.insert(0, '.')\n",
    "from src import config\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import pokebase\n",
    "import requests\n",
    "from bs4 import BeautifulSoup\n",
    "\n",
    "import urllib.parse"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we will first download the [original dataset](https://www.kaggle.com/datasets/rounakbanik/pokemon) from Kaggle. Then we will improve this dataset before saving it for further use."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Starters and pseudo-legendaries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv(config.path_data / \"pokemon.csv\")\n",
    "df = df.drop(\n",
    "    [\n",
    "        \"abilities\",\n",
    "        \"attack\",\n",
    "        \"base_egg_steps\",\n",
    "        \"base_happiness\",\n",
    "        \"capture_rate\",\n",
    "        \"classfication\",\n",
    "        \"defense\",\n",
    "        \"experience_growth\",\n",
    "        \"height_m\",\n",
    "        \"hp\",\n",
    "        \"japanese_name\",\n",
    "        \"percentage_male\",\n",
    "        \"sp_attack\",\n",
    "        \"sp_defense\",\n",
    "        \"speed\",\n",
    "        \"weight_kg\",\n",
    "    ],\n",
    "    axis=1,\n",
    ")\n",
    "df[\"name\"] = df[\"name\"].str.lower()\n",
    "df[\"is_starter\"] = df[\"name\"].isin(config.list_starters_final_stage).astype(int)\n",
    "df[\"is_pseudo_legendary\"] = df[\"name\"].isin(config.list_pseudo_legendaries).astype(int)\n",
    "df[\"is_fossil\"] = df[\"name\"].isin(config.list_fossils).astype(int)\n",
    "\n",
    "df = df.rename(columns={\"against_fight\": \"against_fighting\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_type_main(pkmn):\n",
    "    pkmn_ = pkmn.copy()\n",
    "    pkmn_data = pokebase.pokemon(pkmn_[\"name\"])\n",
    "    pkmn_[\"type1\"] = pkmn_data.types[0].type.name\n",
    "    pkmn_[\"type2\"] = pkmn_data.types[1].type.name if len(pkmn_data.types) == 2 else np.nan\n",
    "    return pkmn_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "df.to_csv(config.path_data / \"pokemon_starter_pl_fossil.csv\", index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Remake datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The dataset available from Kaggle is missing some crucial information:\n",
    "- it does not care about game availability of Pokemon,\n",
    "- it does not care about types being introduced or modified later in the series.\n",
    "\n",
    "To fix this, we remake new datasets containing the type information and base total for each game version as well as a \"national\" pokedex version which contains all Pokemon with gen IX types. The type data is obtained from the PokeAPI while the Pokemon data is scraped from [pokemondb](https://pokemondb.net)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Taken from https://pokebase.readthedocs.io/en/latest/examples/index.html#making-a-type-chart\n",
    "def type_multiplayer(attack, defense):\n",
    "    # Get API data for the attcking type.\n",
    "    dfs_type = pokebase.type_(defense)\n",
    "\n",
    "    # Check which damage_relation list the defense is in. Matches by name\n",
    "    if attack in [t.name for t in dfs_type.damage_relations.no_damage_from]:\n",
    "        return 0.0\n",
    "    elif attack in [t.name for t in dfs_type.damage_relations.half_damage_from]:\n",
    "        return 0.5\n",
    "    elif attack in [t.name for t in dfs_type.damage_relations.double_damage_from]:\n",
    "        return 2.0\n",
    "    else:\n",
    "        return 1.0\n",
    "\n",
    "\n",
    "# Build type matrix (Uncomment below to recompute -- SLOW)\n",
    "types_mat = pd.read_csv(config.path_data / \"attack_defense_types_mat.csv\", index_col=0)\n",
    "# types_mat = pd.DataFrame(\n",
    "#     [[type_multiplayer(type1, type2) for type2 in config.list_types] for type1 in config.list_types],\n",
    "#     columns=config.list_types,\n",
    "#     index=config.list_types,\n",
    "# ).transpose()\n",
    "# types_mat.to_csv(config.path_data / \"attack_defense_types_mat.csv\")\n",
    "\n",
    "# Adapt type matrix for previous generations\n",
    "# See https://pokemondb.net/type for changes pre-gen VI\n",
    "types_mat_prefairy = types_mat.copy(deep=True)\n",
    "types_mat_prefairy = types_mat_prefairy.drop(columns=[\"fairy\"], index=[\"fairy\"])\n",
    "types_mat_prefairy.loc[\"steel\", \"ghost\"] = 0.5\n",
    "types_mat_prefairy.loc[\"steel\", \"dark\"] = 0.5\n",
    "# types_mat_prefairy.to_csv(config.path_data / \"attack_defense_types_mat_gen2to5.csv\")\n",
    "\n",
    "types_mat_rby = types_mat_prefairy.copy(deep=True)\n",
    "types_mat_rby = types_mat_rby.drop(columns=[\"dark\", \"steel\"], index=[\"dark\", \"steel\"])\n",
    "types_mat_rby.loc[\"psychic\", \"ghost\"] = 0\n",
    "types_mat_rby.loc[\"bug\", \"poison\"] = 2\n",
    "types_mat_rby.loc[\"poison\", \"bug\"] = 2\n",
    "types_mat_rby.loc[\"fire\", \"ice\"] = 1\n",
    "# types_mat_rby.to_csv(config.path_data / \"attack_defense_types_mat_gen1.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "BASE_DOMAIN = \"https://pokemondb.net\"\n",
    "\n",
    "\n",
    "# Finds the generation a Pokemon was introduced in using its national pokedex ID\n",
    "def get_gen(x):\n",
    "    upper_bound = [y for y in config.cutoffs if y >= x][0]\n",
    "    return config.cutoffs.index(upper_bound) + 1\n",
    "\n",
    "\n",
    "# Gets general information on all Pokemon (base total, generation) using the national Pokedex\n",
    "def get_stats_table():\n",
    "    pageStats = requests.get(urllib.parse.urljoin(BASE_DOMAIN, \"pokedex/all\"))\n",
    "    soupStats = BeautifulSoup(pageStats.content, \"html.parser\")\n",
    "    tbody = soupStats.select_one(\"tbody\")\n",
    "    trs = tbody.find_all(\"tr\", recursive=True)\n",
    "    results = pd.DataFrame(\n",
    "        [\n",
    "            {\n",
    "                \"name\": tr.select_one(\"a.ent-name\").text.lower(),\n",
    "                \"base_total\": tr.select_one(\"td.cell-total\").text,\n",
    "                \"gen\": get_gen(int(tr.select_one(\"span.infocard-cell-data\").text)),\n",
    "                \"img\": tr.select_one(\"img.icon-pkmn\")[\"src\"],\n",
    "            }\n",
    "            for tr in trs\n",
    "        ]\n",
    "    )\n",
    "    results = results.drop_duplicates(subset=\"name\", keep=\"first\")\n",
    "    return results\n",
    "\n",
    "\n",
    "# Gets all the Pokemon names and types available in a specific version group. Uses the national pokedex data computed only once for base total and generation.\n",
    "def get_names_pokedex(url):\n",
    "    page = requests.get(urllib.parse.urljoin(BASE_DOMAIN, \"pokedex/\" + url))\n",
    "    soup = BeautifulSoup(page.content, \"html.parser\")\n",
    "    base_total_table = get_stats_table()\n",
    "    pokemons = soup.find_all(\"span\", class_=\"infocard-lg-data\")\n",
    "    pokemons_lst = [\n",
    "        {\n",
    "            \"name\": pkmn.select_one(\"a.ent-name\").text.lower(),\n",
    "            \"types\": [x.text.lower() for x in pkmn.findAll(\"a\", class_=\"itype\")],\n",
    "        }\n",
    "        for pkmn in pokemons\n",
    "    ]\n",
    "    pokemons_lst_2 = pd.DataFrame(\n",
    "        [\n",
    "            {\n",
    "                \"name\": pkmn[\"name\"],\n",
    "                \"base_total\": base_total_table.loc[base_total_table[\"name\"] == pkmn[\"name\"], \"base_total\"].iloc[0],\n",
    "                \"generation\": base_total_table.loc[base_total_table[\"name\"] == pkmn[\"name\"], \"gen\"].iloc[0],\n",
    "                \"img\": base_total_table.loc[base_total_table[\"name\"] == pkmn[\"name\"], \"img\"].iloc[0],\n",
    "                \"type1\": pkmn[\"types\"][0],\n",
    "                \"type2\": np.nan if len(pkmn[\"types\"]) <= 1 else pkmn[\"types\"][1],\n",
    "            }\n",
    "            for pkmn in pokemons_lst\n",
    "        ]\n",
    "    )\n",
    "    return pokemons_lst_2\n",
    "\n",
    "\n",
    "# Returns the list of damage multiplier coefficients for a given Pokemon's types and a type matrix\n",
    "def get_type_pkmn(type_1, type_2, type_mat):\n",
    "    type_factor_1 = type_mat.loc[type_1]\n",
    "    type_factor_2 = 1 if pd.isna(type_2) else type_mat.loc[type_2]\n",
    "    series = type_factor_1 * type_factor_2\n",
    "    series_renamed = series.rename({typev: f\"against_{typev}\" for typev in config.list_types})\n",
    "    return series_renamed\n",
    "\n",
    "\n",
    "# Builds the full dataset for a version group including information about fossils, legendaries, etc...\n",
    "def build_dataset(url, gen):\n",
    "    if gen == 1:\n",
    "        type_mat = pd.read_csv(config.path_data / \"attack_defense_types_mat_gen1.csv\", index_col=0)\n",
    "    elif gen >= 2 and gen < 6:\n",
    "        type_mat = pd.read_csv(\n",
    "            filepath_or_buffer=config.path_data / \"attack_defense_types_mat_gen2to5.csv\", index_col=0\n",
    "        )\n",
    "    else:\n",
    "        type_mat = pd.read_csv(filepath_or_buffer=config.path_data / \"attack_defense_types_mat.csv\", index_col=0)\n",
    "\n",
    "    pkms = get_names_pokedex(url)\n",
    "    types = pkms.apply(lambda x: get_type_pkmn(x[\"type1\"], x[\"type2\"], type_mat), axis=1)\n",
    "    pkms[\"is_starter\"] = pkms[\"name\"].isin(config.list_starters_final_stage).astype(int)\n",
    "    pkms[\"is_legendary\"] = pkms[\"name\"].isin(config.list_legendaries).astype(int)\n",
    "    pkms[\"is_pseudo_legendary\"] = pkms[\"name\"].isin(config.list_pseudo_legendaries).astype(int)\n",
    "    pkms[\"is_fossil\"] = pkms[\"name\"].isin(config.list_fossils).astype(int)\n",
    "    return pd.concat([pkms, types], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# RUN TO RECOMPUTE THE DATASETS\n",
    "for gen_id, details in config.list_games.items():\n",
    "    df = build_dataset(**details)\n",
    "    df = df.drop_duplicates(subset=[\"name\"])\n",
    "    df.to_csv(config.get_file_loc(gen_id), index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "pkmn_nat = pd.read_csv(config.get_file_loc(\"nat\"))\n",
    "for k in config.list_games.keys():\n",
    "    if k == \"nat\":\n",
    "        continue\n",
    "    pkmn_k = pd.read_csv(config.get_file_loc(k))\n",
    "    pkmn_nat[f\"is_{k}\"] = False\n",
    "    names_pkmn = pkmn_k[\"name\"]\n",
    "    pkmn_nat.loc[pkmn_nat[\"name\"].isin(names_pkmn), f\"is_{k}\"] = True\n",
    "\n",
    "pkmn_nat.to_csv(config.path_data / \"pokemon_full.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pokemon-optimization_py",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}