{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "mysterious-calgary",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading environment football failed: No module named 'gfootball'\n"
     ]
    }
   ],
   "source": [
    "from kaggle_environments import make"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "complex-setting",
   "metadata": {},
   "outputs": [],
   "source": [
    "from lux.game import Game\n",
    "from lux.game_map import Cell, RESOURCE_TYPES, Position\n",
    "from lux.constants import Constants\n",
    "from lux.game_constants import GAME_CONSTANTS\n",
    "from lux import annotate\n",
    "import math\n",
    "\n",
    "### Define helper functions\n",
    "\n",
    "# this snippet finds all resources stored on the map and puts them into a list so we can search over them\n",
    "def find_resources(game_state):\n",
    "    resource_tiles: list[Cell] = []\n",
    "    width, height = game_state.map_width, game_state.map_height\n",
    "    for y in range(height):\n",
    "        for x in range(width):\n",
    "            cell = game_state.map.get_cell(x, y)\n",
    "            if cell.has_resource():\n",
    "                resource_tiles.append(cell)\n",
    "    return resource_tiles\n",
    "\n",
    "# the next snippet finds the closest resources that we can mine given position on a map\n",
    "def find_closest_resources(pos, player, resource_tiles):\n",
    "    closest_dist = math.inf\n",
    "    closest_resource_tile = None\n",
    "    for resource_tile in resource_tiles:\n",
    "        # we skip over resources that we can't mine due to not having researched them\n",
    "        if resource_tile.resource.type == Constants.RESOURCE_TYPES.COAL and not player.researched_coal(): continue\n",
    "        if resource_tile.resource.type == Constants.RESOURCE_TYPES.URANIUM and not player.researched_uranium(): continue\n",
    "        dist = resource_tile.pos.distance_to(pos)\n",
    "        if dist < closest_dist:\n",
    "            closest_dist = dist\n",
    "            closest_resource_tile = resource_tile\n",
    "    return closest_resource_tile\n",
    "\n",
    "def find_closest_city_tile(pos, player):\n",
    "    closest_city_tile = None\n",
    "    if len(player.cities) > 0:\n",
    "        closest_dist = math.inf\n",
    "        # the cities are stored as a dictionary mapping city id to the city object, which has a citytiles field that\n",
    "        # contains the information of all citytiles in that city\n",
    "        for k, city in player.cities.items():\n",
    "            for city_tile in city.citytiles:\n",
    "                dist = city_tile.pos.distance_to(pos)\n",
    "                if dist < closest_dist:\n",
    "                    closest_dist = dist\n",
    "                    closest_city_tile = city_tile\n",
    "    return closest_city_tile\n",
    "\n",
    "game_state = None\n",
    "def agent(observation, configuration):\n",
    "    global game_state\n",
    "\n",
    "    ### Do not edit ###\n",
    "    if observation[\"step\"] == 0:\n",
    "        game_state = Game()\n",
    "        game_state._initialize(observation[\"updates\"])\n",
    "        game_state._update(observation[\"updates\"][2:])\n",
    "        game_state.id = observation.player\n",
    "    else:\n",
    "        game_state._update(observation[\"updates\"])\n",
    "    \n",
    "    actions = []\n",
    "\n",
    "    ### AI Code goes down here! ### \n",
    "    player = game_state.players[observation.player]\n",
    "    opponent = game_state.players[(observation.player + 1) % 2]\n",
    "    width, height = game_state.map.width, game_state.map.height\n",
    "\n",
    "    resource_tiles = find_resources(game_state)\n",
    "    for unit in player.units:\n",
    "        # if the unit is a worker (can mine resources) and can perform an action this turn\n",
    "        if unit.is_worker() and unit.can_act():\n",
    "            # we want to mine only if there is space left in the worker's cargo\n",
    "            if unit.get_cargo_space_left() > 0:\n",
    "                # find the closest resource if it exists to this unit\n",
    "                closest_resource_tile = find_closest_resources(unit.pos, player, resource_tiles)\n",
    "                if closest_resource_tile is not None:\n",
    "                    # create a move action to move this unit in the direction of the closest resource tile and add to our actions list\n",
    "                    action = unit.move(unit.pos.direction_to(closest_resource_tile.pos))\n",
    "                    actions.append(action)\n",
    "            else:\n",
    "                # find the closest citytile and move the unit towards it to drop resources to a citytile to fuel the city\n",
    "                closest_city_tile = find_closest_city_tile(unit.pos, player)\n",
    "                if closest_city_tile is not None:\n",
    "                    # create a move action to move this unit in the direction of the closest resource tile and add to our actions list\n",
    "                    action = unit.move(unit.pos.direction_to(closest_city_tile.pos))\n",
    "                    actions.append(action)\n",
    "    \n",
    "    return actions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "valid-applicant",
   "metadata": {},
   "outputs": [],
   "source": [
    "env = make(\"lux_ai_2021\", configuration={\"loglevel\": 0, \"annotations\": False}, debug=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "modular-density",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== Episode 2\n",
      "=== Episode 3\n",
      "=== Episode 4\n",
      "=== Episode 5\n",
      "=== Episode 6\n",
      "=== Episode 7\n",
      "=== Episode 8\n",
      "=== Episode 9\n",
      "=== Episode 10\n",
      "=== Episode 11\n",
      "=== Episode 12\n",
      "=== Episode 13\n",
      "=== Episode 14\n",
      "=== Episode 15\n",
      "=== Episode 16\n",
      "=== Episode 17\n",
      "=== Episode 18\n",
      "=== Episode 19\n",
      "=== Episode 20\n",
      "=== Episode 21\n",
      "=== Episode 22\n",
      "=== Episode 23\n",
      "=== Episode 24\n",
      "=== Episode 25\n",
      "=== Episode 26\n",
      "=== Episode 27\n",
      "=== Episode 28\n",
      "CPU times: user 1min 22s, sys: 468 ms, total: 1min 23s\n",
      "Wall time: 1min 31s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# env.state = data[\"steps\"][100]\n",
    "# env.train([None, \"simple_agent\"])\n",
    "trainer = env.train([None, agent])\n",
    "# env.state = data[\"steps\"][100]\n",
    "# env.steps = data[\"steps\"][:100]\n",
    "obs = trainer.reset()\n",
    "eps = 1\n",
    "# env.run([\"simple_agent\", \"simple_agent\"])\n",
    "\n",
    "# exit()\n",
    "maxstep=10000\n",
    "for i in range(maxstep + 1):\n",
    "\n",
    "    action = ['m u_1 n']\n",
    "    \n",
    "#     print(\"=== Episode {} - Step {} === \".format(eps, i + 1))\n",
    "    obs, reward, done, info = trainer.step(action)\n",
    "#     print({\"reward\": reward, \"obs_reward\": obs[\"reward\"]})\n",
    "    # print(repr(game_state.game_map.map[0][0]))\n",
    "    # update the game state for next step\n",
    "    # game_state._update(obs[\"updates\"])\n",
    "    if done:\n",
    "#         with open(\"replay.json\", \"w\") as f:\n",
    "#             f.write(env.render(mode=\"json\"))\n",
    "        # if episode is done, reset env and get new observations and new game state\n",
    "        obs = trainer.reset()\n",
    "        # game_state = Game()\n",
    "        # game_state._initialize(obs[\"updates\"])\n",
    "        eps += 1\n",
    "        print(\"=== Episode {}\".format(eps))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "cubic-mentor",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "109.89010989010988"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "10000/(60+31)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "administrative-yield",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.18 s ± 162 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "steps = env.run([agent, agent])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "productive-knitting",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
