{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc4d0a41-153b-4ed9-a171-6ca5bf1cd9fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import geopandas as gpd\n",
    "import numpy as np\n",
    "import gurobipy as gp\n",
    "from gurobipy import GRB\n",
    "import networkx as nx\n",
    "import itertools\n",
    "import sys\n",
    "sys.path.append('..')\n",
    "from security_game.target import Target\n",
    "\n",
    "from utils.visualization import plot_animal_locations\n",
    "\n",
    "from security_game.green_security_game import GreenSecurityGame\n",
    "from security_game.infra_security_game import InfraSecurityGame\n",
    "\n",
    "from solvers.mip import mip\n",
    "from solvers.nash import nash\n",
    "from solvers.double_oracle import double_oracle\n",
    "from solvers.double_oracle_sf import double_oracle_sf\n",
    "from solvers.no_regret import regret_matching\n",
    "from solvers.simple_sse_lp import solve_sse_lp\n",
    "from solvers.nfg_sse_lp import solve_general_sum_normal_form\n",
    "import time\n",
    "from matplotlib import pyplot as plt\n",
    "from collections import defaultdict, Counter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5352bbf7-f245-45a3-ac18-546273648ac5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import csv\n",
    "df = pd.read_csv('etosha.csv', header = 0, delimiter=\",\", quoting=csv.QUOTE_NONE, encoding='utf-8')\n",
    "df = df[[\"individual-local-identifier\", \"location-lat\", \"location-long\", \"timestamp\"]]\n",
    "df.columns = [\"animal_id\",\"lat\", \"long\", \"timestamp\"]\n",
    "\n",
    "lat_min, lon_min = -19.416379222301813, 16.225642353847586\n",
    "lat_max, lon_max = -19.062243389734533, 16.834273063191954\n",
    "\n",
    "coordinate_rectangle = [lat_min, lat_max, lon_min, lon_max]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "aee209b9-764c-47a1-bc64-a35b81f3cfd3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[-19.416379222301813,\n",
       " -19.062243389734533,\n",
       " 16.225642353847586,\n",
       " 16.834273063191954]"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "coordinate_rectangle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "284e0121-3007-46bc-88f9-495cd5997032",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 444268 characters]",
      "text/plain": [
       "<Figure size 1000x800 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_animal_locations(df, coordinate_rectangle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d623665-60c3-4b32-91f9-0bb04eb8a7c1",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "halali = (-19.036839154412093, 16.471708788452602)\n",
    "omutati = (-19.31668834814194, 16.87791731016715)\n",
    "ondongab = (-19.205402906436444, 16.194226231381293)\n",
    "\n",
    "schedule_form_kwargs = {\n",
    "    \"schedule_form\": True,\n",
    "    \"simple\": False,\n",
    "    \"attacker_penalty_factor\": 5,\n",
    "    \"defender_penalty_factor\": 5\n",
    "}\n",
    "\n",
    "general_sum_kwargs = {\n",
    "    \"general_sum\": False,\n",
    "    \"attacker_animal_value\":  1, \n",
    "    \"defender_animal_value\": 1, \n",
    "    \"defender_step_cost\": 0, \n",
    "}\n",
    "\n",
    "gsg = GreenSecurityGame(df, coordinate_rectangle, \"centroid\", num_clusters=7, num_rows=6, num_columns=6, escape_line_points=None)\n",
    "gsg.generate(num_attackers=1, num_defenders=2, \n",
    "             home_base_assignments=[(halali, omutati, ondongab),(halali, omutati, ondongab)], \n",
    "             num_timesteps=8, \n",
    "             generate_utility_matrix=True, \n",
    "             defense_time_threshold=2, \n",
    "             generate_actions=True, \n",
    "             force_return=True,  \n",
    "             **schedule_form_kwargs,**general_sum_kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "394e3255-dfd4-4ed6-a0fc-ead776ac63d3",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    }
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 89592 characters]",
      "text/plain": [
       "<Figure size 1200x1000 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "gsg.draw_graph()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "de8b8b16-93a7-46f5-b557-0a6ab7879abb",
   "metadata": {},
   "outputs": [],
   "source": [
    "D_a, D_d, u = nash(gsg.schedule_form_dict[\"defender_utility_matrix\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "9d49228e-d21d-4ddf-bdd8-83039d7fcc03",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(53, 0.05470733511771794),\n",
       " (63, 0.06126649436876683),\n",
       " (87, 0.024068249067060954),\n",
       " (115, 0.24611125886401253),\n",
       " (121, 0.3904464753516927),\n",
       " (144, 0.19292395373083557),\n",
       " (146, 0.030476233499913463)]"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[(i,p) for i,p in enumerate(D_d) if p != 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "87be3441-82a1-4886-83ac-b6f0b61e2187",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-0.2437658680941743"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "u"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8151a8ba-8d4e-4f1e-af08-b7e2dc3b7a8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "gsg.to_pkl(\"etosha_sfg_gsg_1.pkl\",schedule_form=True)"
   ]
  }
 ],
 "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.12.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
