{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Launch a virtual experiment\n",
    "\n",
    "This tutorial demonstrates how to develop your own custom behavioral modules.\n",
    "\n",
    "Let's walk through the steps to create two custom behavioral module implementations: a custom olfactor and a custom thermosensor.\n",
    "\n",
    "First define all required imports: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext param.ipython\n",
    "from param.ipython import ParamPager\n",
    "import params.IPython\n",
    "\n",
    "import larvaworld as lw\n",
    "from larvaworld.lib import reg, sim\n",
    "from larvaworld.lib.reg.generators import ExpConf\n",
    "\n",
    "# Tutorial safety switches (avoid GUI/media by default)\n",
    "RUN_CUSTOM_MODULE_DEMO = False\n",
    "RUN_GUI_DEMO = False\n",
    "SAVE_MEDIA = False\n",
    "MEDIA_DIR = \"videos\"\n",
    "\n",
    "lw.VERBOSE = 1\n",
    "\n",
    "# Interfaces we're about to customize\n",
    "from larvaworld.lib.model.modules.sensor import Olfactor, Thermosensor\n",
    "from larvaworld.lib.model.modules.turner import Turner\n",
    "from larvaworld.lib.model.modules.module_modes import BrainModuleDB\n",
    "\n",
    "import random"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Behavioral modules in larvaworld are the building blocks used to run and control an individual agents. Larvaworld tries to be as modular as possible to allow you to write your own custom modules and plug them in.\n",
    "\n",
    "In essence a behavioral module is a Python class that maps some input state to some output state. The output state of one module may then be fed as input state to some other behavioral module. To see how computed states of modules flow through the modular architecture please refer to our technical paper.\n",
    "\n",
    "The basic steps to implement a custom module are:\n",
    "- create a new Python class\n",
    "- inherit from the appropriate larvaworld module base class (ie. Olfactor)\n",
    "- implement a ´update()´ function\n",
    "- compute some new output state - potentially using sensory input via ´self.input´ \n",
    "- store the computed output state in ´self.output´ property\n",
    "\n",
    "\n",
    "In the example below we will implement 2 custom modules:\n",
    "\n",
    "### Custom Olfactor\n",
    "This will read the currently sensed odor value from the sensor. The output state will be the current absolute odor value multiplied by some random number\n",
    "\n",
    "We will use that class as implementation for the olfactor sensor.\n",
    "\n",
    "### Custom Behavior Module\n",
    "This will not make use of any input. The output state will be computed as a simple random variable.\n",
    "\n",
    "We will use that module as implementation for the Thermosensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomOlfactor(Olfactor):\n",
    "    def __init__(self, **kwargs):\n",
    "        self.last_osn_activity = None\n",
    "        print(\"**** CustomOlfactor ****\")\n",
    "        print(kwargs)\n",
    "        super().__init__(**kwargs)\n",
    "\n",
    "    def update(self):\n",
    "        agent_id = (\n",
    "            self.brain.agent.unique_id if self.brain is not None else self.agent_id\n",
    "        )\n",
    "        sim_id = self.brain.agent.model.id if self.brain is not None else self.sim_id\n",
    "\n",
    "        # self.X.values() provides an array of all odor types, where the index is the odor_id\n",
    "        # lets read the currently sensed values for the first odor:\n",
    "        olfactory_input = {\n",
    "            \"odor_id\": 0,\n",
    "            # absolute concentration of 1st odor\n",
    "            \"concentration_mmol\": self.input.values()[0],\n",
    "            # change in concentration of 1st odor\n",
    "            \"concentration_change_mmol\": self.first_odor_concentration_change,\n",
    "        }\n",
    "\n",
    "        # set the output value of this custom Olfaction module\n",
    "        # that gets passed as input to other larvaworld modules\n",
    "        self.output = olfactory_input[\"concentration_mmol\"] * random.random()\n",
    "        print(\n",
    "            f\"LocalOlfactor output: {self.output} sim_id: {sim_id} agent_id: {agent_id}\"\n",
    "        )\n",
    "\n",
    "\n",
    "class CustomBehaviorModule(Thermosensor):\n",
    "    def __init__(self, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "\n",
    "    def update(self):\n",
    "        agent_id = (\n",
    "            self.brain.agent.unique_id if self.brain is not None else self.agent_id\n",
    "        )\n",
    "        sim_id = self.brain.agent.model.id if self.brain is not None else self.sim_id\n",
    "        # ignore input and use some randomness\n",
    "        self.output = random.random()\n",
    "        print(\n",
    "            f\"CustomBehaviorModule output: {self.output} sim_id: {sim_id} agent_id: {agent_id}\"\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have the implementations for our custom modules we need to instruct larvaworld to use these for the respective behavioral modules instead of larvaworld's default implementations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# overwrite mode 'osn' to use our custom CustomOlfactor class\n",
    "BrainModuleDB.BrainModuleModes.olfactor.osn = CustomOlfactor\n",
    "\n",
    "# instruct larvaworld to use our CustomBehaviorModule class as the Thermosensor\n",
    "BrainModuleDB.BrainModuleModes.thermosensor.custom = CustomBehaviorModule\n",
    "\n",
    "# for a custom turner implementation you would use:\n",
    "# BrainModuleDB.BrainModuleModes.turner.custom = CustomTurnerModule"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's almost it - all that's left is to configure an experiment to run. \n",
    "\n",
    "And instruct larvaworld to use the custom implementations when creating larva agent instances.\n",
    "\n",
    "This will be done by adjusting the model configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load predefined experiment\n",
    "expID = \"chemorbit\"\n",
    "exp_conf = reg.conf.Exp.getID(expID)\n",
    "# explore some experiment settings\n",
    "print(\n",
    "    f\"Experiment odor intensity: {exp_conf.env_params.food_params.source_units.Source.odor}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# customize odor intensity of above source\n",
    "exp_conf.env_params.food_params.source_units.Source.odor.spread = 0.01\n",
    "print(\n",
    "    f\"Adjusted odor intensity: {exp_conf.env_params.food_params.source_units.Source.odor}\"\n",
    ")\n",
    "larva_group = exp_conf.larva_groups\n",
    "# Print the entire larva group config\n",
    "print(\"Larva group configuration:\")\n",
    "print(larva_group)\n",
    "\n",
    "print(f\"Available larva group IDs/names: {larva_group.keylist}\")\n",
    "# get the \"name\" of the first larva group\n",
    "# this defaults to the larva agent model ID used by this group if none is provided\n",
    "larva_group_id = larva_group.keylist[0]\n",
    "# larva_group_id and model ID are both the same thing if no custom larva_group name is provided\n",
    "print(larva_group_id, larva_group[larva_group_id].model)\n",
    "\n",
    "# access the model_id that is used by larva agent instances created by this group:\n",
    "print(f\"model_id of group #1: {larva_group[larva_group_id].model}\")\n",
    "\n",
    "# change the model type to use by this larva group:\n",
    "exp_conf.larva_groups[larva_group_id].model = \"Levy_navigator\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fetch and update the model configuration and set the behavioral module mode's to use our custom implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# retrieve and print the model configuration of a given model ID\n",
    "mm = reg.conf.Model.getID(larva_group[larva_group_id].model)\n",
    "print(f\"larva agent model config: {mm}\")\n",
    "\n",
    "# set the model's module implementations to 'custom' to activate our custom implementations:\n",
    "mm.brain.olfactor.mode = \"osn\"\n",
    "mm.brain.thermosensor.mode = \"custom\"\n",
    "# to activate the custom turner you'd use:\n",
    "# mm.brain.turner.mode = 'custom'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally run the simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Launch a simulation run of the customized experiment\n",
    "\n",
    "run_id = \"my-custom-modules-run\"\n",
    "\n",
    "screen_kws = {}\n",
    "if RUN_GUI_DEMO or SAVE_MEDIA:\n",
    "    screen_kws = {\n",
    "        \"vis_mode\": \"video\",\n",
    "        \"show_display\": RUN_GUI_DEMO,\n",
    "        \"save_video\": SAVE_MEDIA,\n",
    "        \"fps\": 20,\n",
    "        \"video_file\": f\"larva-sim-{run_id}\",\n",
    "        \"media_dir\": MEDIA_DIR,\n",
    "    }\n",
    "\n",
    "erun = sim.ExpRun(\n",
    "    experiment=expID,\n",
    "    modelIDs=[\"navigator\", \"Levy_navigator\"],\n",
    "    screen_kws=screen_kws,\n",
    "    N=2,\n",
    "    duration=0.5,  # minutes\n",
    ")\n",
    "\n",
    "if RUN_CUSTOM_MODULE_DEMO:\n",
    "    erun.simulate()\n",
    "    print(f\"Run_id: {run_id} completed\")\n",
    "    # erun.analyze()"
   ]
  }
 ],
 "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.13.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}