{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6312bf3e-4b94-4c8b-8958-c1a5d7205984",
   "metadata": {},
   "source": [
    "# Free-Space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcb09255-de20-47ef-a89f-5e00558b700d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import pathlib\n",
    "\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1f4721c-212c-4c79-a22c-5711f03c9ca9",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p img\n",
    "\n",
    "plt.rcParams.update(\n",
    "    {\n",
    "        \"text.usetex\": True,\n",
    "        \"font.family\": \"sans-serif\",\n",
    "        \"figure.figsize\": (3, 2.5),\n",
    "        \"figure.dpi\": 140,\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6012f989-b9c1-4990-93e8-6eb99fc9666b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_trajectories_data(file_name=None, path=None):\n",
    "    path = pathlib.Path().resolve() if path is None else pathlib.Path(path)\n",
    "    file_name = path / file_name\n",
    "    assert file_name.is_file(), file_name\n",
    "\n",
    "    trajectories = json.load(open(file_name, \"r\"))[\"trajectories\"]\n",
    "    return [np.array(trj) for trj in trajectories]\n",
    "\n",
    "\n",
    "TRAJECTORIES = {\n",
    "    \"p\": load_trajectories_data(\"p.json\"),\n",
    "    \"q\": load_trajectories_data(\"q.json\"),\n",
    "}\n",
    "len(TRAJECTORIES[\"p\"]), len(TRAJECTORIES[\"q\"]), TRAJECTORIES[\"q\"][0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e98db75f-c5a7-4a1b-bc2d-ad55b4a25571",
   "metadata": {},
   "outputs": [],
   "source": [
    "TEST_TRAJECTORIES = {\n",
    "    \"p\": np.array([[[0, 0], [3, 4], [7, 4], [7, 8], [5, 3], [9, 3], [11, 5]]]),\n",
    "    \"q\": np.array([[[3, 0], [1, 6], [9, 6], [11, 8]]]),\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d65c9ad-13f5-466e-a698-1c5a8e503ffe",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_length(p):\n",
    "    n = [trj.shape[0] for trj in p]\n",
    "\n",
    "    fig, ax = plt.subplots()\n",
    "    ax.hist(n, rwidth=0.9)\n",
    "    ax.set_xlabel(\"Trajectory length\")\n",
    "    ax.set_ylabel(\"Entries\")\n",
    "\n",
    "    fig.savefig(\"img/trj_length.png\")\n",
    "\n",
    "\n",
    "plot_length(TRAJECTORIES[\"p\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "843b7fd9-3654-496e-b90e-9369fa2171b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_trajectories(p, q, *, p_idx, show_legend=False):\n",
    "    fig, ax = plt.subplots()\n",
    "\n",
    "    qx, qy = q[0].T\n",
    "    ax.plot(qx[0], qy[0], \"k.\", alpha=0.5)\n",
    "    ax.plot(qx[-1], qy[-1], \"k.\", alpha=0.5)\n",
    "    ax.plot(qx, qy, \"k\", alpha=0.5, label=\"$q$\")\n",
    "\n",
    "    for n, i in enumerate(p_idx, start=1):\n",
    "        px, py = p[i].T\n",
    "        ax.plot(px[0], py[0], f\"C{n}.\", alpha=0.5)\n",
    "        ax.plot(px[-1], py[-1], f\"C{n}.\", alpha=0.5)\n",
    "        ax.plot(px, py, f\"C{n}-\", alpha=0.5, label=f\"$p_{i}$\")\n",
    "\n",
    "    ax.set_aspect(\"equal\")\n",
    "    ax.set_xlabel(\"$x$\")\n",
    "    ax.set_ylabel(\"$y$\")\n",
    "\n",
    "    if show_legend:\n",
    "        ax.legend()\n",
    "\n",
    "    fig.tight_layout()\n",
    "    fig.savefig(\"img/trajectories.png\")\n",
    "\n",
    "\n",
    "plot_trajectories(TEST_TRAJECTORIES[\"p\"], TEST_TRAJECTORIES[\"q\"], p_idx=[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e501ff35-94e0-411f-9df8-e776a6de9a7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_trajectories(TRAJECTORIES[\"p\"], TRAJECTORIES[\"q\"], p_idx=[0, 1, 2, 3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6f5358e-2070-4902-9bbe-6b80b272c231",
   "metadata": {},
   "outputs": [],
   "source": [
    "def interpolate(points, *, n):\n",
    "    t = np.linspace(0, 1, n, endpoint=False)[:, np.newaxis]\n",
    "    return np.concatenate(\n",
    "        [a + (b - a) * t for a, b in zip(points[:-1], points[1:])]\n",
    "        + [np.array([points[-1]])]\n",
    "    )\n",
    "\n",
    "\n",
    "def metric(p, q):\n",
    "    px, py = p.T\n",
    "    qx, qy = q.T\n",
    "\n",
    "    return np.hypot(np.subtract.outer(px, qx), np.subtract.outer(py, qy))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "758f5131-ab4b-4c52-9ab0-1a7e3c01e6b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def frechet_distance(d):\n",
    "    P, Q = d.shape\n",
    "\n",
    "    d = np.copy(d)\n",
    "    d[:, 0] = np.maximum.accumulate(d[:, 0])\n",
    "    d[0, :] = np.maximum.accumulate(d[0, :])\n",
    "\n",
    "    path = np.zeros((P, Q, 2), dtype=np.int64)\n",
    "    path[1:, 0, 0] = np.arange(P - 1)\n",
    "    path[0, 1:, 1] = np.arange(Q - 1)\n",
    "\n",
    "    for i in range(1, P):\n",
    "        for j in range(1, Q):\n",
    "            idx = np.array([[i - 1, j], [i - 1, j - 1], [i, j - 1]])\n",
    "\n",
    "            min_idx = np.argmin(d[*idx.T])\n",
    "            min_idx = idx[min_idx]\n",
    "\n",
    "            d[i, j] = max(d[*min_idx], d[i, j])\n",
    "            path[i, j] = min_idx\n",
    "\n",
    "    optimal_path = [np.array([P - 1, Q - 1])]\n",
    "    while not np.all(optimal_path[-1] == np.array([0, 0])):\n",
    "        optimal_path.append(path[*optimal_path[-1]])\n",
    "\n",
    "    optimal_path = np.stack(optimal_path[::-1]).astype(np.float64)\n",
    "\n",
    "    return d[-1, -1], optimal_path / optimal_path[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff8c0fde-0007-4bc3-8816-e97d350c5d76",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_distance(p, q, *, p_idx=None, show_path=True, n=1):\n",
    "    if p_idx is not None:\n",
    "        p = p[p_idx]\n",
    "\n",
    "    p = interpolate(p, n=n)\n",
    "    q = interpolate(q, n=n)\n",
    "    d = metric(p, q)\n",
    "\n",
    "    fd, path = frechet_distance(d)\n",
    "    print(f\"{fd=}\")\n",
    "\n",
    "    tp = np.linspace(0, 1, d.shape[0])\n",
    "    tq = np.linspace(0, 1, d.shape[1])\n",
    "    x, y = np.meshgrid(tp, tq)\n",
    "    z = d.T\n",
    "\n",
    "    fig, ax = plt.subplots()\n",
    "    cset = ax.contourf(x, y, z)\n",
    "    ax.contour(x, y, z, cset.levels, colors=\"k\")\n",
    "    cbar = fig.colorbar(cset)\n",
    "\n",
    "    if show_path:\n",
    "        ax.plot(path[:, 0], path[:, 1], \"--\", color=\"red\")\n",
    "\n",
    "    ax.set_xlim(-0.02, 1.02)\n",
    "    ax.set_ylim(-0.02, 1.02)\n",
    "\n",
    "    ax.set_xlabel(\n",
    "        r\"$\\alpha(t)$\" if p_idx is None else r\"$\\alpha_{\" + str(p_idx) + \"}(t)$\"\n",
    "    )\n",
    "    ax.set_ylabel(r\"$\\beta(t)$\")\n",
    "    cbar.set_label(r\"Fr\\'echet distance\")\n",
    "\n",
    "    fig.tight_layout()\n",
    "    fig.savefig(\"img/distance.png\")\n",
    "\n",
    "\n",
    "plot_distance(TEST_TRAJECTORIES[\"p\"][0], TEST_TRAJECTORIES[\"q\"][0], n=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a1025e9-3f81-43bf-bdb1-bcef86146524",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_distance(TRAJECTORIES[\"p\"][0][:30], TRAJECTORIES[\"q\"][0][:30], n=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8aaf959-4c58-469c-8fa0-5dcd0bfecad9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_freespace(p, q, *, p_idx=None, threshold=None, show_path=True, n=1):\n",
    "    if p_idx is not None:\n",
    "        p = p[p_idx]\n",
    "\n",
    "    p = interpolate(p, n=n)\n",
    "    q = interpolate(q, n=n)\n",
    "    d = metric(p, q)\n",
    "\n",
    "    fd, path = frechet_distance(d)\n",
    "    if threshold is None:\n",
    "        threshold = fd\n",
    "\n",
    "    tp = np.linspace(0, 1, d.shape[0])\n",
    "    tq = np.linspace(0, 1, d.shape[1])\n",
    "    x, y = np.meshgrid(tp, tq)\n",
    "    z = d.T\n",
    "\n",
    "    fig, ax = plt.subplots()\n",
    "\n",
    "    norm = mpl.colors.Normalize(vmin=np.min(z), vmax=np.max(z))\n",
    "    thr = norm(threshold)\n",
    "    sdata = [\n",
    "        [0.0, 1.0, 1.0],\n",
    "        [thr - 0.03, 1.0, 1.0],\n",
    "        [thr + 0.03, 0.5, 0.5],\n",
    "        [1.0, 0.5, 0.5],\n",
    "    ]\n",
    "    my_cmap = mpl.colors.LinearSegmentedColormap(\n",
    "        \"my_cmap\",\n",
    "        segmentdata={\"red\": sdata, \"green\": sdata, \"blue\": sdata},\n",
    "        N=1024,\n",
    "    )\n",
    "    cset = ax.pcolormesh(x, y, z, cmap=my_cmap, norm=norm)\n",
    "    cbar = fig.colorbar(cset)\n",
    "\n",
    "    if show_path:\n",
    "        ax.plot(path[:, 0], path[:, 1], \"--\", color=\"red\")\n",
    "\n",
    "    ax.set_xlim(-0.02, 1.02)\n",
    "    ax.set_ylim(-0.02, 1.02)\n",
    "\n",
    "    ax.set_xlabel(\n",
    "        r\"$\\alpha(t)$\" if p_idx is None else \"$\\alpha_{\" + str(p_idx) + \"}(t)$\"\n",
    "    )\n",
    "    ax.set_ylabel(r\"$\\beta(t)$\")\n",
    "    cbar.set_label(r\"Fr\\'echet distance\")\n",
    "\n",
    "    fig.tight_layout()\n",
    "    fig.savefig(\"img/freespace.png\")\n",
    "\n",
    "\n",
    "plot_freespace(TEST_TRAJECTORIES[\"p\"][0], TEST_TRAJECTORIES[\"q\"][0], n=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31a8de2d-6201-49bb-a89c-d62763e4c1ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_freespace(TRAJECTORIES[\"p\"][0][:30], TRAJECTORIES[\"q\"][0][:30], n=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1d021c4-a273-41f2-b4ef-189d1eab9100",
   "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
}