{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true,
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Follow the simulation using websockets\n",
    "\n",
    "In this notebook we'll show you how to open a websocket connection with the simulation to\n",
    "follow its progress. Additionally we will also use a websocket to follow the waterlevels\n",
    "of a node over time.\n",
    "\n",
    "In this notebook we will use the package [websockets](https://websockets.readthedocs.io/en/stable/).\n",
    "Make sure it's installed if you want to run this notebook locally."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import json\n",
    "import asyncio\n",
    "from datetime import datetime\n",
    "from getpass import getpass\n",
    "\n",
    "import numpy as np\n",
    "import ipywidgets as widgets\n",
    "import websockets\n",
    "from IPython.display import display\n",
    "from websockets.http import Headers\n",
    "\n",
    "from threedi_api_client.api import ThreediApi\n",
    "from threedi_api_client.openapi.api.v3_api import V3Api"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Provide authentication details\n",
    "API_HOST = \"https://api.3di.live\"\n",
    "PERSONAL_API_KEY = getpass(\"Personal API token\")  # https://management.3di.live/personal_api_keys\n",
    "\n",
    "config = {\n",
    "    \"THREEDI_API_HOST\": API_HOST,\n",
    "    \"THREEDI_API_PERSONAL_API_TOKEN\": PERSONAL_API_KEY\n",
    "}\n",
    "\n",
    "api_client: V3Api = ThreediApi(config=config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Specify the threedi-model and organisation\n",
    "organisation_uuid = \"b08433fa47c1401eb9cbd4156034c679\"\n",
    "threedimodel = api_client.threedimodels_list(name__icontains='v2_bergermeer').results[0]\n",
    "\n",
    "# Retrieve first simulation template\n",
    "simulation_templates = api_client.simulation_templates_list(simulation__threedimodel__id=threedimodel.id)\n",
    "assert simulation_templates.count > 0, f\"No simulation templates found for threedimodel {threedimodel.name}\"\n",
    "simulation_template_id = simulation_templates.results[0].id\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "simulation = api_client.simulations_from_template(\n",
    "    data={\n",
    "        \"template\": simulation_template_id,\n",
    "        \"name\": \"just some simulation\",\n",
    "        \"organisation\": organisation_uuid,\n",
    "        \"start_datetime\": datetime.now(),\n",
    "        \"duration\": 3600  # in seconds, so we simulate for 1 hour\n",
    "    }\n",
    ")\n",
    "print(simulation)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The simulation needs to be initialized before we can start a websocket connection.\n",
    "We'll initialize the simulation but we won't start it yet."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "api_client.simulations_actions_create(\n",
    "    simulation_pk=simulation.id, data={\"name\": \"initialize\"}\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "The websocket uses JWT authentication so lets get our JWT token."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "token = api_client.api_client.configuration.get_basic_auth_token()\n",
    "headers = Headers(authorization=token)\n",
    "uri = f'wss://api.3di.live/v3/simulations/{simulation.id}/'"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "First we create the progress-bar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "progress = widgets.IntProgress(value=0, min=0, max=simulation.duration)\n",
    "\n",
    "async def update_progress_bar():\n",
    "    async with websockets.connect(uri, extra_headers=headers) as websocket:\n",
    "        print(\"Connected to the websocket\")\n",
    "        async for message in websocket:\n",
    "            message = json.loads(message)\n",
    "            if message.get(\"type\") == \"time\":\n",
    "                progress.value = message['data']['time']\n",
    "        print(\"Websocket connection closed\")\n",
    "\n",
    "asyncio.tasks.create_task(update_progress_bar())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Next we print the waterlevels over time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "waterlevel_websocket_url = api_client.simulations_visualisations_water_level_graph_create(\n",
    "    simulation_pk=simulation.id,\n",
    "    data={\"start_time\": 0, \"subscribe\": True, \"node_id\": 1, \"subscribe_rate_limit\": 2}\n",
    ")\n",
    "\n",
    "async def print_waterlevels():\n",
    "    async with websockets.connect(waterlevel_websocket_url.url, extra_headers=headers) as websocket:\n",
    "        print(\"Connected to the websocket\")\n",
    "        async for message in websocket:\n",
    "            data = np.frombuffer(message, dtype=np.float32)\n",
    "            print(f\"Time: {int(data[0]):<4} waterlevel: {data[1]}\")\n",
    "        print(\"Websocket connection closed\")\n",
    "\n",
    "asyncio.tasks.create_task(print_waterlevels())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Once we start the simulation we should see updates in the progress bar and the waterlevel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "display(progress)\n",
    "api_client.simulations_actions_create(\n",
    "    simulation_pk=simulation.id, data={\"name\": \"start\"}\n",
    ")"
   ]
  }
 ],
 "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}