{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "12fb5aef",
   "metadata": {},
   "source": [
    "# WIP: save reflection information in stream\n",
    "See: https://github.com/bluesky/hklpy/issues/158"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2dcbeb2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from bluesky import RunEngine\n",
    "from bluesky.callbacks.best_effort import BestEffortCallback\n",
    "import bluesky.plans as bp\n",
    "import bluesky.plan_stubs as bps\n",
    "import bluesky.preprocessors as bpp\n",
    "import databroker\n",
    "import hkl\n",
    "from hkl import *  # TODO: wildcard import, yikes!\n",
    "import numpy as np\n",
    "import pyRestTable\n",
    "from ophyd import Component, Device, EpicsSignal, Signal\n",
    "from ophyd.signal import AttributeSignal, ArrayAttributeSignal\n",
    "from ophyd.sim import *\n",
    "import pandas as pd\n",
    "\n",
    "bec = BestEffortCallback()\n",
    "bec.disable_plots()\n",
    "cat = databroker.temp().v2\n",
    "\n",
    "RE = RunEngine({})\n",
    "RE.subscribe(bec)\n",
    "RE.subscribe(cat.v1.insert)\n",
    "RE.md[\"notebook\"] = \"tst_UB_in_stream\"\n",
    "RE.md[\"objective\"] = \"Demonstrate UB matrix save & restore in stream of bluesky run\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c7fbcbc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import namedtuple\n",
    "\n",
    "Lattice = namedtuple(\"LatticeTuple\", \"a b c alpha beta gamma\")\n",
    "Position = namedtuple(\"PositionTuple\", \"omega chi phi tth\")\n",
    "Reflection = namedtuple(\"ReflectionTuple\", \"h k l position\")\n",
    "\n",
    "class Holder:\n",
    "    samples = {}\n",
    "\n",
    "class Reflections:\n",
    "    reflections = []\n",
    "\n",
    "class MyDevice(Device):\n",
    "    uptime = Component(EpicsSignal, \"gp:UPTIME\", kind=\"normal\")\n",
    "    apple = Component(Signal, value=\"Fuji\", kind=\"omitted\")\n",
    "    orange = Component(Signal, value=\"Valencia\", kind=\"omitted\")\n",
    "    octopus = Component(Signal, value=\"spotted\", kind=\"omitted\")\n",
    "\n",
    "    stream_name = Component(AttributeSignal, attr=\"_stream_name\", doc=\"stream name\", kind=\"omitted\")\n",
    "    stream_attrs = Component(AttributeSignal, attr=\"_stream_attrs\", doc=\"attributes in stream\", kind=\"omitted\")\n",
    "\n",
    "    _samples = {}\n",
    "    _stream_name = \"bozo\"\n",
    "    # _stream_attrs = \"orange octopus samples stream_name stream_attrs\".split()\n",
    "    _stream_attrs = \"orange octopus stream_name stream_attrs\".split()\n",
    "\n",
    "    # cannot make AttributeSignal from these that can be written by bluesky\n",
    "    paddle = Holder()\n",
    "    spots = Reflections()\n",
    "\n",
    "    def other_streams(self, label=None):\n",
    "        label = label or self._stream_name\n",
    "        yield from bps.create(name=label)\n",
    "        for attr in self._stream_attrs:\n",
    "            yield from bps.read(getattr(self, attr))\n",
    "        yield from bps.save()\n",
    "\n",
    "        yield from bps.create(\"fruit\")\n",
    "        yield from bps.read(self.apple)\n",
    "        yield from bps.read(self.orange)\n",
    "        yield from bps.save()\n",
    "\n",
    "        yield from bps.create(\"animal\")\n",
    "        yield from bps.read(self.octopus)\n",
    "        yield from bps.save()\n",
    "\n",
    "\n",
    "nitwit = MyDevice(\"\", name=\"nitwit\")\n",
    "nitwit.paddle.samples[\"main\"] = Lattice(1,1,1,30,60,90)\n",
    "nitwit.paddle.samples[\"second\"] = Lattice(2,2,2, 2,2,2)\n",
    "\n",
    "def try_it():\n",
    "    yield from bps.open_run()\n",
    "\n",
    "    yield from bps.create()\n",
    "    yield from bps.read(nitwit)\n",
    "    yield from bps.save()\n",
    "\n",
    "    yield from nitwit.other_streams()\n",
    "\n",
    "    yield from bps.close_run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0af36aa4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OrderedDict([('nitwit_uptime',\n",
       "              {'value': '11 days, 06:21:15', 'timestamp': 1655096814.592612})])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nitwit.wait_for_connection()\n",
    "nitwit.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "69cc74b2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "Transient Scan ID: 1     Time: 2022-06-13 00:06:55\n",
      "Persistent Unique Scan ID: '8abd6968-1ce9-443d-a4ef-f9ad98f52d6c'\n",
      "New stream: 'primary'\n",
      "+-----------+------------+\n",
      "|   seq_num |       time |\n",
      "+-----------+------------+\n",
      "|         1 | 00:06:55.2 |\n",
      "New stream: 'bozo'\n",
      "New stream: 'fruit'\n",
      "New stream: 'animal'\n",
      "+-----------+------------+\n",
      "generator try_it ['8abd6968'] (scan num: 1)\n",
      "\n",
      "\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "('8abd6968-1ce9-443d-a4ef-f9ad98f52d6c',)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# RE(bp.count([nitwit]))\n",
    "RE(try_it())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "26b80f9a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BlueskyRun\n",
       "  uid='8abd6968-1ce9-443d-a4ef-f9ad98f52d6c'\n",
       "  exit_status='success'\n",
       "  2022-06-13 00:06:55.235 -- 2022-06-13 00:06:55.253\n",
       "  Streams:\n",
       "    * animal\n",
       "    * fruit\n",
       "    * primary\n",
       "    * bozo\n"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cat[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8e616824",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "stream = 'animal'\n",
      "<xarray.Dataset>\n",
      "Dimensions:         (time: 1)\n",
      "Coordinates:\n",
      "  * time            (time) float64 1.655e+09\n",
      "Data variables:\n",
      "    nitwit_octopus  (time) <U7 'spotted'\n",
      "stream = 'fruit'\n",
      "<xarray.Dataset>\n",
      "Dimensions:        (time: 1)\n",
      "Coordinates:\n",
      "  * time           (time) float64 1.655e+09\n",
      "Data variables:\n",
      "    nitwit_orange  (time) <U8 'Valencia'\n",
      "    nitwit_apple   (time) <U4 'Fuji'\n",
      "stream = 'primary'\n",
      "<xarray.Dataset>\n",
      "Dimensions:        (time: 1)\n",
      "Coordinates:\n",
      "  * time           (time) float64 1.655e+09\n",
      "Data variables:\n",
      "    nitwit_uptime  (time) <U17 '11 days, 06:21:15'\n",
      "stream = 'bozo'\n",
      "<xarray.Dataset>\n",
      "Dimensions:              (time: 1, dim_0: 4)\n",
      "Coordinates:\n",
      "  * time                 (time) float64 1.655e+09\n",
      "Dimensions without coordinates: dim_0\n",
      "Data variables:\n",
      "    nitwit_stream_attrs  (time, dim_0) <U12 'orange' ... 'stream_attrs'\n",
      "    nitwit_octopus       (time) <U7 'spotted'\n",
      "    nitwit_stream_name   (time) <U4 'bozo'\n",
      "    nitwit_orange        (time) <U8 'Valencia'\n"
     ]
    }
   ],
   "source": [
    "run = cat[-1]\n",
    "for stream in list(run):\n",
    "    print(f\"{stream = }\")\n",
    "    print(getattr(run, stream).read())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c4a1fa76",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 1, 1, 30, 60, 90)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nitwit.paddle.samples[\"main\"][:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "38fc5de6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'nitwit_octopus': {'value': 'spotted', 'timestamp': 1655096815.1244874}}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nitwit.octopus.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "cc7c70e4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "k = 'main'\n",
      "v[:] = (1, 1, 1, 30, 60, 90)\n",
      "k = 'second'\n",
      "v[:] = (2, 2, 2, 2, 2, 2)\n"
     ]
    }
   ],
   "source": [
    "for k, v in nitwit.paddle.samples.items():\n",
    "    print(f\"{k = }\")\n",
    "    print(f\"{v[:] = }\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "01ea2641",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sig.read() = {'sig': {'value': {'main': {'a': 1, 'b': 1, 'c': 1, 'alpha': 30, 'beta': 60, 'gamma': 90}}, 'timestamp': 1655096815.6220567}}\n",
      "nitwit.paddle.samples['main'] = LatticeTuple(a=1, b=1, c=1, alpha=30, beta=60, gamma=90)\n"
     ]
    }
   ],
   "source": [
    "sig = Signal(name=\"sig\", value=dict(main=nitwit.paddle.samples[\"main\"]._asdict()))\n",
    "print(f\"{sig.read() = }\")\n",
    "print(f\"{nitwit.paddle.samples['main'] = }\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c4fbeee9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(-145.451, 0.0, 0.0, 69.0966)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "r = Reflection(4.0, 0., 0., Position(omega=-145.451, chi=0.0, phi=0.0, tth=69.0966))\n",
    "r[-1][:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "907dad59",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "Transient Scan ID: 2     Time: 2022-06-13 00:06:55\n",
      "Persistent Unique Scan ID: '3554e9e1-2afc-4384-8d53-e4108099f85b'\n",
      "New stream: 'samples'\n",
      "New stream: 'reflections'\n",
      "\n",
      "\n",
      "\n",
      "nm = 'samples'\n",
      "<xarray.Dataset>\n",
      "Dimensions:  (time: 1, dim_0: 6, dim_1: 6, dim_2: 6)\n",
      "Coordinates:\n",
      "  * time     (time) float64 1.655e+09\n",
      "Dimensions without coordinates: dim_0, dim_1, dim_2\n",
      "Data variables:\n",
      "    main     (time, dim_0) int64 1 1 1 30 60 90\n",
      "    _keys    (time, dim_1) <U5 'a' 'b' 'c' 'alpha' 'beta' 'gamma'\n",
      "    second   (time, dim_2) int64 2 2 2 2 2 2\n",
      "------------------------------\n",
      "nm = 'reflections'\n",
      "<xarray.Dataset>\n",
      "Dimensions:         (time: 1, dim_0: 4, dim_1: 3, dim_2: 4, dim_3: 2, dim_4: 4,\n",
      "                     dim_5: 3, dim_6: 4)\n",
      "Coordinates:\n",
      "  * time            (time) float64 1.655e+09\n",
      "Dimensions without coordinates: dim_0, dim_1, dim_2, dim_3, dim_4, dim_5, dim_6\n",
      "Data variables:\n",
      "    r1              (time, dim_0) object 4.0 ... [-145.451, 0.0, 0.0, 69.0966]\n",
      "    r2_hkl          (time, dim_1) float64 0.0 4.0 0.0\n",
      "    r2              (time, dim_2) object 0.0 ... [-145.451, 0.0, 90.0, 69.0966]\n",
      "    _keys           (time, dim_3) <U2 'r1' 'r2'\n",
      "    r2_axis_values  (time, dim_4) float64 -145.5 0.0 90.0 69.1\n",
      "    r1_hkl          (time, dim_5) float64 4.0 0.0 0.0\n",
      "    r2_wavelength   (time) float64 2.101\n",
      "    r1_wavelength   (time) float64 2.101\n",
      "    r1_axis_values  (time, dim_6) float64 -145.5 0.0 0.0 69.1\n",
      "------------------------------\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/prjemian/.conda/envs/training_2022/lib/python3.9/site-packages/dask/array/core.py:4333: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray.\n",
      "  a = np.asarray(a, dtype=dtype, order=order)\n",
      "/home/prjemian/.conda/envs/training_2022/lib/python3.9/site-packages/bluesky_live/conversion.py:380: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray.\n",
      "  out[k] = numpy.asarray(out[k])\n"
     ]
    }
   ],
   "source": [
    "def read_soft_signal(key, value):\n",
    "    yield from bps.read(Signal(name=key, value=value))\n",
    "\n",
    "def stream_dict(dictionary, label):\n",
    "    yield from bps.create(label)\n",
    "    for k, v in dictionary.items():\n",
    "        yield from read_soft_signal(k, v)\n",
    "    yield from bps.save()\n",
    "\n",
    "\n",
    "def stream_samples(samples, label=\"samples\"):\n",
    "    if len(samples):\n",
    "        yield from bps.create(label)\n",
    "        for sname, lattice in samples.items():\n",
    "            yield from read_soft_signal(sname, lattice[:])\n",
    "        yield from read_soft_signal(\"_keys\", list(lattice._fields))\n",
    "        yield from bps.save()\n",
    "    else:\n",
    "        # because you have to yield _something_\n",
    "        yield from bps.null()\n",
    "\n",
    "\n",
    "def stream_test(reflections, label=\"reflections\"):\n",
    "    if len(reflections):\n",
    "        yield from bps.create(label)\n",
    "        keys = []\n",
    "        for i, refl in enumerate(reflections):\n",
    "            key = f\"r{i+1}\"\n",
    "            keys.append(key)\n",
    "            yield from read_soft_signal(key, (*refl[:3], refl[3][:]))\n",
    "            yield from read_soft_signal(key+\"_hkl\", refl[:3])\n",
    "            yield from read_soft_signal(key+\"_axis_values\", refl[3])\n",
    "            yield from read_soft_signal(key+\"_wavelength\", 2.101)\n",
    "        yield from read_soft_signal(\"_keys\", keys)\n",
    "        yield from bps.save()\n",
    "    else:\n",
    "        # because you have to yield _something_\n",
    "        yield from bps.null()\n",
    "\n",
    "\n",
    "def stream_reflections(self, label=\"reflections\"):\n",
    "    reflections = self.calc.sample._sample.reflections_get()\n",
    "    if len(reflections):\n",
    "        yield from bps.create(label)\n",
    "        orient_refls = self.calc.sample._orientation_reflections\n",
    "        keys = []\n",
    "        for i, refl in enumerate(reflections):\n",
    "            key = f\"r{i+1}\"\n",
    "            keys.append(key)\n",
    "            hkl_tuple = refl.hkl_get()\n",
    "            geom = refl.geometry_get()\n",
    "            yield from read_soft_signal(key, (*hkl_tuple[:], geom.axis_values_get(1)))\n",
    "            yield from read_soft_signal(key+\"_hkl\", hkl_tuple[:])\n",
    "            yield from read_soft_signal(key+\"_axis_names\", geom.axis_names_get(1))\n",
    "            yield from read_soft_signal(key+\"_axis_values\", geom.axis_values_get(1))\n",
    "            yield from read_soft_signal(key+\"_wavelength\", geom.wavelength_get(1))\n",
    "            yield from read_soft_signal(key+\"_for_calcUB\", refl in orient_refls)\n",
    "            # ignore `flag`, no documentation for it, always 1 (?used by libhkl's GUI?)\n",
    "        yield from read_soft_signal(\"_keys\", keys)\n",
    "        yield from bps.save()\n",
    "    else:\n",
    "        # because you have to yield _something_\n",
    "        yield from bps.null()\n",
    "\n",
    "\n",
    "def stream_multi(label=\"multi\"):\n",
    "    for i in range(3):\n",
    "        yield from bps.create(label)\n",
    "        yield from read_soft_signal(\"a\", 1.2345 + i)\n",
    "        yield from read_soft_signal(\"b\", f\"4.{i}56\")\n",
    "        yield from read_soft_signal(\"arr\", [-1, i , 1.1])\n",
    "        yield from bps.save()\n",
    "    else:\n",
    "        # because you have to yield _something_\n",
    "        yield from bps.null()\n",
    "\n",
    "\n",
    "def streams():\n",
    "    yield from bps.open_run()\n",
    "\n",
    "    yield from stream_samples(nitwit.paddle.samples)\n",
    "    # yield from stream_reflections(\n",
    "    yield from stream_test(\n",
    "        [\n",
    "            Reflection(4.0, 0., 0., Position(omega=-145.451, chi=0.0, phi=0.0, tth=69.0966)),\n",
    "            Reflection(0., 4.0, 0., Position(omega=-145.451, chi=0.0, phi=90.0, tth=69.0966))\n",
    "        ]\n",
    "    )\n",
    "    # yield from stream_multi()\n",
    "\n",
    "    yield from bps.close_run()\n",
    "\n",
    "\n",
    "RE(streams())\n",
    "for nm, stream in cat[-1].items():\n",
    "    print(f\"{nm = }\")\n",
    "    print(stream.read())\n",
    "    print(\"-\"*30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "8cf7e9bc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(4.0, 0.0, 0.0, array([-145.451 ,    0.    ,    0.    ,   69.0966]))"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "k = \"r1\"\n",
    "r = cat[-1].reflections.read()[k+\"_hkl\"][0]\n",
    "a = cat[-1].reflections.read()[k+\"_axis_values\"][0]\n",
    "(*r.data, a.data)"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "cd09a60d4ca96784847e6d28c64916bc86a437fe6be574606d07ffca69ac8887"
  },
  "kernelspec": {
   "display_name": "Python 3.9.12 ('training_2022')",
   "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}