{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Replaying experiments\n",
    "\n",
    "In this tutorial, we will demonstrate how to reconstruct and visualize previous experiments from stored data. \n",
    "\n",
    "Let's import the relevant classes :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext param.ipython\n",
    "import panel as pn\n",
    "\n",
    "import larvaworld\n",
    "from larvaworld.lib import reg, util\n",
    "\n",
    "# Import the simulation class\n",
    "from larvaworld.lib.sim import ReplayRun\n",
    "\n",
    "# Import the configuration classes\n",
    "from larvaworld.lib.reg.generators import ReplayConf, ReplayConfUnit, ReplayConfGroup\n",
    "\n",
    "larvaworld.VERBOSE = 1\n",
    "\n",
    "# Tutorial safety switches (avoid GUI / media generation by default)\n",
    "RUN_REPLAY_DEMOS = False\n",
    "RUN_DEBUG_CELL = False\n",
    "SAVE_MEDIA = False\n",
    "MEDIA_DIR = \"./media/replay\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Replay Configuration class\n",
    "\n",
    "The configuration parameters passed to the simulation class are managed by three classes :\n",
    "\n",
    "- ReplayConfGroup : group-level parameters\n",
    "- ReplayConfUnit : agent-level parameters\n",
    "- ReplayConf : all parameters (including the above)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show the attributes of the ReplayConfGroup class\n",
    "%params ReplayConfGroup\n",
    "\n",
    "# Show the attributes of the ReplayConfGroup class as a nested dictionary\n",
    "ReplayConfGroup.param"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show the attributes of the ReplayConfUnit class\n",
    "%params ReplayConfUnit\n",
    "\n",
    "# Show the attributes of the ReplayConfUnit class as a nested dictionary\n",
    "ReplayConfUnit.param"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show the attributes of the ReplayConf class\n",
    "%params ReplayConf\n",
    "\n",
    "# Show the attributes of the ReplayConf class as a nested dictionary\n",
    "ReplayConf.param"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Replay examples\n",
    "\n",
    "Now we will specify the dataset to be reconstructed by its unique ID.\n",
    "\n",
    "It is also possible to locate it by the directory where it is stored "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reg.conf.Ref.confIDs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "refID = reg.default_refID\n",
    "# Load full data only when running demos\n",
    "d = reg.loadRef(id=refID, load=(RUN_REPLAY_DEMOS or RUN_DEBUG_CELL))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will specify a number of configuration sets as dictionaries :"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "replay_confs = {\n",
    "    \"normal\": {\"time_range\": (0, 60)},\n",
    "    \"dispersal\": {\"transposition\": \"origin\"},\n",
    "    \"fixed_point\": {\n",
    "        \"agent_ids\": [0],\n",
    "        \"close_view\": True,\n",
    "        \"fix_point\": 6,\n",
    "        \"time_range\": (80, 100),\n",
    "    },\n",
    "    \"fixed_segment\": {\n",
    "        \"agent_ids\": [0],\n",
    "        \"close_view\": True,\n",
    "        \"fix_point\": 6,\n",
    "        \"fix_segment\": \"rear\",\n",
    "        \"time_range\": (100, 130),\n",
    "    },\n",
    "    \"fixed_overlap\": {\n",
    "        \"agent_ids\": [0],\n",
    "        \"close_view\": True,\n",
    "        \"fix_point\": 6,\n",
    "        \"fix_segment\": \"front\",\n",
    "        \"overlap_mode\": True,\n",
    "    },\n",
    "    \"2segs\": {\"draw_Nsegs\": 2, \"time_range\": (80, 100)},\n",
    "    \"all_segs\": {\"draw_Nsegs\": 11, \"time_range\": (80, 100)},\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if RUN_DEBUG_CELL:\n",
    "    # This needs debugging\n",
    "    import numpy as np\n",
    "\n",
    "    d.step_data = d.align_trajectories(transposition=\"origin\", replace=True)\n",
    "    xy_max = 2 * np.max(\n",
    "        d.step_data[util.nam.xy(d.c.point)].dropna().abs().values.flatten()\n",
    "    )\n",
    "\n",
    "    p = ReplayConf(**replay_confs[\"dispersal\"]).nestedConf\n",
    "    dd, bg = d.smaller_dataset(p=p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if RUN_REPLAY_DEMOS:\n",
    "    # A method that runs the replay simulation\n",
    "    def run_replay(mode):\n",
    "        p = ReplayConf(refID=refID, **replay_confs[mode]).nestedConf\n",
    "\n",
    "        screen_kws = {\n",
    "            \"show_display\": False,\n",
    "            \"vis_mode\": \"video\" if SAVE_MEDIA else None,\n",
    "            \"save_video\": SAVE_MEDIA,\n",
    "            \"media_dir\": f\"{MEDIA_DIR}/{mode}\",\n",
    "            \"video_file\": f\"{refID}_replay_{mode}\",\n",
    "        }\n",
    "\n",
    "        rep = ReplayRun(\n",
    "            parameters=p,\n",
    "            id=f\"{refID}_replay_{mode}\",\n",
    "            dir=f\"{MEDIA_DIR}/{mode}\",\n",
    "            screen_kws=screen_kws,\n",
    "        )\n",
    "        _ = rep.run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if RUN_REPLAY_DEMOS:\n",
    "    # Run a normal replay of the dataset\n",
    "    run_replay(\"normal\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if RUN_REPLAY_DEMOS:\n",
    "    # Run a reconstructed dispersal experiment where trajectories have benn transposed to the origin\n",
    "    run_replay(\"dispersal\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if RUN_REPLAY_DEMOS:\n",
    "    # Substitute the larva body contour by a bisegmental body\n",
    "    run_replay(\"2segs\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if RUN_REPLAY_DEMOS:\n",
    "    # ... or by a body with all segments, making use of all the midline points available\n",
    "    run_replay(\"all_segs\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if RUN_REPLAY_DEMOS:\n",
    "    # Now let's examine a single individual. Fixate a midline point of the larva body to the arena center\n",
    "    run_replay(\"fixed_point\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if RUN_REPLAY_DEMOS:\n",
    "    # Now fixate a midline segment along the y axis\n",
    "    run_replay(\"fixed_segment\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if RUN_REPLAY_DEMOS:\n",
    "    # And collapse the entire video to a single image to visualize the flexibility of each segment\n",
    "    run_replay(\"fixed_overlap\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "larvaworld_autoversioning",
   "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}