{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "## Pathfinding\n\nXarray-spatial provides A\\* pathfinding for finding optimal routes through raster surfaces.\nPaths can be computed using geometric distance alone (unweighted) or weighted by a\nfriction/cost surface, which makes the algorithm find the *least-cost* path rather\nthan the geometrically shortest one.\n\n`multi_stop_search` extends A\\* to visit a sequence of waypoints, stitching segments\ninto a single cumulative-cost surface. It can optionally reorder interior waypoints\nto minimize total travel cost (TSP).\n\nAll four backends are supported:\n\n| Backend | Strategy |\n|---------|----------|\n| **NumPy** | Numba-jitted kernel (fast, in-memory) |\n| **Dask** | Sparse Python A\\* with LRU chunk cache — loads chunks on demand so the full grid never needs to fit in RAM |\n| **CuPy** | CPU fallback (transfers to numpy, runs the numba kernel, transfers back) |\n| **Dask+CuPy** | Same sparse A\\* as Dask, with automatic cupy-to-numpy chunk conversion |\n\n**Note:** `snap_start` and `snap_goal` are not supported with Dask-backed arrays\nbecause the brute-force nearest-pixel scan is O(h*w). Ensure the start and goal\npixels are valid before calling `a_star_search`.\n\n- [Unweighted A\\*](#Unweighted-A*): find the shortest geometric path through a line network\n- [Weighted A\\*](#Weighted-A*): find the least-cost path through a friction surface\n- [Multi-Stop Search](#Multi-Stop-Search): route through multiple waypoints with optional reordering\n- [Dask (out-of-core)](#Dask-(out-of-core)): pathfinding on datasets that don't fit in RAM"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Importing Packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "import numpy as np\nimport pandas as pd\nimport xarray as xr\n\nimport datashader as ds\nfrom datashader.transfer_functions import shade\nfrom datashader.transfer_functions import stack\nfrom datashader.transfer_functions import dynspread\nfrom datashader.transfer_functions import set_background\nfrom datashader.colors import Elevation\n\nimport xrspatial\nfrom xrspatial import a_star_search, multi_stop_search, generate_terrain, slope, cost_distance"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Unweighted A*\n",
    "\n",
    "A\\* is an informed search algorithm that finds the shortest path from a start\n",
    "to a goal through a graph.  In the unweighted case, edge costs are purely\n",
    "geometric (Euclidean distance scaled by cellsize).\n",
    "\n",
    "We'll demonstrate this on a synthetic line raster with barriers."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create a line raster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define range of x and y\n",
    "xrange = (0, 4)\n",
    "yrange = (0, 4)\n",
    "\n",
    "# create line raster\n",
    "ys = [1, 1, 3, 3, 1, 1, np.nan, 1, 3, np.nan, 1, 3, np.nan, 1, 3, np.nan, 2, 2]\n",
    "xs = [1, 3, 3, 1, 1, 3, np.nan, 1, 3, np.nan, 3, 1, np.nan, 1, 2, np.nan, 1, 3]\n",
    "line_df = pd.DataFrame(dict(x=xs, y=ys))\n",
    "\n",
    "W = 800\n",
    "H = 600\n",
    "cvs = ds.Canvas(plot_width=W, plot_height=H, x_range=xrange, y_range=yrange)\n",
    "line_agg = cvs.line(line_df, x=\"x\", y=\"y\").astype(int)\n",
    "line_shaded = dynspread(shade(line_agg, cmap=[\"black\", \"salmon\"]))\n",
    "\n",
    "# pick start and goal locations\n",
    "start = (3, 1)\n",
    "goal = (1, 3)\n",
    "\n",
    "start_df = pd.DataFrame({\"x\": [start[1]], \"y\": [start[0]]})\n",
    "start_agg = cvs.points(start_df, \"x\", \"y\")\n",
    "start_shaded = dynspread(shade(start_agg, cmap=[\"red\"]), threshold=1, max_px=5)\n",
    "\n",
    "goal_df = pd.DataFrame({\"x\": [goal[1]], \"y\": [goal[0]]})\n",
    "goal_agg = cvs.points(goal_df, \"x\", \"y\")\n",
    "goal_shaded = dynspread(shade(goal_agg, cmap=[\"lime\"]), threshold=1, max_px=5)\n",
    "\n",
    "set_background(stack(line_shaded, start_shaded, goal_shaded), \"black\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 8-connectivity shortest path\n",
    "\n",
    "Cells with value 0 (off the lines) are barriers.  The path is highlighted in green."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path_agg_8 = a_star_search(\n",
    "    line_agg, start, goal, barriers=[0], snap_start=True, snap_goal=True\n",
    ")\n",
    "\n",
    "path_shaded = dynspread(shade(path_agg_8, cmap=[\"green\"]))\n",
    "set_background(stack(line_shaded, path_shaded, start_shaded, goal_shaded), \"black\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 4-connectivity shortest path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path_agg_4 = a_star_search(\n",
    "    line_agg, start, goal, barriers=[0],\n",
    "    snap_start=True, snap_goal=True, connectivity=4\n",
    ")\n",
    "\n",
    "path_shaded = dynspread(shade(path_agg_4, cmap=[\"green\"]))\n",
    "set_background(stack(line_shaded, path_shaded, start_shaded, goal_shaded), \"black\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Weighted A*\n",
    "\n",
    "When a **friction surface** is provided, A\\* finds the *least-cost* path rather\n",
    "than the geometrically shortest one.  Each edge cost becomes\n",
    "`geometric_distance * mean_friction_of_endpoints`, the same cost model used\n",
    "by `cost_distance()`.\n",
    "\n",
    "This is useful for terrain routing: a path over flat ground may be longer\n",
    "in distance but much cheaper than a short path over steep slopes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Generate terrain and friction surface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "# Generate a terrain raster\nW, H = 600, 400\ntemplate = xr.DataArray(np.zeros((H, W)), dims=['y', 'x'])\nterrain = generate_terrain(template, seed=12345)\n\nshade(terrain, cmap=Elevation, how=\"linear\")"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "# Derive friction from slope: steeper = higher cost\nterrain_slope = slope(terrain)\n\n# slope() produces NaN at edges; fill with 0 so friction = 1 there (passable, flat)\n# Add 1 so flat areas have friction=1 rather than 0\nfriction = terrain_slope.fillna(0.0) + 1.0\n\nshade(friction, cmap=[\"lightgreen\", \"yellow\", \"red\"], how=\"eq_hist\")"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Find paths: unweighted vs weighted\n",
    "\n",
    "We pick a start and goal on opposite sides of a ridge, then compare\n",
    "the path found with and without friction weighting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Choose start and goal coordinates\n",
    "start_coord = (400.0, 50.0)\n",
    "goal_coord = (100.0, 450.0)\n",
    "\n",
    "# Unweighted path (geometric distance only)\n",
    "path_unweighted = a_star_search(terrain, start_coord, goal_coord)\n",
    "\n",
    "# Weighted path (friction from slope)\n",
    "path_weighted = a_star_search(terrain, start_coord, goal_coord, friction=friction)\n",
    "\n",
    "print(f\"Unweighted path cost at goal: {np.nanmax(path_unweighted):.2f}\")\n",
    "print(f\"Weighted path cost at goal:   {np.nanmax(path_weighted):.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Visualise both paths on the terrain\n",
    "\n",
    "Red = unweighted (geometrically shortest), green = weighted (least-cost).\n",
    "Notice how the weighted path avoids steep ridges even if it's longer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "terrain_shaded = shade(terrain, cmap=Elevation, how=\"linear\", alpha=180)\n",
    "\n",
    "unweighted_shaded = dynspread(\n",
    "    shade(path_unweighted, cmap=[\"red\", \"red\"], how=\"linear\", min_alpha=255),\n",
    "    threshold=1, max_px=2,\n",
    ")\n",
    "weighted_shaded = dynspread(\n",
    "    shade(path_weighted, cmap=[\"lime\", \"lime\"], how=\"linear\", min_alpha=255),\n",
    "    threshold=1, max_px=2,\n",
    ")\n",
    "\n",
    "stack(terrain_shaded, unweighted_shaded, weighted_shaded)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Validate against cost_distance\n",
    "\n",
    "For a single source, the A\\* path cost at the goal should match the\n",
    "`cost_distance()` value at that same pixel.  This confirms both algorithms\n",
    "use the same cost model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a source raster with a single target at the start pixel\n",
    "source_raster = terrain.copy(data=np.zeros_like(terrain.data))\n",
    "# Find the start pixel index\n",
    "from xrspatial.pathfinding import _get_pixel_id\n",
    "spy, spx = _get_pixel_id(start_coord, terrain)\n",
    "source_raster.data[spy, spx] = 1.0\n",
    "\n",
    "cd = cost_distance(source_raster, friction)\n",
    "\n",
    "# Compare: A* cost at goal vs cost_distance at goal\n",
    "gpy, gpx = _get_pixel_id(goal_coord, terrain)\n",
    "cd_val = float(cd.values[gpy, gpx])\n",
    "astar_val = float(path_weighted.values[gpy, gpx])\n",
    "\n",
    "print(f\"cost_distance at goal: {cd_val:.4f}\")\n",
    "print(f\"A* path cost at goal:  {astar_val:.4f}\")\n",
    "print(f\"Match: {np.isclose(cd_val, astar_val, rtol=1e-5)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "source": "### Multi-Stop Search\n\n`multi_stop_search` routes through a list of waypoints in order, calling A\\* for\neach consecutive pair and stitching the segments into one cumulative-cost surface.\n\nSet `optimize_order=True` to let the solver reorder the interior waypoints (keeping\nthe first and last fixed) to minimize total travel cost. For up to 12 waypoints it\nuses exact Held-Karp; beyond that it falls back to nearest-neighbor + 2-opt.",
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": "#### Route through three waypoints\n\nWe pick three waypoints across the terrain and let `multi_stop_search` find\nthe least-cost route visiting them in order.",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "source": "# Three waypoints across the terrain (re-using the terrain and friction from above)\nwp_start = (400.0, 50.0)\nwp_mid   = (250.0, 250.0)\nwp_end   = (100.0, 450.0)\n\n# Naive order\npath_multi = multi_stop_search(\n    terrain, [wp_start, wp_mid, wp_end], friction=friction\n)\n\nprint(f\"Segment costs: {path_multi.attrs['segment_costs']}\")\nprint(f\"Total cost:    {path_multi.attrs['total_cost']:.2f}\")\n\n# Visualise the multi-stop path\nterrain_shaded = shade(terrain, cmap=Elevation, how=\"linear\", alpha=180)\nmulti_shaded = dynspread(\n    shade(path_multi, cmap=[\"cyan\", \"cyan\"], how=\"linear\", min_alpha=255),\n    threshold=1, max_px=2,\n)\nstack(terrain_shaded, multi_shaded)",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": "#### Optimize waypoint order\n\nWith four or more waypoints, the given order may not be the cheapest.\n`optimize_order=True` solves a TSP over the pairwise A\\* costs, keeping the\nfirst and last waypoints fixed.",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "source": "# Four waypoints in a deliberately suboptimal order (zigzag)\nwp_a = (400.0, 50.0)    # start (fixed)\nwp_b = (100.0, 450.0)   # far corner\nwp_c = (350.0, 200.0)   # back near start\nwp_d = (100.0, 350.0)   # end (fixed)\n\nwaypoints = [wp_a, wp_b, wp_c, wp_d]\n\n# Without optimization\npath_naive = multi_stop_search(terrain, waypoints, friction=friction)\n\n# With optimization\npath_opt = multi_stop_search(\n    terrain, waypoints, friction=friction, optimize_order=True\n)\n\nprint(f\"Naive order:     {path_naive.attrs['waypoint_order']}\")\nprint(f\"Naive cost:      {path_naive.attrs['total_cost']:.2f}\")\nprint(f\"Optimized order: {path_opt.attrs['waypoint_order']}\")\nprint(f\"Optimized cost:  {path_opt.attrs['total_cost']:.2f}\")\nprint(f\"Savings:         {path_naive.attrs['total_cost'] - path_opt.attrs['total_cost']:.2f}\")",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### References\n",
    "\n",
    "- A\\* search algorithm: https://en.wikipedia.org/wiki/A*_search_algorithm\n",
    "- Red Blob Games — A\\* implementation: https://www.redblobgames.com/pathfinding/a-star/implementation.html\n",
    "- Cost distance (weighted proximity): `xrspatial.cost_distance`"
   ]
  },
  {
   "cell_type": "markdown",
   "source": "### Dask (out-of-core)\n\nFor rasters that don't fit in memory, wrap your data in a Dask array.\nA\\* explores only the pixels it needs (typically O(path_length) to\nO(path_length^2)), loading chunks on demand through an LRU cache.\nThe output is a lazy Dask array of the same chunk structure.",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "source": "import dask.array as da\n\n# Re-use the terrain and friction from above, but as dask arrays\nterrain_dask = terrain.copy()\nterrain_dask.data = da.from_array(terrain.data, chunks=(200, 300))\n\nfriction_dask = friction.copy()\nfriction_dask.data = da.from_array(friction.data, chunks=(200, 300))\n\n# Run A* — chunks are loaded on demand, not all at once\npath_dask = a_star_search(terrain_dask, start_coord, goal_coord, friction=friction_dask)\nprint(f\"Result type: {type(path_dask.data)}\")\nprint(f\"Chunks: {path_dask.data.chunks}\")\n\n# Compute to verify it matches the in-memory result\ndask_cost = float(path_dask.values[gpy, gpx])\nprint(f\"Dask A* cost at goal:   {dask_cost:.4f}\")\nprint(f\"NumPy A* cost at goal:  {astar_val:.4f}\")\nprint(f\"Match: {np.isclose(dask_cost, astar_val, rtol=1e-5)}\")",
   "metadata": {},
   "execution_count": null,
   "outputs": []
  }
 ],
 "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.14.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}