{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualizing sensory landscapes\n",
    "\n",
    "Sensory landscapes of diverse modalities are available :\n",
    "- Olfactory (Odorscape)\n",
    "- Temperature (Thermoscape)\n",
    "- Wind (Windscape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initialize the larvaworld registry. This loads some components from disc and builds the rest on the fly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from larvaworld.lib import reg, sim, util\n",
    "\n",
    "# Tutorial safety switches (set True to run interactive demos)\n",
    "RUN_SINGLE_DEMO = False\n",
    "RUN_ALL_COMBINATIONS = False\n",
    "RUN_OLD_METHOD = False\n",
    "SAVE_MEDIA = False\n",
    "MEDIA_DIR = \"./media\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from larvaworld.lib.param.composition import Odor\n",
    "\n",
    "\n",
    "def get_odorscape(m):\n",
    "    if m == \"Diffusion\":\n",
    "        oS = reg.gen.DiffusionValueLayer(gaussian_sigma=(0.95, 0.5), evap_const=0.9)\n",
    "        oR = Odor.oD(id=\"Odor_R\")\n",
    "        oL = Odor.oD(id=\"Odor_L\")\n",
    "    elif m == \"Gaussian\":\n",
    "        oS = reg.gen.GaussianValueLayer()\n",
    "        oR = Odor.oG(id=\"Odor_R\")\n",
    "        oL = Odor.oG(id=\"Odor_L\")\n",
    "    else:\n",
    "        raise ValueError(\"Not implemented\")\n",
    "\n",
    "    kws = {\n",
    "        \"group\": \"Source\",\n",
    "        \"radius\": 0.003,\n",
    "        \"amount\": 0.0,\n",
    "    }\n",
    "\n",
    "    sus = {\n",
    "        **reg.gen.Food(\n",
    "            unique_id=\"Source_L\", c=\"blue\", odor=oL, pos=(-0.01, 0.0), **kws\n",
    "        ).entry(),\n",
    "        **reg.gen.Food(\n",
    "            unique_id=\"Source_R\", c=\"cyan\", odor=oR, pos=(0.01, 0.0), **kws\n",
    "        ).entry(),\n",
    "    }\n",
    "\n",
    "    return util.AttrDict(\n",
    "        {\"food_params\": reg.gen.FoodConf(source_units=sus), \"odorscape\": oS}\n",
    "    )\n",
    "\n",
    "\n",
    "def get_windscape(m):\n",
    "    kws = {\n",
    "        \"duration\": 5,\n",
    "        \"speed\": 50,\n",
    "    }\n",
    "\n",
    "    Npuffs = 10\n",
    "    if m == \"single\":\n",
    "        puffs = {\n",
    "            i: reg.gen.AirPuff(\n",
    "                direction=i / Npuffs * 2 * np.pi, start_time=5 + 10 * i, **kws\n",
    "            ).nestedConf\n",
    "            for i in range(Npuffs)\n",
    "        }\n",
    "        ws = 0.0\n",
    "    elif m == \"repetitive\":\n",
    "        puffs = {\n",
    "            \"puff_group\": reg.gen.AirPuff(\n",
    "                direction=np.pi / 4, start_time=5, N=Npuffs, interval=10.0, **kws\n",
    "            ).nestedConf\n",
    "        }\n",
    "        ws = 0.0\n",
    "    elif m == \"no\":\n",
    "        puffs = {}\n",
    "        ws = 10.0\n",
    "    else:\n",
    "        raise ValueError(\"Not implemented\")\n",
    "\n",
    "    wS = reg.gen.WindScape(wind_direction=0.0, wind_speed=ws, puffs=puffs)\n",
    "\n",
    "    bl = {\"Border\": reg.gen.Border(vertices=[(-0.03, 0.02), (0.03, 0.02)])}\n",
    "    bl = {}\n",
    "    return util.AttrDict({\"windscape\": wS, \"border_list\": bl})\n",
    "\n",
    "\n",
    "def get_env(Om=None, Pm=None):\n",
    "    dO = get_odorscape(m=Om) if Om is not None else {}\n",
    "    dW = get_windscape(m=Pm) if Pm is not None else {}\n",
    "    return reg.gen.Env(**dO, **dW)\n",
    "\n",
    "\n",
    "def get_id(Om=None, Pm=None, Wm=None):\n",
    "    if Om is not None and Pm is None:\n",
    "        filename = f\"{Om}_odorscape\"\n",
    "    elif Om is None and Pm is not None:\n",
    "        filename = f\"{Pm}_air-puffs_variable_wind_{Wm}\"\n",
    "    else:\n",
    "        filename = f\"{Om}_odorscape_{Pm}_air-puffs_variable_wind_{Wm}\"\n",
    "    return filename"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# New method\n",
    "\n",
    "\n",
    "def run_scape(Om=None, Pm=None, Wm=None, duration=0.15, **kwargs):\n",
    "    def func1(model):\n",
    "        model.windscape.set_wind_direction((model.t / 10 / np.pi) % (2 * np.pi))\n",
    "\n",
    "    def func2(model):\n",
    "        model.windscape.wind_speed = model.t % 100\n",
    "\n",
    "    dic = {\"direction\": func1, \"speed\": func2}\n",
    "\n",
    "    func = dic[Wm] if Wm in dic.keys() else None\n",
    "\n",
    "    id = get_id(Om=Om, Pm=Pm, Wm=Wm)\n",
    "    env = get_env(Om=Om, Pm=Pm)\n",
    "    env.visualize(\n",
    "        duration=duration,\n",
    "        id=id,\n",
    "        screen_kws={\n",
    "            \"vis_mode\": \"video\",\n",
    "            \"save_video\": SAVE_MEDIA,\n",
    "            \"media_dir\": MEDIA_DIR,\n",
    "            \"video_file\": id,\n",
    "        },\n",
    "        func=func,\n",
    "        **kwargs,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if RUN_SINGLE_DEMO:\n",
    "    # Run single scape\n",
    "\n",
    "    run_scape(Om=\"Diffusion\", Pm=\"repetitive\", Wm=\"speed\", duration=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if RUN_ALL_COMBINATIONS:\n",
    "    # Run all scape combinations\n",
    "\n",
    "    # exp='Wind&Odorscape visualization'\n",
    "    # exp='dish'\n",
    "\n",
    "    for Om in [\"Gaussian\", \"Diffusion\", None]:\n",
    "        for Pm in [None, \"single\", \"repetitive\", \"no\"]:\n",
    "            for Wm in [None, \"direction\", \"speed\", \"no\"]:\n",
    "                if Pm is None and Wm is not None:\n",
    "                    continue\n",
    "                elif Pm is not None and Wm is None:\n",
    "                    continue\n",
    "                elif Om is None and Pm is None:\n",
    "                    continue\n",
    "                run_scape(Om=Om, Pm=Pm, Wm=Wm)\n",
    "                # print(conf.screen_kws.video_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if RUN_OLD_METHOD:\n",
    "    # Old method\n",
    "\n",
    "    def run_scape(Om=None, Pm=None, Wm=None, duration=0.3, **kwargs):\n",
    "        p = util.AttrDict({\"env_params\": get_env(Om=Om, Pm=Pm).nestedConf})\n",
    "        id = get_id(Om=Om, Pm=Pm, Wm=Wm)\n",
    "        conf = util.AttrDict(\n",
    "            {\n",
    "                \"id\": id,\n",
    "                \"parameters\": p,\n",
    "                \"screen_kws\": {\n",
    "                    \"vis_mode\": \"video\",\n",
    "                    \"fps\": 10,\n",
    "                    \"show_display\": True,\n",
    "                    \"save_video\": True,\n",
    "                    \"media_dir\": \"./media\",\n",
    "                    \"video_file\": id,\n",
    "                    \"odor_aura\": True,\n",
    "                },\n",
    "            }\n",
    "        )\n",
    "\n",
    "        m = sim.base_run.BaseRun(\n",
    "            **conf, runtype=\"Exp\", experiment=\"dish\", duration=duration, **kwargs\n",
    "        )\n",
    "        m.build_env(m.p.env_params)\n",
    "        # print(m.p.steps)\n",
    "        m.sim_setup(steps=m.p.steps, seed=None)\n",
    "        if Om is not None:\n",
    "            m.odor_layers[\"Odor_L\"].visible = True\n",
    "            # env.odor_layers['Odor_R'].visible = True\n",
    "            # env.screen_manager.odor_aura = True\n",
    "        if Pm is not None:\n",
    "            m.windscape.visible = True\n",
    "        while m.running:\n",
    "            if Wm == \"direction\":\n",
    "                m.windscape.set_wind_direction((m.t / 10 / np.pi) % (2 * np.pi))\n",
    "            elif Wm == \"speed\":\n",
    "                m.windscape.wind_speed = m.t % 100\n",
    "            m.sim_step()\n",
    "        m.end()\n",
    "        m.screen_manager.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from larvaworld.lib.model.envs import *\n",
    "\n",
    "(\n",
    "    GaussianValueLayer.default_color,\n",
    "    WindScape.default_color,\n",
    "    DiffusionValueLayer.default_color,\n",
    "    ThermoScape.default_color,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# GaussianValueLayer.__bases__[0].default_color='yellow'\n",
    "\n",
    "OdorScape().default_color, GaussianValueLayer().default_color\n",
    "# OdorScape.default_color='grey'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "larvaworld_venv_new",
   "language": "python",
   "name": "larvaworld_venv_new"
  },
  "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.10.12"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}