{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0f22d444-d81e-4723-a365-1f801d0e93ef",
   "metadata": {},
   "source": [
    "# Agent-based Transportaion Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "355b8300-db9a-4a23-9e48-ab8952b70779",
   "metadata": {},
   "source": [
    "Let's create a simple agent-based transportation model. First, start with a transportation network. Our transportation network can be created as a Networkx Graph or a (Geo)Pandas (Geo)DataFrame. There are some sample transportation networks available in dpd.mapping.samples. Real-life networks can be imported from OpenStreetMap via pyrosm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16f97dfd-48b5-4be6-a5b0-ab59dcc9aac9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from networkx import draw\n",
    "\n",
    "from dpd.mapping.samples import cross_box\n",
    "\n",
    "graph = cross_box\n",
    "\n",
    "pos = {\n",
    "    node: (graph.nodes()[node][\"geometry\"].x, graph.nodes()[node][\"geometry\"].y)\n",
    "    for node in graph.nodes\n",
    "}\n",
    "\n",
    "node_color = [\n",
    "    (\n",
    "        \"red\"\n",
    "        if graph.nodes()[node].get(\"type\") == \"stop_sign\"\n",
    "        else (\n",
    "            \"yellow\"\n",
    "            if graph.nodes()[node].get(\"type\") == \"yield_sign\"\n",
    "            else (\n",
    "                \"blue\"\n",
    "                if graph.nodes()[node].get(\"type\") == \"stop_light\"\n",
    "                else \"orange\" if graph.nodes()[node].get(\"type\") == \"stop\" else \"green\"\n",
    "            )\n",
    "        )\n",
    "    )\n",
    "    for node in graph.nodes()\n",
    "]\n",
    "\n",
    "draw(graph, pos=pos, node_color=node_color, with_labels=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c69a066-c66d-4b7a-bb1a-21f7489fb7cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from geopandas import GeoDataFrame\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "edges_df = GeoDataFrame(\n",
    "    [graph.edges[edge] for edge in graph.edges], index=list(graph.edges)\n",
    ")\n",
    "nodes_df = GeoDataFrame([graph.nodes[node] for node in graph.nodes])\n",
    "\n",
    "fig = plt.figure(figsize=(9, 8))\n",
    "ax = fig.add_subplot(111)\n",
    "\n",
    "nodes_df[\"geometry\"].plot(ax=ax, color=node_color, markersize=1000)\n",
    "nodes_df.apply(\n",
    "    lambda x: ax.annotate(text=x.name, xy=x.geometry.coords[0], ha=\"center\", size=20),\n",
    "    axis=1,\n",
    ")\n",
    "edges_df[\"geometry\"].plot(ax=ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b452d47-5619-4c32-940c-29bac9101399",
   "metadata": {},
   "source": [
    "networkx provides the ability to compute a path from any node to another node. When using OpenStreetMap, the same can be accomplished via the Open Source Routing Machine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d8382e4-2a6f-430c-b106-7d4a92cbfda2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from networkx import shortest_path\n",
    "\n",
    "node_ids = shortest_path(graph, 0, 1)\n",
    "print(\"Node IDs:\", node_ids)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18926268-367e-4e4f-a7d2-df46a48dbac5",
   "metadata": {},
   "source": [
    "To create a couple agents, we can create some transportation zones. For simplicity, we will create one zone per node with three Production and three Attraction so each zone has one person that goes to each other zone. Below are some other DataFrames we can generate from the Zones DataFrame. Also, we will create a path for each person using networkx."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5619a4c-c801-45cf-8bcb-1a7d962f8a9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dpd.modeling import Zones\n",
    "\n",
    "zones = Zones(\n",
    "    data=[\n",
    "        {\"Name\": \"Zone 0\", \"Production\": 3, \"Attraction\": 3},\n",
    "        {\"Name\": \"Zone 1\", \"Production\": 3, \"Attraction\": 3},\n",
    "        {\"Name\": \"Zone 2\", \"Production\": 3, \"Attraction\": 3},\n",
    "        {\"Name\": \"Zone 3\", \"Production\": 3, \"Attraction\": 3},\n",
    "    ],\n",
    "    index=nodes_df.index,\n",
    ")\n",
    "zones[\"geometry\"] = nodes_df[\"geometry\"]\n",
    "zones"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16026c7b-4bc2-4099-9734-3e1c38422d0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "distance_dataframe = zones.calculate_distance_dataframe()\n",
    "distance_dataframe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74a56222-87f0-43b9-96e8-cffae5807ea4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy\n",
    "\n",
    "from dpd.modeling import TripDataFrame\n",
    "\n",
    "trip_dataframe = TripDataFrame(\n",
    "    data=numpy.ones([4, 4]), index=zones.index, columns=zones.index\n",
    ").map(int)\n",
    "\n",
    "trip_dataframe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23f487ff-612c-40a4-bc6a-c7f6e60fe07f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dpd.modeling import Population\n",
    "\n",
    "population = Population.from_trip_dataframe(trip_dataframe)\n",
    "population = population[\n",
    "    population.origin != population.destination\n",
    "]  # remove rows where the origin and destination zone are equal\n",
    "population"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a993bfd3-3413-4618-a4c7-062eaa265a65",
   "metadata": {},
   "outputs": [],
   "source": [
    "population[\"node_ids\"] = population.apply(\n",
    "    lambda x: shortest_path(graph, x[\"origin\"], x[\"destination\"]), axis=1\n",
    ")\n",
    "population"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2031ee07-f156-4e4f-99f6-bd3e65f45466",
   "metadata": {},
   "source": [
    "Next, we setup and run our agent-based model. We need to tranform our transportation network in to Python objects for Edges and Nodes. Again, this can be done with either a Graph or a (Geo)DataFrame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d60bf60f-cc1e-4ebb-8520-7473258c5100",
   "metadata": {},
   "outputs": [],
   "source": [
    "from uuid import uuid4\n",
    "\n",
    "from dpd.driving import EdgesLanesNodesDriver\n",
    "from dpd.mapping import add_object_to_edges_and_nodes\n",
    "from dpd.mapping.nodes import NodeModel\n",
    "from dpd.mechanics import KinematicBodyWithAcceleration\n",
    "from dpd.mechanics.datacollection import BODY_AGENT_REPORTERS\n",
    "from dpd.modeling import TransportationModel\n",
    "\n",
    "body_model = TransportationModel(\n",
    "    agent_reporters=BODY_AGENT_REPORTERS | {\"geometry\": \"geometry\"}\n",
    ")\n",
    "\n",
    "node_model = NodeModel()\n",
    "\n",
    "graph = add_object_to_edges_and_nodes(graph, node_model)\n",
    "\n",
    "for index, row in population.iterrows():\n",
    "    kbwas = KinematicBodyWithAcceleration(\n",
    "        initial_acceleration=0.1,\n",
    "        initial_velocity=0.1,\n",
    "        initial_position=0,\n",
    "        max_deceleration=0.1,\n",
    "        min_velocity=0,\n",
    "        unique_id=uuid4(),\n",
    "        model=body_model,\n",
    "    )\n",
    "    elnd = EdgesLanesNodesDriver.from_node_ids(\n",
    "        nodes_dict=graph.nodes,\n",
    "        edges_dict=graph.edges,\n",
    "        node_ids=row[\"node_ids\"],\n",
    "        body=kbwas,\n",
    "        driver_final_velocity=0,\n",
    "        unique_id=uuid4(),\n",
    "        model=body_model,\n",
    "    )\n",
    "    body_model.schedule.add(elnd)\n",
    "\n",
    "while body_model.running:\n",
    "    body_model.step()\n",
    "    node_model.step()\n",
    "\n",
    "df = GeoDataFrame(body_model.get_dataframe())\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a61fb878-121f-478d-912d-bfe6734f07a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "\n",
    "from geopandas import GeoDataFrame\n",
    "from movingpandas import TrajectoryCollection\n",
    "from pandas import to_timedelta\n",
    "\n",
    "START_TIME = datetime(1970, 1, 1, 0, 0, 0)\n",
    "timedelta = to_timedelta(df.index.levels[0], unit=\"s\")\n",
    "index = START_TIME + timedelta\n",
    "gdf = df\n",
    "gdf.index = gdf.index.set_levels(index, level=0)\n",
    "gdf.reset_index(level=\"AgentID\", inplace=True)\n",
    "tc = TrajectoryCollection(gdf, \"AgentID\")\n",
    "tc.add_speed()\n",
    "tc.hvplot(line_width=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd4cb4f1-d922-4ee4-8c2b-31b75320cde2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "\n",
    "def random_color():\n",
    "    return (\n",
    "        \"#\"\n",
    "        + hex(random.randint(0, 0xFF))[2:]\n",
    "        + hex(random.randint(0, 0xFF))[2:]\n",
    "        + hex(random.randint(0, 0xFF))[2:]\n",
    "    )\n",
    "\n",
    "\n",
    "features = []\n",
    "for trajectory in tc.trajectories:\n",
    "    color = random_color()\n",
    "    df = trajectory.df.copy()\n",
    "    df[\"previous_geometry\"] = df[\"geometry\"].shift()\n",
    "    df[\"time\"] = df.index\n",
    "    df[\"previous_time\"] = df[\"time\"].shift()\n",
    "    for _, row in df.iloc[1:].iterrows():\n",
    "        coordinates = [\n",
    "            [row[\"previous_geometry\"].xy[0][0], row[\"previous_geometry\"].xy[1][0]],\n",
    "            [row[\"geometry\"].xy[0][0], row[\"geometry\"].xy[1][0]],\n",
    "        ]\n",
    "        times = [row[\"previous_time\"].isoformat(), row[\"time\"].isoformat()]\n",
    "        features.append(\n",
    "            {\n",
    "                \"type\": \"Feature\",\n",
    "                \"geometry\": {\n",
    "                    \"type\": \"LineString\",\n",
    "                    \"coordinates\": coordinates,\n",
    "                },\n",
    "                \"properties\": {\n",
    "                    \"times\": times,\n",
    "                    \"style\": {\n",
    "                        \"color\": color,\n",
    "                        \"weight\": 5,\n",
    "                    },\n",
    "                },\n",
    "            }\n",
    "        )\n",
    "\n",
    "import folium\n",
    "from folium.plugins import TimestampedGeoJson\n",
    "\n",
    "m = folium.Map(location=[0, 0], zoom_start=16)\n",
    "\n",
    "TimestampedGeoJson(\n",
    "    {\n",
    "        \"type\": \"FeatureCollection\",\n",
    "        \"features\": features,\n",
    "    },\n",
    "    period=\"PT1S\",\n",
    "    add_last_point=True,\n",
    "    transition_time=1000,\n",
    ").add_to(m)\n",
    "\n",
    "m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30a658ed-a8e9-48f7-aed8-90cfe5f8c12f",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}