{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Driving"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some other relations to use are:\n",
    "* expo_line_west: 2351006\n",
    "* blue_line_north: 2351005\n",
    "* c_line_east: 2377889\n",
    "* s1_ost_flug: 1773071\n",
    "* bart_red_south: 2851509\n",
    "* caltrain_local: 2768041\n",
    "* san_bernardino_line: 1287719\n",
    "* orange_county_line: 2812899\n",
    "* ca_hsr: 7739486\n",
    "\n",
    "If using BART, set `gague = 1.676 * meter` when creating the Route."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Maybe Network is just a Route with a MultiIndex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pandas import isna\n",
    "\n",
    "from dpd.driving import Route\n",
    "\n",
    "relation = 1287719  # san_bernardino_line\n",
    "route = Route.from_osm_relation(relation=relation)\n",
    "\n",
    "for name in [\n",
    "    \"San Bernardino–Downtown\",\n",
    "    \"Rialto\",\n",
    "    \"Fontana\",\n",
    "    \"Upland\",\n",
    "    \"Claremont\",\n",
    "    \"Pomona–North\",\n",
    "    \"Baldwin Park\",\n",
    "    \"El Monte\",\n",
    "]:\n",
    "    route.remove_stop(name)\n",
    "\n",
    "route.drop([0, 1, 2], inplace=True)\n",
    "route[\"type\"] = route[\"name\"].map(lambda x: \"node\" if isna(x) else \"stop\")\n",
    "route[\"dwell_time\"] = route[\"type\"].map(lambda x: 45 if x == \"stop\" else None)\n",
    "route[\"distance_to_point\"] = route.distance_to_point\n",
    "route.stops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from folium import Map\n",
    "\n",
    "map_ = Map(location=[34, -117.75], zoom_start=10)\n",
    "\n",
    "route.geometry.explore(m=map_)\n",
    "route.stops.geometry.explore(m=map_, color=\"black\", marker_kwds={\"radius\": 20})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dpd.driving import Vehicles\n",
    "\n",
    "vehicles = Vehicles(\n",
    "    url=\"https://raw.githubusercontent.com/davidbailey/Notes/trunk/Trains.csv\",\n",
    "    vehicle=\"Stadler KISS\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from uuid import uuid4\n",
    "\n",
    "from astropy.constants import g0\n",
    "from astropy.units import meter, second\n",
    "\n",
    "from dpd.driving import EdgesLanesNodesDriver\n",
    "from dpd.driving.datacollection import DYNAMIC_BODY_DRIVER_REPORTERS\n",
    "from dpd.mapping.edges import object_for_edge\n",
    "from dpd.mapping.nodes import NodeModel\n",
    "from dpd.modeling import TransportationModel\n",
    "\n",
    "mu = 0.25  # https://railroads.dot.gov/sites/fra.dot.gov/files/fra_net/17468/A%20Survey%20of%20Wheel-Rail%20Friction.pdf\n",
    "\n",
    "node_model = NodeModel(velocity_unit=meter / second)\n",
    "\n",
    "edges = route.edges\n",
    "edges[\"object\"] = edges.apply(object_for_edge, axis=1)\n",
    "route[\"object\"] = route.apply(node_model.object_for_node, axis=1)\n",
    "\n",
    "body_model = TransportationModel(\n",
    "    time_unit=second / 1,\n",
    "    agent_reporters=DYNAMIC_BODY_DRIVER_REPORTERS | {\"geometry\": \"geometry\"},\n",
    ")\n",
    "\n",
    "vehicle = vehicles.vehicle(\n",
    "    initial_acceleration=mu * g0,\n",
    "    max_acceleration=mu * g0,\n",
    "    max_deceleration=1 * meter / second**2,\n",
    "    initial_position=0 * meter,\n",
    "    initial_velocity=0 * meter / second,\n",
    "    unique_id=uuid4(),\n",
    "    model=body_model,\n",
    ")\n",
    "\n",
    "driver = EdgesLanesNodesDriver.from_node_ids(\n",
    "    edges_dict=edges.T,\n",
    "    nodes_dict=route.T,\n",
    "    node_ids=route.index,\n",
    "    body=vehicle,\n",
    "    driver_final_velocity=0 * meter / second,\n",
    "    driver_max_velocity=vehicle.max_velocity,\n",
    "    unique_id=uuid4(),\n",
    "    model=body_model,\n",
    ")\n",
    "\n",
    "body_model.schedule.add(driver)\n",
    "\n",
    "while body_model.running:\n",
    "    body_model.step()\n",
    "    node_model.step()\n",
    "\n",
    "# run the data collector one more time because we never get the final data point\n",
    "body_model.schedule.add(driver)\n",
    "body_model.datacollector.collect(body_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Add to beginning of edges in the driver to get geometry to work.\n",
    "\"\"\"\n",
    "\n",
    "from shapely.geometry import LineString, Point\n",
    "\n",
    "LineString([Point(0, 0), Point(0, 0)]).interpolate(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pandas import merge\n",
    "\n",
    "df = body_model.get_dataframe()\n",
    "df = merge(\n",
    "    df,\n",
    "    route.stops[[\"name\", \"distance_to_point\"]],\n",
    "    left_on=\"position\",\n",
    "    right_on=\"distance_to_point\",\n",
    "    how=\"left\",\n",
    ").set_index(df.index)\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Redo Trip to support MultiIndex and multiple Trips, Trip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "\n",
    "from dpd.driving import Trip\n",
    "\n",
    "trip = Trip.from_model(df, route)\n",
    "fig = plt.figure(figsize=(18, 16))\n",
    "ax = fig.add_subplot(111)\n",
    "trip.plot_schedule(ax=ax)\n",
    "trip.stops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "\n",
    "trip.index = trip.index + datetime(year=1970, month=1, day=1)\n",
    "trip[\"distance\"] = trip[\"total_distance\"].map(lambda x: x.value)\n",
    "trajectory = trip.to_trajectory(index=0)\n",
    "trajectory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from astropy.units import m, s\n",
    "from pandas import Series\n",
    "\n",
    "trajectory.add_speed(overwrite=True)\n",
    "speed = trajectory.df[\"speed\"]\n",
    "speed.index = df.position\n",
    "speed.index = speed.index.map(lambda x: x.value)\n",
    "\n",
    "speed_limits = Series(\n",
    "    map(lambda x: x.value, route.speed_limits), index=route.distance_to_point[:-1]\n",
    ")\n",
    "speed_limits.index.name = \"position\"\n",
    "speed_limits.index = speed_limits.index.map(lambda x: x.value)\n",
    "\n",
    "fig = plt.figure(figsize=(18, 16))\n",
    "ax = fig.add_subplot(111)\n",
    "speed.plot(ax=ax)\n",
    "speed_limits.plot(ax=ax)\n",
    "ax.set_ylim([0, 50])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dpd.driving import Schedule\n",
    "\n",
    "schedule = Schedule.from_trip(trip)\n",
    "\n",
    "schedule.schedule[schedule.schedule.index.isin([\"Departure\"], level=1)].applymap(\n",
    "    lambda x: x.round(\"s\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from numpy import concatenate\n",
    "\n",
    "route_reversed = Route(route.reversed[[\"geometry\", \"name\", \"object\"]].copy())\n",
    "route_reversed[\"distance_to_point\"] = route_reversed.distance_to_point\n",
    "route_reversed.stops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "edges = route_reversed.edges\n",
    "edges[\"object\"] = edges.apply(object_for_edge, axis=1)\n",
    "\n",
    "body_model = TransportationModel(\n",
    "    time_unit=second / 1, agent_reporters=DYNAMIC_BODY_DRIVER_REPORTERS\n",
    ")\n",
    "\n",
    "vehicle = vehicles.vehicle(\n",
    "    initial_acceleration=mu * g0,\n",
    "    max_acceleration=mu * g0,\n",
    "    max_deceleration=1 * meter / second**2,\n",
    "    initial_position=0 * meter,\n",
    "    initial_velocity=0 * meter / second,\n",
    "    unique_id=uuid4(),\n",
    "    model=body_model,\n",
    ")\n",
    "\n",
    "driver = EdgesLanesNodesDriver.from_node_ids(\n",
    "    edges_dict=edges.T,\n",
    "    nodes_dict=route_reversed.T,\n",
    "    node_ids=route_reversed.index,\n",
    "    body=vehicle,\n",
    "    driver_final_velocity=0 * meter / second,\n",
    "    driver_max_velocity=vehicle.max_velocity,\n",
    "    unique_id=uuid4(),\n",
    "    model=body_model,\n",
    ")\n",
    "\n",
    "body_model.schedule.add(driver)\n",
    "\n",
    "while body_model.running:\n",
    "    body_model.step()\n",
    "    node_model.step()\n",
    "\n",
    "# run the data collector one more time because we never get the final data point\n",
    "body_model.schedule.add(driver)\n",
    "body_model.datacollector.collect(body_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = body_model.get_dataframe()\n",
    "df = merge(\n",
    "    df,\n",
    "    route_reversed.stops[[\"name\", \"distance_to_point\"]],\n",
    "    left_on=\"position\",\n",
    "    right_on=\"distance_to_point\",\n",
    "    how=\"left\",\n",
    ").set_index(df.index)\n",
    "\n",
    "trip_reversed = Trip.from_model(df, route, include_geometry=False)\n",
    "trip_reversed.index = trip_reversed.index + datetime(year=1970, month=1, day=1)\n",
    "\n",
    "schedule_reversed = Schedule.from_trip(trip_reversed)\n",
    "schedule_reversed.reverse_distance()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(18, 16))\n",
    "ax = fig.add_subplot(111)\n",
    "schedule.plot_schedule(ax=ax)\n",
    "schedule_reversed.plot_schedule(ax=ax)\n",
    "schedule.trips[list(schedule.trips.keys())[0]].plot_schedule(ax=ax)  # to add labels\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gtfs_kit\n",
    "\n",
    "feed = gtfs_kit.read_feed(\n",
    "    \"http://www.bart.gov/dev/schedules/google_transit.zip\", dist_units=\"mi\"\n",
    ")\n",
    "\n",
    "trip = Trip.from_gtfs(feed, trip_id=\"1508825\")\n",
    "from astropy import units\n",
    "\n",
    "units.imperial.enable()\n",
    "trip[\"total_distance\"] = trip.total_distance.map(\n",
    "    lambda x: (x.value * units.imperial.foot).to(units.meter)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "\n",
    "fig = plt.figure(figsize=(18, 16))\n",
    "ax = fig.add_subplot(111)\n",
    "trip.plot_schedule(ax=ax)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dpd.driving import Schedule\n",
    "\n",
    "schedule = Schedule.from_gtfs(feed, route_id=\"6\", direction_id=0)\n",
    "schedule_reversed = Schedule.from_gtfs(feed, route_id=\"5\", direction_id=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from astropy import units\n",
    "\n",
    "units.imperial.enable()\n",
    "for trip_id in schedule.trips:\n",
    "    trip = schedule.trips[trip_id]\n",
    "    trip[\"total_distance\"] = trip.total_distance.map(\n",
    "        lambda x: (x.value * units.imperial.foot).to(units.meter)\n",
    "    )\n",
    "for trip_id in schedule_reversed.trips:\n",
    "    trip = schedule_reversed.trips[trip_id]\n",
    "    trip[\"total_distance\"] = trip.total_distance.map(\n",
    "        lambda x: (x.value * units.imperial.foot).to(units.meter)\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "schedule_reversed.reverse_distance()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "\n",
    "fig = plt.figure(figsize=(18, 16))\n",
    "ax = fig.add_subplot(111)\n",
    "schedule.plot_schedule(ax=ax)\n",
    "schedule.trips[list(schedule.trips.keys())[-1]].plot_schedule(ax=ax)  # to add labels\n",
    "schedule_reversed.plot_schedule(ax=ax)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "schedule_reversed.schedule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": 4
}