{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Staging: Ophyd's Hook for Device-Specific Setup and Cleanup\n",
    "\n",
    "In this tutorial we will see how the `stage` and `unstage` methods can be used to get a device ready and then return it to its previous state.\n",
    "\n",
    "## Set up for tutorial\n",
    "\n",
    "First, let's ensure our simulated IOCs are running.\n",
    "\n",
    "The IOCs may already be running in the background. Run this command to verify\n",
    "that they are running: it should produce output with STARTING or RUNNING on each line.\n",
    "In the event of a problem, edit this command to replace `status` with `restart all` and run again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!../supervisor/start_supervisor.sh status"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Motivation\n",
    "\n",
    "There are many situations when we want to ensure that one or more PVs are set up in a well-defined and safe state, prior to doing things such as triggering a detector or collecting data. Likewise, once we are done, we may want to return everything back to the way it was, to get it back to a neutral state rather than leaving it primed for the specific task we were doing.\n",
    "\n",
    "In Bluesky, this sort of priming or \"setting the stage\" is conventionally done using the `stage` and `unstage` methods.\n",
    "\n",
    "This is typically used with detectors, but it can be used with any Device.\n",
    "\n",
    "Consider the following example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ophyd import Device, Component, EpicsSignal, EpicsSignalRO, set_and_wait, DeviceStatus\n",
    "import time\n",
    "\n",
    "class TriggeredDetector(Device):\n",
    "    \"\"\"\n",
    "    A detector that requires triggering\n",
    "    \"\"\"\n",
    "    gain = Component(EpicsSignal, ':gain', kind='config')\n",
    "    exposure_time = Component(EpicsSignal, ':exposure_time', kind='config')\n",
    "    reading = Component(EpicsSignalRO, ':reading', kind='normal')\n",
    "    acquire = Component(EpicsSignal, ':acquire', kind='omitted', put_complete=True)\n",
    "    enabled = Component(EpicsSignal, ':enabled', kind='omitted')\n",
    "\n",
    "    def trigger(self):\n",
    "        \"\"\"\n",
    "        Trigger the detector and return a Status object.\n",
    "        \"\"\"\n",
    "        status = DeviceStatus(self)\n",
    "        self.acquire.put(1, callback=status._finished)\n",
    "        return status\n",
    "    \n",
    "    def stage(self):\n",
    "        \"\"\"\n",
    "        Set the stage, storing previous value for later.\n",
    "        \"\"\"\n",
    "        self.initial_enabled_state = self.enabled.get()\n",
    "        set_and_wait(self.enabled, 1)\n",
    "        return super().stage()\n",
    "    \n",
    "    def unstage(self):\n",
    "        \"\"\"\n",
    "        Restore previously saved state.\n",
    "        \"\"\"\n",
    "        ret =  super().unstage()\n",
    "        set_and_wait(self.enabled, self.initial_enabled_state)\n",
    "        return ret\n",
    "    \n",
    "triggered_detector = TriggeredDetector('trigger_with_pc', name='triggered_detector')\n",
    "triggered_detector"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When bluesky obtains a reading from some `device` it typically:\n",
    "\n",
    "* Calls `device.trigger()` and receives back a status object\n",
    "* Waits for that status object to complete (while potentially doing other things, like triggering other detectors in parallel)\n",
    "* Calls `device.read()`\n",
    "\n",
    "If it obtains multiple readings in sequence, it repeats this trigger/wait/read cycle. \n",
    "\n",
    "Here we can see simple examples of `stage` and `unstage` methods, which can be used to perform a choreographed sequence of steps necessary to make `device` ready for use and some corresponding sequence to put in back safely into a resting state. Bluesky plans typically call `device.stage()` *once* before first using a device in a plan and then `device.unstage()` at the end. There are mechanisms to ensure that, even if a plan is interrupted by the user or by an exception is raised, `device.unstage()` will be called *if* `device.stage()` was previously called as part of a plan executed by `RE`.\n",
    "\n",
    "**Important:** Note the use of `set_and_wait` in these methods. Unlike many of the other mechanisms in bluesky, we do not want these methods to return until we are certain that the desired state is fully set and complete. So we use those blocking calls to ensure the PV values are set to the indicated values.\n",
    "\n",
    "There is no mechanism to automatically call `device.stage()` just because it was defined. You must invoke it explicitly, either as part of a plan or manually from the command line. If invoked manually, you must also manually `unstage`. Let's have a look at what this might look like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up our \"pre-stage\" state to have the `enabled` signal at 0\n",
    "triggered_detector.enabled.put(0)\n",
    "triggered_detector.enabled.get()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Staging should set it to 1, and ensure it is fully set prior to returning\n",
    "triggered_detector.stage()\n",
    "triggered_detector.enabled.get()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Trigger the detector and obtain status object to monitor\n",
    "status = triggered_detector.trigger()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Wait until it is done prior to reading\n",
    "while not status.done:\n",
    "    time.sleep(0.01)\n",
    "    \n",
    "triggered_detector.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We can trigger and read again; RE would typically do this in a loop internally\n",
    "status = triggered_detector.trigger()\n",
    "\n",
    "while not status.done:\n",
    "    time.sleep(0.01)\n",
    "    \n",
    "triggered_detector.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Unstage when we are done\n",
    "triggered_detector.unstage()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's check that `enabled` really went back to its previous state:\n",
    "triggered_detector.enabled.get()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A convenient shorthand for common simple cases: `stage_sigs`\n",
    "\n",
    "Storing and restoring individual signal values by hand is a bit tedious. The base implementations of `stage` and `unstage` allow us to do this with far less manual labor. In fact, in this simple scenario, we can do away with overriding these methods entirely: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TriggeredDetector(Device):\n",
    "    \"\"\"\n",
    "    A detector that requires triggering\n",
    "    \"\"\"\n",
    "    gain = Component(EpicsSignal, ':gain', kind='config')\n",
    "    exposure_time = Component(EpicsSignal, ':exposure_time', kind='config')\n",
    "    reading = Component(EpicsSignalRO, ':reading', kind='hinted')\n",
    "    acquire = Component(EpicsSignal, ':acquire', kind='omitted', put_complete=True)\n",
    "    enabled = Component(EpicsSignal, ':enabled', kind='omitted')\n",
    "    \n",
    "    def __init__(self, *args, **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.stage_sigs['enabled'] = 1  # OrderedDict mapping component name to desired state\n",
    "\n",
    "    def trigger(self):\n",
    "        \"\"\"\n",
    "        Trigger the detector and return a Status object.\n",
    "        \"\"\"\n",
    "        status = DeviceStatus(self)\n",
    "        self.acquire.put(1, callback=status._finished)\n",
    "        return status\n",
    "    \n",
    "triggered_detector = TriggeredDetector('trigger_with_pc', name='triggered_detector')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Keep in mind this is a fairly trivial example and, if you need to perform more sophisticated setup and cleanup, you will still need to define your own `stage` and `unstage` methods."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise\n",
    "\n",
    "Try staging the device twice in a row. Then try unstaging it twice in a row."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Customizing cleanup via `stop`, `resume`, `pause`\n",
    "\n",
    "These optional methods can be used to further customize a Device's cleanup:\n",
    "\n",
    "* `stop` -- called by bluesky when a plan is paused or exits (successfully or in error)\n",
    "* `pause` -- called when the RunEngine is paused\n",
    "* `resume` -- called when the RunEngine resumes from a pause"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TriggeredDetector(Device):\n",
    "    \"\"\"\n",
    "    A detector that requires triggering\n",
    "    \"\"\"\n",
    "    gain = Component(EpicsSignal, ':gain', kind='config')\n",
    "    exposure_time = Component(EpicsSignal, ':exposure_time', kind='config')\n",
    "    reading = Component(EpicsSignalRO, ':reading', kind='hinted')\n",
    "    acquire = Component(EpicsSignal, ':acquire', kind='omitted', put_complete=True)\n",
    "    enabled = Component(EpicsSignal, ':enabled', kind='omitted')\n",
    "    \n",
    "    def __init__(self, *args, **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.stage_sigs['enabled'] = 1  # OrderedDict mapping component name to desired state\n",
    "\n",
    "    def trigger(self):\n",
    "        \"\"\"\n",
    "        Trigger the detector and return a Status object.\n",
    "        \"\"\"\n",
    "        status = DeviceStatus(self)\n",
    "        self.acquire.put(1, callback=status._finished)\n",
    "        return status\n",
    "    \n",
    "    def resume(self):\n",
    "        ...\n",
    "        \n",
    "    def pause(self):\n",
    "        ...\n",
    "        \n",
    "    def stop(self, success=False):\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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}